From e2fe233efc3c74d6f0d9bdc0a1b9eede57c1cd5c Mon Sep 17 00:00:00 2001 From: RheagalFire Date: Tue, 9 Jun 2026 22:23:45 +0530 Subject: [PATCH 1/2] feat: add LiteLLM AI gateway integration --- burr/integrations/litellm.py | 298 +++++++++++++++++++++++++++++++++++ pyproject.toml | 4 + 2 files changed, 302 insertions(+) create mode 100644 burr/integrations/litellm.py diff --git a/burr/integrations/litellm.py b/burr/integrations/litellm.py new file mode 100644 index 000000000..f546101bc --- /dev/null +++ b/burr/integrations/litellm.py @@ -0,0 +1,298 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""LiteLLM integration for Burr. + +This module provides Action classes for invoking LLM models through +the LiteLLM AI gateway within Burr applications. LiteLLM supports +100+ providers (OpenAI, Anthropic, Google, Azure, AWS Bedrock, Ollama, +Groq, Mistral, and more) through a unified interface. + +Example usage: + from burr.integrations.litellm import LiteLLMAction + + def prompt_mapper(state): + return { + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": state["user_input"]}, + ], + } + + action = LiteLLMAction( + model="anthropic/claude-sonnet-4-6", + input_mapper=prompt_mapper, + reads=["user_input"], + writes=["response"], + ) +""" + +import logging +from typing import Any, Generator, Optional, Protocol + +from burr.core.action import SingleStepAction, StreamingAction +from burr.core.state import State +from burr.integrations.base import require_plugin + +logger = logging.getLogger(__name__) + +try: + import litellm +except ImportError as e: + require_plugin(e, "litellm") + + +class StateToMessagesMapper(Protocol): + """Protocol for mapping Burr state to LiteLLM messages format.""" + + def __call__(self, state: State) -> dict[str, Any]: + ... # noqa: E704 + + +class _LiteLLMCore: + """Shared LiteLLM configuration and request building.""" + + def __init__( + self, + model: str, + input_mapper: StateToMessagesMapper, + reads: list[str], + writes: list[str], + name: str, + api_key: Optional[str], + temperature: Optional[float], + max_tokens: Optional[int], + extra_kwargs: Optional[dict[str, Any]], + ): + self._model = model + self._input_mapper = input_mapper + self._reads = reads + self._writes = writes + self._name = name + self._api_key = api_key + self._temperature = temperature + self._max_tokens = max_tokens + self._extra_kwargs = extra_kwargs or {} + + @property + def reads(self) -> list[str]: + return self._reads + + @property + def writes(self) -> list[str]: + return self._writes + + @property + def name(self) -> str: + return self._name + + def build_completion_kwargs(self, state: State, stream: bool = False) -> dict[str, Any]: + """Build kwargs for litellm.completion from current state.""" + prompt = self._input_mapper(state) + kwargs: dict[str, Any] = { + "model": self._model, + "messages": prompt["messages"], + "stream": stream, + "drop_params": True, + } + if self._api_key: + kwargs["api_key"] = self._api_key + if self._temperature is not None: + kwargs["temperature"] = self._temperature + if self._max_tokens is not None: + kwargs["max_tokens"] = self._max_tokens + kwargs.update(self._extra_kwargs) + return kwargs + + +def _result_for_writes( + text: str, + usage: dict[str, Any], + writes: list[str], +) -> dict[str, Any]: + result: dict[str, Any] = { + "response": text, + "usage": usage, + } + for w in writes: + if w == "usage": + result[w] = usage + else: + result[w] = text + return result + + +class _LiteLLMBase: + """Shared LiteLLM wiring for action subclasses.""" + + _llm: _LiteLLMCore + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + + def _init_litellm_core( + self, + model: str, + input_mapper: StateToMessagesMapper, + reads: list[str], + writes: list[str], + name: str, + api_key: Optional[str], + temperature: Optional[float], + max_tokens: Optional[int], + extra_kwargs: Optional[dict[str, Any]], + ) -> None: + self._llm = _LiteLLMCore( + model=model, + input_mapper=input_mapper, + reads=reads, + writes=writes, + name=name, + api_key=api_key, + temperature=temperature, + max_tokens=max_tokens, + extra_kwargs=extra_kwargs, + ) + + @property + def reads(self) -> list[str]: + return self._llm.reads + + @property + def writes(self) -> list[str]: + return self._llm.writes + + @property + def name(self) -> str: + return self._llm.name + + +class LiteLLMAction(_LiteLLMBase, SingleStepAction): + """Action that invokes LLM models through the LiteLLM AI gateway. + + :param model: LiteLLM model string (e.g. ``anthropic/claude-sonnet-4-6``, + ``openai/gpt-4o``, ``groq/llama-4-scout-17b-16e-instruct``). + :param input_mapper: Callable mapping :class:`~burr.core.state.State` to + a dict with a ``messages`` key (OpenAI chat format). + :param reads: State keys this action reads. + :param writes: State keys to update (typically include ``response``). + :param name: Action name for the graph. + :param api_key: Optional API key (falls back to provider env vars). + :param temperature: Optional sampling temperature. + :param max_tokens: Optional max tokens for the response. + :param extra_kwargs: Additional kwargs passed to ``litellm.completion``. + """ + + def __init__( + self, + model: str, + input_mapper: StateToMessagesMapper, + reads: list[str], + writes: list[str], + name: str = "litellm_invoke", + api_key: Optional[str] = None, + temperature: Optional[float] = None, + max_tokens: Optional[int] = None, + extra_kwargs: Optional[dict[str, Any]] = None, + ): + super().__init__() + self._init_litellm_core( + model=model, + input_mapper=input_mapper, + reads=reads, + writes=writes, + name=name, + api_key=api_key, + temperature=temperature, + max_tokens=max_tokens, + extra_kwargs=extra_kwargs, + ) + + def run_and_update(self, state: State, **run_kwargs) -> tuple[dict, State]: + kwargs = self._llm.build_completion_kwargs(state) + response = litellm.completion(**kwargs) + + text = response.choices[0].message.content or "" + usage = dict(response.usage) if response.usage else {} + + result = _result_for_writes(text, usage, self._llm.writes) + updates = {key: result[key] for key in self._llm.writes if key in result} + new_state = state.update(**updates) + + return result, new_state + + +class LiteLLMStreamingAction(_LiteLLMBase, StreamingAction): + """Streaming LiteLLM action. + + Parameters match :class:`LiteLLMAction` except the default ``name`` is + ``litellm_stream``. Yields chunk dicts from :meth:`stream_run` and merges + the final response in :meth:`update`. + """ + + def __init__( + self, + model: str, + input_mapper: StateToMessagesMapper, + reads: list[str], + writes: list[str], + name: str = "litellm_stream", + api_key: Optional[str] = None, + temperature: Optional[float] = None, + max_tokens: Optional[int] = None, + extra_kwargs: Optional[dict[str, Any]] = None, + ): + super().__init__() + self._init_litellm_core( + model=model, + input_mapper=input_mapper, + reads=reads, + writes=writes, + name=name, + api_key=api_key, + temperature=temperature, + max_tokens=max_tokens, + extra_kwargs=extra_kwargs, + ) + + def stream_run(self, state: State, **run_kwargs) -> Generator[dict, None, None]: + kwargs = self._llm.build_completion_kwargs(state, stream=True) + response = litellm.completion(**kwargs) + + text_parts: list[str] = [] + for chunk in response: + choices = getattr(chunk, "choices", None) + if not choices: + continue + delta = getattr(choices[0], "delta", None) + if delta is None: + continue + content = getattr(delta, "content", None) + if isinstance(content, str) and content: + text_parts.append(content) + full_response = "".join(text_parts) + yield {"chunk": content, "response": full_response} + + full_text = "".join(text_parts) + payload = _result_for_writes(full_text, {}, self._llm.writes) + yield {"chunk": "", "complete": True, **payload} + + def update(self, result: dict, state: State) -> State: + if result.get("complete"): + updates = {key: result[key] for key in self._llm.writes if key in result} + return state.update(**updates) + return state diff --git a/pyproject.toml b/pyproject.toml index c2ab5d202..296390ecc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -136,6 +136,10 @@ bedrock = [ "boto3" ] +litellm = [ + "litellm" +] + tracking-server-s3 = [ "aerich", "aiobotocore", From 2202e52f7600741f26e840f411313c361ffcf6e9 Mon Sep 17 00:00:00 2001 From: Aarish Irani Date: Wed, 1 Jul 2026 19:58:23 +0530 Subject: [PATCH 2/2] Add E2E example for LiteLLM integration Add examples/integrations/litellm/ with application.py demonstrating both LiteLLMAction (non-streaming) and LiteLLMStreamingAction, following the same structure as the existing Bedrock integration example. --- examples/README.md | 1 + examples/integrations/litellm/README.md | 56 ++++++++ examples/integrations/litellm/__init__.py | 16 +++ examples/integrations/litellm/application.py | 127 ++++++++++++++++++ .../integrations/litellm/requirements.txt | 18 +++ 5 files changed, 218 insertions(+) create mode 100644 examples/integrations/litellm/README.md create mode 100644 examples/integrations/litellm/__init__.py create mode 100644 examples/integrations/litellm/application.py create mode 100644 examples/integrations/litellm/requirements.txt diff --git a/examples/README.md b/examples/README.md index 9d4447cc2..530893bae 100644 --- a/examples/README.md +++ b/examples/README.md @@ -49,5 +49,6 @@ Note we have a few more in [other-examples](other-examples/), but those do not y - [multi-modal-chatbot](multi-modal-chatbot/) - This example shows how to use Burr to create a multi-modal chatbot. This demonstrates how to use a model to delegate to other models conditionally. - [streaming-overview](streaming-overview/) - This example shows how we can use the streaming API to respond to return quicker results to the user and build a seamless experience - [integrations/bedrock](integrations/bedrock/) - Minimal graphs using Amazon Bedrock (`BedrockAction` and `BedrockStreamingAction`). +- [integrations/litellm](integrations/litellm/) - Minimal graphs using the LiteLLM AI gateway (`LiteLLMAction` and `LiteLLMStreamingAction`). - [tracing-and-spans](tracing-and-spans/) - This example shows how to use Burr to create a simple chatbot with additional visibility. This is a good starting point for understanding how to use Burr's tracing functionality. - [web-server](web-server/) - This example shows how to use Burr in a web server. This is a good starting point for understanding how to use Burr for interaction. diff --git a/examples/integrations/litellm/README.md b/examples/integrations/litellm/README.md new file mode 100644 index 000000000..d3ac90d75 --- /dev/null +++ b/examples/integrations/litellm/README.md @@ -0,0 +1,56 @@ + + +# LiteLLM integration + +This example shows how to use Burr's LiteLLM helpers: + +- `LiteLLMAction` - single-step completion call via the LiteLLM AI gateway. +- `LiteLLMStreamingAction` - streaming completion with `Application.stream_result`. + +LiteLLM provides a unified interface to 100+ LLM providers (OpenAI, Anthropic, Google, Azure, AWS Bedrock, Ollama, Groq, Mistral, and more). + +## Setup + +1. Install dependencies (from the repo root): + + ```bash + pip install -r examples/integrations/litellm/requirements.txt + ``` + +2. Set the API key for the provider you want to use. For example: + + ```bash + export OPENAI_API_KEY="sk-..." # for OpenAI models + export ANTHROPIC_API_KEY="sk-ant-..." # for Anthropic models + ``` + +3. Optionally override the default model with `LITELLM_MODEL` (default is `openai/gpt-4o-mini`): + + ```bash + export LITELLM_MODEL="anthropic/claude-sonnet-4-6" + ``` + +## Run + +```bash +python examples/integrations/litellm/application.py +``` + +The script runs a non-streaming call, then a streaming call, using two small Burr graphs defined in `application()` and `streaming_application()`. diff --git a/examples/integrations/litellm/__init__.py b/examples/integrations/litellm/__init__.py new file mode 100644 index 000000000..13a83393a --- /dev/null +++ b/examples/integrations/litellm/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/examples/integrations/litellm/application.py b/examples/integrations/litellm/application.py new file mode 100644 index 000000000..7cc02cb5a --- /dev/null +++ b/examples/integrations/litellm/application.py @@ -0,0 +1,127 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Minimal Burr apps using :class:`~burr.integrations.litellm.LiteLLMAction` +and :class:`~burr.integrations.litellm.LiteLLMStreamingAction`. + +LiteLLM supports 100+ LLM providers through a unified interface. +Set the appropriate provider API key (e.g. ``OPENAI_API_KEY``, +``ANTHROPIC_API_KEY``) and optionally ``LITELLM_MODEL`` before running. +""" + +from __future__ import annotations + +import os + +from burr.core import Application, ApplicationBuilder, State, default +from burr.core.action import action +from burr.integrations.litellm import LiteLLMAction, LiteLLMStreamingAction + + +def _default_model() -> str: + return os.environ.get("LITELLM_MODEL", "openai/gpt-4o-mini") + + +def prompt_mapper(state: State) -> dict: + """Map Burr state to LiteLLM messages format (OpenAI chat format).""" + return { + "messages": [ + {"role": "system", "content": "You are a concise assistant."}, + {"role": "user", "content": state["user_input"]}, + ], + } + + +@action(reads=[], writes=["user_input"]) +def set_user_input(state: State, user_input: str) -> State: + return state.update(user_input=user_input) + + +def application(model: str | None = None) -> Application: + """Builds a graph with :class:`~burr.integrations.litellm.LiteLLMAction` (non-streaming).""" + invoke = LiteLLMAction( + model=model or _default_model(), + input_mapper=prompt_mapper, + reads=["user_input"], + writes=["response"], + name="invoke_litellm", + max_tokens=512, + ) + return ( + ApplicationBuilder() + .with_actions(set_prompt=set_user_input, invoke_litellm=invoke) + .with_transitions( + ("set_prompt", "invoke_litellm", default), + ) + .with_state(user_input="", response="") + .with_entrypoint("set_prompt") + .build() + ) + + +def streaming_application(model: str | None = None) -> Application: + """Builds a graph with :class:`~burr.integrations.litellm.LiteLLMStreamingAction`.""" + stream = LiteLLMStreamingAction( + model=model or _default_model(), + input_mapper=prompt_mapper, + reads=["user_input"], + writes=["response"], + name="stream_litellm", + max_tokens=512, + ) + return ( + ApplicationBuilder() + .with_actions(set_prompt=set_user_input, stream_litellm=stream) + .with_transitions( + ("set_prompt", "stream_litellm", default), + ) + .with_state(user_input="", response="") + .with_entrypoint("set_prompt") + .build() + ) + + +def _demo_invoke() -> None: + app = application() + _, _, state = app.run( + halt_after=["invoke_litellm"], + inputs={"user_input": "Explain what Burr is in one short sentence."}, + ) + print(state["response"]) + + +def _demo_stream() -> None: + app = streaming_application() + _, streaming_result = app.stream_result( + halt_after=["stream_litellm"], + inputs={"user_input": "Count from 1 to 5, separated by commas."}, + ) + for item in streaming_result: + chunk = item.get("chunk") or "" + if chunk: + print(chunk, end="", flush=True) + print() + _, state = streaming_result.get() + print("Final response:", state["response"]) + + +if __name__ == "__main__": + print("--- LiteLLMAction (non-streaming) ---") + _demo_invoke() + print() + print("--- LiteLLMStreamingAction (streaming) ---") + _demo_stream() diff --git a/examples/integrations/litellm/requirements.txt b/examples/integrations/litellm/requirements.txt new file mode 100644 index 000000000..407f37651 --- /dev/null +++ b/examples/integrations/litellm/requirements.txt @@ -0,0 +1,18 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +burr[litellm]