FEAT: add LiteLLMChatTarget for multi-provider access via LiteLLM by RheagalFire · Pull Request #2154 · microsoft/PyRIT

RheagalFire · 2026-07-09T19:23:47Z

Description

Adds LiteLLMChatTarget, a new prompt target that uses the LiteLLM SDK (litellm.acompletion()) to reach 100+ providers (Anthropic, AWS Bedrock, Google Vertex, Cohere, etc.) directly — no separate proxy server required. LiteLLM speaks the OpenAI Chat Completions wire format, so the target shares its request-building and response-parsing logic with OpenAIChatTarget instead of reinventing it.

from pyrit.prompt_target import LiteLLMChatTarget

# Anthropic (LiteLLM reads ANTHROPIC_API_KEY from env)
target = LiteLLMChatTarget(model_name="anthropic/claude-sonnet-4-6")

# AWS Bedrock (reads AWS credentials from env)
target = LiteLLMChatTarget(model_name="bedrock/anthropic.claude-v2")

# Self-hosted LiteLLM gateway / proxy
target = LiteLLMChatTarget(
    model_name="anthropic/claude-sonnet-4-6",
    endpoint="http://localhost:4000",
    api_key="sk-proxy-key",
)

Key features

Multi-provider via LiteLLM — one target routes to any LiteLLM-supported provider through litellm.acompletion().
Capabilities derived from the model — input/output modalities and JSON support are read from LiteLLM's own model metadata (supports_vision, supports_audio_input/supports_audio_output, supports_response_schema, get_supported_openai_params) at construction, so multimodal (image/audio) input and audio output are enabled only when the resolved model actually supports them. Falls back to a text-only default when metadata is unavailable.
Auth flexibility — explicit api_key, a sync/async token provider callable (Entra-style), the LITELLM_API_KEY env var, or LiteLLM's own provider-specific env var lookup.
Curated OpenAI params + passthrough — temperature, top_p, max_tokens, frequency_penalty, presence_penalty, seed, n, stop, plus arbitrary provider params via extra_body_parameters.
Single max_tokens knob — LiteLLM normalizes max_tokens to whatever each model/provider expects (e.g. max_completion_tokens for OpenAI reasoning/gpt-5 models), so callers set one value and it works cross-provider. Provider-specific token params remain reachable via extra_body_parameters.
drop_unsupported_params init arg (default True) — first-class control over LiteLLM's drop_params, letting a single target send the full OpenAI parameter set across providers with differing support; set False for strict validation. Overridable per request via extra_body_parameters.
Token usage capture — provider usage payloads are parsed into a new provider-neutral TokenUsage value object and persisted to message metadata.
Retry + exception translation — honors PyRIT's retry configuration via LiteLLM's num_retries; litellm.exceptions.* are mapped to PyRIT's RateLimitException / PyritException hierarchy (rate-limit and transient errors become retryable, no bare Exception catches).
Consistent identifier — behavioral params are recorded on the target identifier; the API key is intentionally excluded.

Shared Chat Completions helpers (refactor)

To avoid duplicating logic between the OpenAI and LiteLLM targets, the Chat Completions request-building and response-parsing code is extracted into shared modules, and OpenAIChatTarget is refactored to use them (net removal of ~250 lines there):

pyrit/prompt_target/common/chat_completions_message_builder.py — builds Chat Completions request messages from PyRIT messages.
pyrit/prompt_target/common/chat_completions_response_parser.py — parses Chat Completions responses (text/tool calls/audio, finish-reason validation, token usage).
pyrit/models/token_usage.py — new provider-neutral TokenUsage value object (input/output vocabulary aligned with the Responses API/Anthropic/Gemini), with metadata (de)serialization. Provider→TokenUsage conversion lives in the parser/target that knows the wire format, not in the value object.

Dependency

litellm>=1.83.0,<2.0.0 added as an optional dependency: a named litellm extra and included in the all group (pip install pyrit[litellm] or pip install pyrit[all]). It is intentionally not a core dependency because current LiteLLM releases ship Linux-only wheels plus a Rust-based sdist, which would break pip install pyrit on Windows/macOS. import litellm is lazy, so users without the package are unaffected.

Tests and Documentation

Unit tests (all passing):

tests/unit/prompt_target/target/test_litellm_chat_target.py — 54 tests (construction/env fallback, auth resolution, capability derivation, param/body construction, drop_unsupported_params, max_tokens normalization + passthrough, audio, token usage, finish-reason and empty/malformed handling, exception translation, retries).
tests/unit/prompt_target/target/test_chat_completions_helpers.py — 39 tests for the shared message builder / response parser.
tests/unit/models/test_token_usage.py — 7 tests for the TokenUsage value object.
tests/unit/prompt_target/target/test_openai_chat_target.py — updated for the shared-helper refactor.

Integration tests:

tests/integration/targets/test_litellm_chat_target_integration.py — new, reusing existing Azure OpenAI deployments (routed through LiteLLM) so LiteLLM is exercised against real endpoints.
tests/integration/targets/test_openai_chat_target_integration.py — updated (default integration model wired to an Azure GPT-5 deployment with Entra auth).

Pre-commit (ruff format/check, ty type check, async-suffix and Sphinx-role hooks) passes clean.

rlundeen2 · 2026-07-10T00:21:19Z

I've been wanting to do this also; @RheagalFire this is a great start. I'll likely push to your branch to flush it out and support all the things (multi-modal, identifiers, underlying model, integration tests, capabilities detection, other gaps). So from your perspective I would consider this "merged" and I'll take it and try to get it in before the next release.

TY for the help and nudge!

…s, token usage) Extends the LiteLLM target for parity with OpenAIChatTarget and shares logic instead of reinventing it: - Extract shared Chat Completions helpers (chat_completions_message_builder, chat_completions_response_parser) used by both OpenAIChatTarget and LiteLLMChatTarget for request building and response parsing (text, image, audio, tool calls, content-filter handling). - Add multimodal support (image + audio input, audio output via audio_response_config) with capabilities derived from LiteLLM's model registry and a conservative text-only fallback. - Support the full OpenAI parameter set plus an extra_body_parameters passthrough; auth via sync/async token providers; identifiers that exclude the api_key; underlying_model capability lookup; and LiteLLM-owned retry (num_retries from PyRIT's global convention) to avoid double-retrying. - Add a provider-neutral TokenUsage value object (input/output/total/reasoning/cached + extra) and capture it for both targets; capture LiteLLM per-call cost. - Add litellm as an optional extra and include it in the 'all' extra. - Modernize type syntax (X | None), tidy docstrings per the style guide, and add unit + integration tests (image/audio on a gpt-5 deployment). Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

… object Make TokenUsage a pure value object (fields + to_metadata/from_metadata) and move Chat Completions usage parsing into chat_completions_response_parser as token_usage_from_chat_completion, explicit to the one wire shape both chat targets actually send. - Drop the speculative Responses-API sniffing (no caller sends that shape); a Responses target should parse its own usage in its own module. - Tolerate dict-or-attribute usage payloads so a mapping no longer silently yields all-None counts. - Capture LiteLLM/Anthropic top-level cache fields (cache_read_input_tokens -> cached_tokens, cache_creation_input_tokens -> extra), preserving a zero cached count. - Move/expand parsing tests into test_chat_completions_helpers; test_token_usage now covers only metadata round-tripping. - Convert remaining Sphinx roles to plain double-backtick references. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

…ts to gpt-5.4 Revert the PLATFORM_OPENAI_CHAT_MODEL / PLATFORM_OPENAI_AUDIO_MODEL additions to .env_example and default the chat/vision integration fixtures to the deployed gpt-5.4 model instead of the generic "gpt-5". Also convert two stray Sphinx roles in the integration test docstrings to plain double-backtick references. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Rename the platform_* chat fixtures to azure_gpt5_* and point them at the Azure OpenAI GPT-5.4 deployment via LiteLLM's OpenAI-compatible openai/ prefix. The deployment is keyless, so fall back to a DefaultAzureCredential bearer-token provider when no key is set, which also exercises callable/Entra auth support. Audio tests stay on the platform OpenAI gpt-audio deployment. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: c01c3489-dd68-4db5-bad7-025b9e700379

Resolve conflicts from main's chat-target modernization and module relocations against this PR's shared chat-completion helper extraction: - Keep litellm optional: add a named `litellm` extra and include it in the `all` group (not a core dependency), avoiding the Rust-only litellm 1.92.0 build on Windows/macOS. Regenerate uv.lock accordingly. - Point shared helpers and the LiteLLM target at main's relocated APIs: data_serializer_factory/DataTypeSerializer now from pyrit.memory, convert_local_image_to_data_url_async from pyrit.memory.storage, _JsonResponseConfig from pyrit.prompt_target.common.json_response_config (json_config.json_schema), and async serializer methods (save_data_async / save_formatted_audio_async / read_data_base64_async). - ComponentIdentifier now imported from pyrit.models; get_known_capabilities used as the module-level function; _construct_message_from_response_async renamed for the async-suffix rule. - Drop the optional-install "pip install pyrit[litellm]" garble from the docstring and reference the clean extra. - Update affected unit tests to the relocated import paths and async mocks. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: c01c3489-dd68-4db5-bad7-025b9e700379

rlundeen2

I wrote a lot of it though so needs another review

Promote LiteLLM's core drop_params behavior to a first-class, documented constructor argument (drop_unsupported_params, default True). It controls whether provider-unsupported request params are silently dropped (the cross-provider default) or raise for strict validation. A per-request extra_body_parameters={"drop_params": ...} still overrides it. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: c01c3489-dd68-4db5-bad7-025b9e700379

Remove max_completion_tokens and the mutual-exclusivity raise. LiteLLM normalizes max_tokens to the parameter each model/provider expects (e.g. max_completion_tokens for gpt-5/o-series), so a single knob works consistently across providers. Provider-specific token params remain reachable via extra_body_parameters. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: c01c3489-dd68-4db5-bad7-025b9e700379

litellm 1.92.0 dropped the universal py3-none-any wheel in favor of manylinux-only wheels backed by a Rust/PyO3 extension. That breaks installs everywhere except Linux cp310-cp313: Windows/macOS have no wheels, and the sdist fails to build on Python 3.14 (PyO3 0.23.5 caps at 3.13). Pin to >=1.83.0,<1.92.0 (1.91.3 is the last pure-Python line) in both the litellm extra and the all group, and relock (1.92.0 -> 1.91.3). Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: c01c3489-dd68-4db5-bad7-025b9e700379

behnam-o

overall looks good. Looks like quite a bit of deviation/overrides from the OpenAI Target patterns, but maybe this is even better and more generic (probably)...

behnam-o · 2026-07-13T22:45:29Z

+    Returns:
+        Message: The constructed error Message with ``error="blocked"``.
+    """
+    logger.warning("Output content filtered by content policy.")


I don't think this log warning should happen here. probably right before constructing the message here ?

behnam-o · 2026-07-13T22:48:36Z

+    return pieces
+
+
+def build_text_and_tool_pieces(*, response: Any, request: MessagePiece) -> list[MessagePiece]:


This is not used anywhere?

behnam-o · 2026-07-13T22:50:58Z

+    return audio_serializer.value
+
+
+async def build_audio_pieces_async(


this seems very internal? add _ maybe?

behnam-o · 2026-07-13T23:05:43Z

+        self,
+        *,
+        messages: list[dict[str, Any]],
+        json_config: Any,


why not restrict this type to _JsonResponseConfig ?

behnam-o · 2026-07-13T23:25:15Z

+        n: int | None = None,
+        stop: str | list[str] | None = None,
+        audio_response_config: OpenAIChatAudioConfig | None = None,
+        drop_unsupported_params: bool = True,


nit: I'm wondering having this default to True will make it really hard to debug ? I can see benefits of it too, increasing the probability of it just working and less manual work with changing LITELLM_MODEL between different providers ... I'm 50/50 , so as long as this is intentional ...

feat: add LiteLLM as AI gateway prompt target

42d212b

rlundeen2 self-assigned this Jul 9, 2026

rlundeen2 and others added 3 commits July 10, 2026 14:00

romanlutz reviewed Jul 10, 2026

View reviewed changes

Comment thread pyproject.toml Outdated

rlundeen2 and others added 2 commits July 13, 2026 09:23

rlundeen2 approved these changes Jul 13, 2026

View reviewed changes

rlundeen2 and others added 2 commits July 13, 2026 12:01

behnam-o self-assigned this Jul 13, 2026

rlundeen2 changed the title ~~feat: add LiteLLM as AI gateway prompt target~~ feat: add LiteLLMChatTarget for multi-provider access via LiteLLM Jul 13, 2026

behnam-o changed the title ~~feat: add LiteLLMChatTarget for multi-provider access via LiteLLM~~ FEAT: add LiteLLMChatTarget for multi-provider access via LiteLLM Jul 13, 2026

behnam-o approved these changes Jul 13, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

FEAT: add LiteLLMChatTarget for multi-provider access via LiteLLM#2154

FEAT: add LiteLLMChatTarget for multi-provider access via LiteLLM#2154
RheagalFire wants to merge 9 commits into
microsoft:mainfrom
RheagalFire:feat/add-litellm-provider

RheagalFire commented Jul 9, 2026 •

edited by rlundeen2

Loading

Uh oh!

rlundeen2 commented Jul 10, 2026 •

edited

Loading

Uh oh!

Uh oh!

rlundeen2 left a comment

Uh oh!

behnam-o left a comment

Uh oh!

behnam-o Jul 13, 2026

Uh oh!

behnam-o Jul 13, 2026

Uh oh!

behnam-o Jul 13, 2026

Uh oh!

behnam-o Jul 13, 2026

Uh oh!

behnam-o Jul 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		return pieces


		def build_text_and_tool_pieces(*, response: Any, request: MessagePiece) -> list[MessagePiece]:

		return audio_serializer.value


		async def build_audio_pieces_async(

Uh oh!

Conversation

RheagalFire commented Jul 9, 2026 • edited by rlundeen2 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Key features

Shared Chat Completions helpers (refactor)

Dependency

Tests and Documentation

Uh oh!

rlundeen2 commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

rlundeen2 left a comment

Choose a reason for hiding this comment

Uh oh!

behnam-o left a comment

Choose a reason for hiding this comment

Uh oh!

behnam-o Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

behnam-o Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

behnam-o Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

behnam-o Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

behnam-o Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

RheagalFire commented Jul 9, 2026 •

edited by rlundeen2

Loading

rlundeen2 commented Jul 10, 2026 •

edited

Loading