#479
Add LiteLLM provider for pentestgpt-legacy
Add LiteLLM provider for pentestgpt-legacy
LiteLLM is an open-source gateway that provides a unified completion() interface for 100+ LLM providers. Adding it as a provider gives pentestgpt-legacy users access to Azure, AWS Bedrock, Google Vertex, HuggingFace, vLLM, and any OpenAI-compatible endpoint, all through a single litellm:<model> prefix, without per-provider code.
Changes
| File | Change |
|---|---|
llm/providers/litellm_provider.py | New LiteLLMProvider(BaseProvider) using litellm.acompletion() with drop_params=True |
llm/registry.py | Added litellm ProviderInfo + litellm: dynamic prefix resolver |
llm/factory.py | Registered LiteLLMProvider in _PROVIDER_CLASSES |
llm/providers/__init__.py | Export LiteLLMProvider |
llm/config.py | Added LITELLM_API_KEY and LITELLM_BASE_URL settings |
pyproject.toml | Added litellm>=1.80.0 as optional dependency |
.env.example | Documented LiteLLM configuration |
tests/legacy/test_litellm_provider.py | 6 unit tests covering registry, resolve, factory, and proxy config |
Usage
# Install with the litellm extra
uv sync --extra litellm
# Use any LiteLLM-supported model
pentestgpt-legacy --reasoning-model litellm:azure/gpt-4 --parsing-model litellm:bedrock/anthropic.claude-v2
# With a LiteLLM proxy
LITELLM_API_KEY=sk-proxy LITELLM_BASE_URL=http://proxy:4000 \
pentestgpt-legacy --reasoning-model litellm:gpt-4Tests
Live E2E - full LLMClient stack via Azure Foundry (anthropic/claude-sonnet-4-6):
from pentestgpt_legacy.llm.factory import get_client
# Factory resolves litellm: prefix -> LiteLLMProvider -> litellm.acompletion()
client = get_client('litellm:anthropic/claude-sonnet-4-6')
# Test 1: New conversation
response, cid = client.send_new_message('Say hello in 5 words.')
print(f"Response: {response}")
print(f"Conversation ID: {cid}")
# Response: Hello there, how are you?
# Conversation ID: <uuid>
# Test 2: Multi-turn (same conversation)
response2 = client.send_message('Now say goodbye in 3 words.', cid)
print(f"Follow-up: {response2}")
# Follow-up: Goodbye for now!Full chain verified: get_client('litellm:...') -> LLMClientFactory -> LiteLLMProvider.acompletion() -> litellm.acompletion(drop_params=True) -> Azure Foundry -> response parsed correctly. Multi-turn conversation state maintained across calls.
Activity
5gemini-code-assist[bot]reviewed4mo ago
Code Review
This pull request introduces integration with LiteLLM as an optional multi-provider gateway, adding configuration options, a dedicated LiteLLMProvider class, and corresponding unit tests. The feedback suggests two improvements: updating the installation error message to recommend installing the optional dependency via 'uv sync --extra litellm' instead of 'uv add litellm', and safely guarding against an empty response.choices list to prevent a potential IndexError.