Add LiteLLM provider for pentestgpt-legacy

.env.example

@@ -26,6 +26,14 @@ GROK_API_KEY=              # xAI (or XAI_API_KEY)
 QWEN_API_KEY=             # Alibaba DashScope (or DASHSCOPE_API_KEY)
 KIMI_API_KEY=            # Moonshot (or MOONSHOT_API_KEY)
 
+# =============================================================================
+# LiteLLM (optional — install with: uv sync --extra litellm)
+# =============================================================================
+# Use any model LiteLLM supports: litellm:azure/gpt-4, litellm:bedrock/claude-v2, etc.
+# When connecting to a LiteLLM proxy, set both:
+# LITELLM_API_KEY=
+# LITELLM_BASE_URL=
+
 # Optional base-URL overrides (defaults are built in per provider)
 # OLLAMA_BASE_URL=http://localhost:11434/v1
 # OPENAI_BASE_URL=

pentestgpt_legacy/llm/config.py

@@ -23,6 +23,7 @@
     "qwen": "qwen_api_key",
     "moonshot": "kimi_api_key",
     "ollama": None,
+    "litellm": "litellm_api_key",
 }
 
 
@@ -52,6 +53,7 @@ class LLMSettings(BaseSettings):
     kimi_api_key: str | None = Field(
         default=None, validation_alias=AliasChoices("KIMI_API_KEY", "MOONSHOT_API_KEY")
     )
+    litellm_api_key: str | None = None
 
     # Optional per-provider base-URL overrides (defaults come from the registry).
     openai_base_url: str | None = None
@@ -60,6 +62,7 @@ class LLMSettings(BaseSettings):
     qwen_base_url: str | None = None
     moonshot_base_url: str | None = None
     ollama_base_url: str | None = None
+    litellm_base_url: str | None = None
 
     def api_key_for(self, provider: ProviderInfo) -> str | None:
         """Return the configured API key for ``provider`` (or ``None``)."""

pentestgpt_legacy/llm/factory.py

@@ -8,6 +8,7 @@
 from pentestgpt_legacy.llm.providers import (
     AnthropicProvider,
     GeminiProvider,
+    LiteLLMProvider,
     OpenAICompatibleProvider,
 )
 from pentestgpt_legacy.llm.registry import (
@@ -23,6 +24,7 @@
     "openai": OpenAICompatibleProvider,
     "anthropic": AnthropicProvider,
     "gemini": GeminiProvider,
+    "litellm": LiteLLMProvider,
 }
 
 

pentestgpt_legacy/llm/providers/__init__.py

@@ -2,6 +2,7 @@
 
 from pentestgpt_legacy.llm.providers.anthropic_provider import AnthropicProvider
 from pentestgpt_legacy.llm.providers.gemini_provider import GeminiProvider
+from pentestgpt_legacy.llm.providers.litellm_provider import LiteLLMProvider
 from pentestgpt_legacy.llm.providers.openai_compatible import OpenAICompatibleProvider
 
-__all__ = ["AnthropicProvider", "GeminiProvider", "OpenAICompatibleProvider"]
+__all__ = ["AnthropicProvider", "GeminiProvider", "LiteLLMProvider", "OpenAICompatibleProvider"]

pentestgpt_legacy/llm/providers/litellm_provider.py

@@ -0,0 +1,63 @@
+"""LiteLLM provider — unified gateway to 100+ LLM providers.
+
+Routes through ``litellm.acompletion`` with ``drop_params=True`` so
+provider-specific parameter mismatches are silently handled. Supports any
+model string that LiteLLM recognises (Azure, Bedrock, Vertex, HuggingFace,
+Ollama, vLLM, custom OpenAI-compatible endpoints, etc.).
+
+When used with a LiteLLM proxy, set ``LITELLM_BASE_URL`` and
+``LITELLM_API_KEY`` in the environment or ``.env`` file.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pentestgpt_legacy.llm.base import BaseProvider, Message
+from pentestgpt_legacy.llm.registry import ModelSpec
+
+
+class LiteLLMProvider(BaseProvider):
+    """Async LiteLLM provider (unified multi-provider gateway)."""
+
+    async def acomplete(
+        self,
+        messages: list[Message],
+        system: str | None,
+        spec: ModelSpec,
+        *,
+        max_output_tokens: int | None = None,
+        temperature: float | None = None,
+    ) -> str:
+        try:
+            import litellm
+        except ImportError as exc:  # pragma: no cover - install-time guard
+            raise RuntimeError(
+                "The 'litellm' package is required. Install with: uv sync --extra litellm"
+            ) from exc
+
+        wire_messages: list[dict[str, Any]] = []
+        if system:
+            wire_messages.append({"role": "system", "content": system})
+        wire_messages.extend(messages)
+
+        kwargs: dict[str, Any] = {
+            "model": spec.api_id,
+            "messages": wire_messages,
+            "drop_params": True,
+        }
+        if temperature is not None:
+            kwargs["temperature"] = temperature
+        if max_output_tokens is not None:
+            kwargs["max_tokens"] = max_output_tokens
+        if self.api_key:
+            kwargs["api_key"] = self.api_key
+        if self.base_url:
+            kwargs["api_base"] = self.base_url
+
+        response = await litellm.acompletion(**kwargs)
+        if not response.choices:
+            raise RuntimeError(
+                f"LiteLLM returned no choices for model '{spec.api_id}'"
+            )
+        return response.choices[0].message.content or ""

pentestgpt_legacy/llm/registry.py

@@ -113,10 +113,18 @@ def __post_init__(self) -> None:
         base_url="http://localhost:11434/v1",
         requires_key=False,
     ),
+    "litellm": ProviderInfo(
+        key="litellm",
+        label="LiteLLM",
+        kind="litellm",
+        env="LITELLM_API_KEY",
+        requires_key=False,
+    ),
 }
 
-# Marker provider used for ``ollama:<model>`` dynamic ids.
+# Marker providers used for dynamic ``<prefix>:<model>`` ids.
 OLLAMA_PREFIX = "ollama:"
+LITELLM_PREFIX = "litellm:"
 
 
 # --------------------------------------------------------------------------- #
@@ -292,6 +300,18 @@ def resolve(name: str) -> ModelSpec | None:
             tier="local",
             notes="User-configured local Ollama model",
         )
+    if name.startswith(LITELLM_PREFIX):
+        model = name[len(LITELLM_PREFIX) :].strip()
+        if not model:
+            return None
+        return ModelSpec(
+            id=name,
+            provider="litellm",
+            api_id=model,
+            context_window=128_000,
+            tier="gateway",
+            notes="LiteLLM-routed model",
+        )
     if name in MODELS:
         return MODELS[name]
     if name in _ALIASES:

pyproject.toml

@@ -39,6 +39,9 @@ dependencies = [
     "google-genai>=1.0.0",
 ]
 
+[project.optional-dependencies]
+litellm = ["litellm>=1.80.0,<1.87.0"]
+
 [project.scripts]
 pentestgpt = "pentestgpt.interface.main:main"
 pentestgpt-legacy = "pentestgpt_legacy.main:main"

tests/legacy/test_litellm_provider.py

@@ -0,0 +1,57 @@
+"""Tests for the LiteLLM provider integration."""
+
+import pytest
+
+from pentestgpt_legacy.llm import factory
+from pentestgpt_legacy.llm.config import LLMSettings
+from pentestgpt_legacy.llm.factory import get_client
+from pentestgpt_legacy.llm.providers import LiteLLMProvider
+from pentestgpt_legacy.llm.registry import PROVIDERS, resolve
+
+pytestmark = pytest.mark.unit
+
+
+def _settings(**kwargs: str) -> LLMSettings:
+    return LLMSettings(_env_file=None, **kwargs)
+
+
+def test_litellm_provider_in_registry() -> None:
+    assert "litellm" in PROVIDERS
+    assert PROVIDERS["litellm"].kind == "litellm"
+    assert PROVIDERS["litellm"].requires_key is False
+
+
+def test_resolve_litellm_dynamic() -> None:
+    spec = resolve("litellm:azure/gpt-4")
+    assert spec is not None
+    assert spec.provider == "litellm"
+    assert spec.api_id == "azure/gpt-4"
+    assert spec.tier == "gateway"
+
+
+def test_resolve_litellm_bedrock() -> None:
+    spec = resolve("litellm:bedrock/anthropic.claude-v2")
+    assert spec is not None
+    assert spec.api_id == "bedrock/anthropic.claude-v2"
+
+
+def test_resolve_litellm_empty_is_none() -> None:
+    assert resolve("litellm:") is None
+
+
+def test_get_client_litellm(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(factory, "get_settings", lambda: _settings())
+    client = get_client("litellm:azure/gpt-4")
+    assert isinstance(client.provider, LiteLLMProvider)
+    assert client.spec.api_id == "azure/gpt-4"
+
+
+def test_get_client_litellm_with_proxy_key(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(
+        factory,
+        "get_settings",
+        lambda: _settings(litellm_api_key="sk-proxy", litellm_base_url="http://proxy:4000"),
+    )
+    client = get_client("litellm:gpt-4")
+    assert client.provider.api_key == "sk-proxy"
+    assert client.provider.base_url == "http://proxy:4000"