Skip to content

feat: add LiteLLM AI gateway integration#803

Open
RheagalFire wants to merge 1 commit into
apache:mainfrom
RheagalFire:feat/add-litellm-provider
Open

feat: add LiteLLM AI gateway integration#803
RheagalFire wants to merge 1 commit into
apache:mainfrom
RheagalFire:feat/add-litellm-provider

Conversation

@RheagalFire

Copy link
Copy Markdown

Adds LiteLLM AI gateway integration, giving Burr users access to 100+ LLM providers (OpenAI, Anthropic, Google, Azure, AWS Bedrock, Ollama, Groq, Mistral, and more) through a single Action class.

Changes

  • burr/integrations/litellm.py - New LiteLLMAction (non-streaming) and LiteLLMStreamingAction (streaming) following the same _Core + _Base + Action pattern as burr/integrations/bedrock.py. Uses litellm.completion() with drop_params=True for cross-provider compatibility.
  • pyproject.toml - Added [litellm] optional dependency group (burr[litellm])

The codebase already references LiteLLM in TODO comments (core/parallelism.py, docs/concepts/parallelism.rst) and uses LiteLLM pricing data for cost tracking (telemetry/ui/scripts/token_costs.py). This PR ships the actual integration.

How I tested this

Live E2E via Anthropic Claude:

from burr.core.state import State
from burr.integrations.litellm import LiteLLMAction

def prompt_mapper(state):
    return {'messages': [
        {'role': 'system', 'content': 'Answer concisely.'},
        {'role': 'user', 'content': state['user_input']},
    ]}

action = LiteLLMAction(
    model='anthropic/claude-sonnet-4-6',
    input_mapper=prompt_mapper,
    reads=['user_input'],
    writes=['response'],
)

state = State({'user_input': 'What is 2+2? Just the number.'})
result, new_state = action.run_and_update(state)
# Response: "4"
# State updated: True

Pin resolution verified: pip install -e '.[litellm]' resolves cleanly.

Notes

  • Follows the exact same architecture as the Bedrock integration (_BedrockCore / _BedrockBase / BedrockAction / BedrockStreamingAction)
  • litellm is an optional extra, base install unaffected
  • drop_params=True silently drops provider-unsupported kwargs so the same action works across OpenAI, Anthropic, Gemini, etc.

Checklist

  • PR has an informative and human-readable title (this will be pulled into the release notes)
  • Changes are limited to a single goal (no scope creep)
  • Code passed the pre-commit check & code is left cleaner/nicer than when first encountered.
  • Any change in functionality is tested
  • New functions are documented (with a description, list of inputs, and expected output)
  • Placeholder code is flagged / future TODOs are captured in comments
  • Project documentation has been updated if adding/changing functionality.

@RheagalFire

Copy link
Copy Markdown
Author

cc @elijahbenizzy @skrawcz

@github-actions github-actions Bot added area/integrations External integrations (LLMs, frameworks) area/ci Workflows, build, release scripts labels Jun 9, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area/ci Workflows, build, release scripts area/integrations External integrations (LLMs, frameworks)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant