feat: route scanners through LiteLLM SDK for multi-provider support

[prodmanpd]

Summary

Routes every Python scanner's LLM calls through the LiteLLM SDK instead of a raw openai client, so AI-Infra-Guard can talk to 100+ providers (OpenAI, Anthropic, Gemini, Bedrock, Vertex, Azure, Groq, Mistral, self hosted, or a LiteLLM proxy) through one interface.

• Additive and fully backward compatible: the current OpenRouter default and any custom base_url keep working byte for byte.
• Adds native provider auth (for example Bedrock SigV4, Vertex ADC) that a plain openai client plus base_url cannot reach without standing up a separate gateway.
• Sets drop_params=True so one config works across providers that reject each other's generation params.

How routing works (to_litellm_params)

• base_url set (default, for example OpenRouter or any OpenAI compatible gateway): the request is sent as a custom OpenAI compatible endpoint (model="openai/<model>", api_base=<base_url>). The wire request is identical to the previous openai.OpenAI(base_url) call, so existing configs are unaffected.
• base_url empty: the model string is passed through for native provider routing (anthropic/..., bedrock/..., gemini/...), resolving credentials from that provider's own env vars, with no gateway or proxy.

Changes

• agent-scan/agent_scan/utils/llm.py: LLM now calls litellm.completion(stream=True); LiteLLM exception handling; to_litellm_params helper.
• mcp-scan/mcp_scan/utils/llm.py: same for the sync LLM; adds LiteLLMAsyncClient, a shim compatible with the OpenAI AsyncClient over litellm.acompletion for the red team engine.
• mcp-scan/mcp_scan/redteam/{orchestrator,evaluator,attacker,target}.py: the async red team path now uses LiteLLMAsyncClient (call sites unchanged; response shape identical).
• skill-scan/skill_scan/utils/llm.py: same for the sync LLM (preserves stream_options usage and custom headers via extra_headers).
• */pyproject.toml, */requirements.txt: add litellm>=1.89.0,<2.0.0
• Tests: */pytests/test_llm_request.py updated to the LiteLLM seam; mcp-scan/pytests/test_redteam_litellm_client.py added.

Tests

1. Unit tests, 15 pass (litellm 1.98.0, openai 2.54.0 transitive):

agent-scan: pytests/test_llm_request.py .....                    5 passed
mcp-scan:   pytests/test_llm_request.py + test_redteam_litellm_client.py ......   6 passed
skill-scan: pytests/test_llm_request.py ....                     4 passed

Coverage: OpenAI compatible routing (openai/ prefix plus api_base), native passthrough when base_url blank, drop_params=True present, blank API key sent as None (so LiteLLM falls back to provider env vars), stream_options and extra_headers preserved, and the async red team shim forwarding to litellm.acompletion.

2. Lint (repo ruff config, py312, line length 100): authored files clean:

agent-scan: ruff check agent_scan/utils/llm.py pytests/test_llm_request.py  -> All checks passed!
skill-scan: ruff check skill_scan/utils/llm.py pytests/test_llm_request.py  -> All checks passed!
mcp-scan:   ruff check mcp_scan/utils/llm.py pytests/*                        -> All checks passed!

3. Pin resolves cleanly: pip install "litellm>=1.89.0,<2.0.0" gives litellm 1.98.0 plus openai 2.54.0.

4. Live E2E (real code path through a LiteLLM proxy): all three scanner LLM paths exercised end to end against a live model (gpt-4.1-mini via a LiteLLM proxy), deterministic prompt "Reply with exactly: OK":

== agent-scan LLM.chat (sync stream) ==            agent-scan -> 'OK'
== mcp-scan  LLM.chat (sync stream + usage) ==     mcp-scan  -> 'OK'  usage: {'prompt_tokens': 12, 'completion_tokens': 1, 'total_tokens': 13}
== mcp-scan  redteam async shim -> acompletion ==  redteam    -> 'OK'
ALL LIVE E2E PASSED

This proves the full chain for both the sync litellm.completion path (agent-scan, mcp-scan, and skill-scan share it) and the async litellm.acompletion red team shim: request routed through LiteLLM, provider returned, response and usage parsed back through each scanner's existing code.

Risk and compatibility

• Additive. Existing OpenRouter and OpenAI compatible configs unchanged.
• openai remains available (LiteLLM depends on it); no scanner imports it directly anymore.
• No DB, schema, or API changes.

Example usage

# Existing OpenAI-compatible gateway (e.g. OpenRouter), unchanged:

# New: native provider, no gateway or proxy (creds from provider env vars):
LLM(model="anthropic/claude-sonnet-4.5", api_key="", base_url="")   # uses ANTHROPIC_API_KEY
LLM(model="bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0", api_key="", base_url="")

# Or point at a LiteLLM proxy:
LLM(model="my-model", api_key=PROXY_KEY, base_url="http://localhost:4000/v1")

[boy-hack]

Thanks @prodmanpd for this — routing every scanner through LiteLLM is a great capability win (100+ providers, native Bedrock/Vertex auth, one config via drop_params=True).

The design is solid: the openai/-prefix + api_base trick keeps the OpenRouter default byte-for-byte compatible, and the LiteLLMAsyncClient shim for the red-team agents is a clean minimal surface. Tests cover both the sync and async paths. A few things worth confirming before merge:

1. Dependency weight / supply-chain: litellm is a large package with many transitive deps. Swapping openai==2.8.1 (one lib) for litellm>=1.89.0,<2.0.0 noticeably increases install size and the security surface of AIG's shipped scanners. Acceptable for the feature, but please make sure the CI install job still passes and consider whether the range <2.0.0 could pull a 1.x release with breaking internal changes (LiteLLM has historically done minor-breaking bumps within 1.x). A tighter floor or a note in the PR would help.
2. Module-load cost: import litellm at the top of llm.py now runs in every scanner process startup. Verify it doesn't add noticeable cold-start latency to mcp-scan/skill-scan CLI invocations (and that importing it doesn't fail under the minimal requirements.txt install).
3. Non-streaming chat() in agent-scan: the diff updates chat_stream to litellm.completion, but chat() (non-stream) still runs litellm.completion synchronously inside asyncio.to_thread — confirm that path is still exercised/needed, otherwise it's dead weight alongside chat_async.
4. Out-of-scope consumers still import openai: AIG-PromptSecurity/cli/model_utils/openailike.py and skills/aig-agent-redteam/scripts/common/llm_client.py still import openai directly. AIG-PromptSecurity has its own pyproject.toml so it's fine, and the redteam skill uses a try/except ImportError fallback — just flagging so nobody assumes the openai dep is gone repo-wide.

None of these are blockers for the design. My main ask is (1)/(2): confirm CI install + import works under the pinned range. Nice work overall!

[liuzhao1225]

I'm concerned about the supply-chain boundary of this LiteLLM integration. The upstream incident report confirms compromised PyPI releases 1.82.7/1.82.8 in March 2026; those releases were removed.

At reviewed head 9466706f10, this PR includes litellm>=1.89.0,<2.0.0 in agent-scan/pyproject.toml; litellm>=1.89.0,<2.0.0 in agent-scan/requirements.txt; litellm>=1.89.0,<2.0.0 in mcp-scan/pyproject.toml; litellm>=1.89.0,<2.0.0 in mcp-scan/requirements.txt; litellm>=1.89.0,<2.0.0 in skill-scan/pyproject.toml.

That upper bound permits later 1.x releases beyond a single tested version when dependencies are resolved afresh or upgraded. Please document tested versions, artifact verification, and whether installation of the SDK is opt-in. A disabled application feature cannot prevent an installed malicious Python-startup hook.

This is one of 49 observed LiteLLM integration PRs from the same account, whose author acknowledged the cross-project effort. The pattern prompted this review; it does not establish malicious intent or connect the author to the incident. The consolidated questions and corrections are in YouDub #130; this PR's review snapshot preserves the revision and scope.

[boy-hack]

Thanks @prodmanpd — the design holds up: the openai/-prefix + api_base trick keeps the OpenRouter default byte-for-byte compatible, drop_params=True is the right lever for cross-provider configs, and the LiteLLMAsyncClient shim is a minimal, clean surface. Tests covering both sync and async paths plus a live E2E through a proxy are strong.

Re my earlier asks (1)-(4): the transitive openai/litellm versions resolve cleanly (1.98.0) and the out-of-scope openai consumers are correctly left alone, so those are satisfied. Two things I'd like tightened before merge, prompted partly by @liuzhao1225's supply-chain note:

1. Pin the tested version (supply-chain hygiene). The compromised PyPI releases (1.82.7/1.82.8, March 2026) were removed, and your floor >=1.89.0 never reaches them — so this PR does not pull the bad versions. But the upper bound <2.0.0 still permits any future 1.x release to be resolved fresh, which is exactly the drift the supply-chain reviewer is worried about. Recommendation: pin to the version you actually tested and verified — e.g. litellm==1.98.0 (or at least >=1.98.0,<1.99.0) — so installs are reproducible, and add a one-line note in each scanner's README that the litellm dependency should be installed from a verified index. That directly answers the "document tested versions / artifact verification" ask without weakening the feature. (For deployment you could also recommend hash-pinning.)

2. Confirm minimal-install import + cold-start. import litellm at module top now runs in every scanner process startup. Please confirm import litellm doesn't fail or add noticeable latency under each scanner's minimal requirements.txt install (litellm pulls a lot of transitive deps) — especially for mcp-scan/skill-scan CLI cold starts. The CI install job passing is the proof point here.

On the broader "49 cross-project LiteLLM PRs" framing: that's context about the reviewer's scope, not evidence of malicious intent in this change. The integration itself is legitimate and a clear capability win. Acting on point 1 keeps us safe regardless.

Also still worth a quick confirm: the non-streaming chat() path still calls litellm.completion inside asyncio.to_thread alongside chat_async — is chat() still exercised, or is it dead weight now? If unused, dropping it reduces surface.

Nice work overall — let's tighten the pin and confirm the install path, then this is good to go.

[prodmanpd]

@boy-hack thanks for the detailed analysis, will tighten this.

[boy-hack]

Thanks @prodmanpd — following up on the supply-chain thread raised by @liuzhao1225 and my earlier review.

The good news: the floor litellm>=1.89.0 already sits above the compromised 1.82.7/1.82.8 releases from the March 2026 incident, so those specific poisoned builds can't be resolved. That materially reduces the acute risk.

Since you said you'd tighten this, here's the concrete ask before I'm comfortable approving:

1. Pin a tested version. Rather than >=1.89.0,<2.0.0, ship the version you actually validated against (litellm==1.98.0 per your E2E log) across all five manifests (agent-scan/mcp-scan/skill-scan pyproject.toml + requirements.txt). If you want a patch-level floor, a tight ==1.98.* or >=1.98.0,<1.99.0 is far safer than an open 1.x ceiling. Document the tested version in the PR and in a comment near the dependency.
2. Artifact verification. Add a short note (PR body or requirements comment) on how the wheel is sourced/verified — e.g. pulled from PyPI over TLS and hash-checked in CI if you have a lockfile. A disabled feature can't stop a malicious installed package, but a pinned + hash-locked install can.
3. Confirm it's effectively opt-in. import litellm lives in each scanner's llm.py, and each scanner ships its own requirements.txt. So the SDK is only pulled when a user installs that specific scanner — not a global always-on hook. Please confirm that's the deployment model (it looks like it is) so the "opt-in" boundary is explicit.

Two smaller things from my first pass that still stand:

• Import cost: import litellm at module top now runs on every scanner CLI startup. The E2E you ran exercised a live call but not cold-start latency — a quick time on mcp-scan --help / skill-scan before vs after would be reassuring, though this is minor.
• The earlier note that AIG-PromptSecurity/.../openailike.py and the redteam skill still import openai directly is fine (separate pyproject.toml / try-except ImportError); just confirming nobody assumes openai is gone repo-wide.

Net: the design is good and the openai/-prefix + api_base routing keeps the OpenRouter default byte-for-byte compatible (tests confirm no double-prefixing). I'm holding approval until the version pin + verification note land. Happy to re-review once pushed.

(Review comment only — not merging.)