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

@elijahbenizzy elijahbenizzy left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! Any chance you could add an E2E example?

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.

2 participants