Feat/add litellm provider

backend/chainlit/__init__.py

@@ -100,6 +100,7 @@
         AsyncLangchainCallbackHandler,
         LangchainCallbackHandler,
     )
+    from chainlit.litellm import instrument_litellm
     from chainlit.llama_index.callbacks import LlamaIndexCallbackHandler
     from chainlit.mistralai import instrument_mistralai
     from chainlit.openai import instrument_openai
@@ -131,6 +132,7 @@ def acall(self):
         "LlamaIndexCallbackHandler": "chainlit.llama_index.callbacks",
         "instrument_openai": "chainlit.openai",
         "instrument_mistralai": "chainlit.mistralai",
+        "instrument_litellm": "chainlit.litellm",
         "SemanticKernelFilter": "chainlit.semantic_kernel",
         "server": "chainlit.server",
     }
@@ -186,6 +188,7 @@ def acall(self):
     "data_layer",
     "header_auth_callback",
     "input_widget",
+    "instrument_litellm",
     "instrument_mistralai",
     "instrument_openai",
     "make_async",

backend/chainlit/litellm/__init__.py

@@ -0,0 +1,162 @@
+import asyncio
+import time
+from typing import Any, Dict, List, Optional
+
+from literalai import ChatGeneration, GenerationMessage
+
+from chainlit.context import get_context
+from chainlit.step import Step
+from chainlit.utils import timestamp_utc
+
+
+def _build_generation(
+    kwargs: Dict[str, Any],
+    response_obj: Any,
+    start_time: float,
+    end_time: float,
+) -> ChatGeneration:
+    """Build a ChatGeneration from litellm call kwargs and response."""
+    model = kwargs.get("model", "")
+    messages = kwargs.get("messages", [])
+    settings: Dict[str, Any] = {
+        "model": model,
+        "temperature": kwargs.get("temperature"),
+        "max_tokens": kwargs.get("max_tokens"),
+        "top_p": kwargs.get("top_p"),
+        "stop": kwargs.get("stop"),
+        "stream": kwargs.get("stream", False),
+    }
+    settings = {k: v for k, v in settings.items() if v is not None}
+
+    gen_messages: List[GenerationMessage] = []
+    for msg in messages:
+        gen_messages.append(
+            GenerationMessage(
+                role=msg.get("role", "user"),
+                content=msg.get("content", ""),
+            )
+        )
+
+    message_completion: Optional[GenerationMessage] = None
+    input_tokens: Optional[int] = None
+    output_tokens: Optional[int] = None
+
+    if response_obj and hasattr(response_obj, "choices") and response_obj.choices:
+        choice = response_obj.choices[0]
+        message = getattr(choice, "message", None)
+        if message:
+            message_completion = GenerationMessage(
+                role=getattr(message, "role", "assistant"),
+                content=getattr(message, "content", "") or "",
+            )
+
+    if response_obj and hasattr(response_obj, "usage") and response_obj.usage:
+        usage = response_obj.usage
+        input_tokens = getattr(usage, "prompt_tokens", None)
+        output_tokens = getattr(usage, "completion_tokens", None)
+
+    duration = end_time - start_time
+
+    provider = model.split("/")[0] if "/" in model else "litellm"
+
+    return ChatGeneration(
+        provider=provider,
+        model=model,
+        settings=settings,
+        messages=gen_messages,
+        message_completion=message_completion,
+        input_token_count=input_tokens,
+        output_token_count=output_tokens,
+        duration=duration,
+    )
+
+
+def instrument_litellm() -> None:
+    """Instrument LiteLLM to automatically log completions as Chainlit Steps.
+
+    Uses litellm's native CustomLogger callback API. Each litellm.completion()
+    call will appear as an LLM step in the Chainlit UI with model name,
+    messages, response, token usage, and timing.
+
+    Example:
+        import chainlit as cl
+        cl.instrument_litellm()
+
+        # All subsequent litellm.completion() calls will be traced
+        import litellm
+        response = litellm.completion(
+            model="anthropic/claude-sonnet-4-20250514",
+            messages=[{"role": "user", "content": "Hello!"}],
+        )
+    """
+    try:
+        import importlib.util
+
+        if importlib.util.find_spec("litellm") is None:
+            raise ImportError
+    except ImportError:
+        raise ValueError("litellm is not installed. Run `pip install litellm`")
+
+    import litellm
+    from litellm.integrations.custom_logger import CustomLogger
+
+    class ChainlitLogger(CustomLogger):
+        def log_success_event(
+            self,
+            kwargs: Dict[str, Any],
+            response_obj: Any,
+            start_time: float,
+            end_time: float,
+        ) -> None:
+            try:
+                import datetime
+
+                def _to_epoch(t: Any) -> float:
+                    if isinstance(t, (int, float)):
+                        return float(t)
+                    if isinstance(t, datetime.datetime):
+                        return t.timestamp()
+                    return time.time()
+
+                start_ts = _to_epoch(start_time)
+                end_ts = _to_epoch(end_time)
+
+                generation = _build_generation(kwargs, response_obj, start_ts, end_ts)
+                context = get_context()
+
+                parent_id = None
+                if context.current_step:
+                    parent_id = context.current_step.id
+
+                step = Step(
+                    name=generation.model or "litellm",
+                    type="llm",
+                    parent_id=parent_id,
+                )
+                step.generation = generation
+                step.start = timestamp_utc(start_ts)
+                step.end = timestamp_utc(end_ts)
+
+                if generation.messages:
+                    step.input = generation.messages  # type: ignore
+                if generation.message_completion:
+                    step.output = generation.message_completion  # type: ignore
+
+                asyncio.create_task(step.send())
+            except Exception:
+                pass
+
+        async def async_log_success_event(
+            self,
+            kwargs: Dict[str, Any],
+            response_obj: Any,
+            start_time: float,
+            end_time: float,
+        ) -> None:
+            self.log_success_event(kwargs, response_obj, start_time, end_time)
+
+    for cb in litellm.callbacks:
+        if type(cb).__name__ == "ChainlitLogger":
+            return
+
+    litellm.callbacks.append(ChainlitLogger())

backend/pyproject.toml

@@ -78,6 +78,7 @@ tests = [
     "pytest-asyncio>=1.0.0,<2.0.0",
     "pytest-cov>=6.0.0,<7.0.0",
     "openai>=2.0.0,<3.0.0",
+    "litellm>=1.55,<2.0",
     "langchain>=1.0.0",
     "llama-index>=0.13.0,<1.0.0",
     "semantic-kernel>=1.24.0,<2.0.0",

backend/tests/litellm/test_litellm.py

@@ -0,0 +1,152 @@
+import time
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from chainlit.litellm import _build_generation, instrument_litellm
+from chainlit.step import Step
+
+
+def _make_response(content="test response", model="anthropic/claude-sonnet-4-6"):
+    message = SimpleNamespace(content=content, role="assistant")
+    choice = SimpleNamespace(message=message, finish_reason="stop", index=0)
+    usage = SimpleNamespace(prompt_tokens=10, completion_tokens=5, total_tokens=15)
+    return SimpleNamespace(choices=[choice], usage=usage, model=model, id="mock")
+
+
+def _make_kwargs(model="anthropic/claude-sonnet-4-6", messages=None):
+    return {
+        "model": model,
+        "messages": messages or [{"role": "user", "content": "Hello"}],
+        "temperature": 0.7,
+        "max_tokens": 100,
+        "stream": False,
+    }
+
+
+class TestBuildGeneration:
+    def test_basic_generation(self):
+        kwargs = _make_kwargs()
+        response = _make_response()
+        gen = _build_generation(kwargs, response, 1000.0, 1002.0)
+
+        assert gen.model == "anthropic/claude-sonnet-4-6"
+        assert gen.provider == "anthropic"
+        assert gen.settings["model"] == "anthropic/claude-sonnet-4-6"
+        assert gen.settings["temperature"] == 0.7
+        assert gen.settings["max_tokens"] == 100
+        assert gen.duration == pytest.approx(2.0)
+
+    def test_messages_converted(self):
+        messages = [
+            {"role": "system", "content": "You are helpful."},
+            {"role": "user", "content": "Hi"},
+        ]
+        kwargs = _make_kwargs(messages=messages)
+        response = _make_response()
+        gen = _build_generation(kwargs, response, 0, 1)
+
+        assert len(gen.messages) == 2
+        assert gen.messages[0]["role"] == "system"
+        assert gen.messages[0]["content"] == "You are helpful."
+        assert gen.messages[1]["role"] == "user"
+        assert gen.messages[1]["content"] == "Hi"
+
+    def test_completion_extracted(self):
+        kwargs = _make_kwargs()
+        response = _make_response(content="The answer is 4.")
+        gen = _build_generation(kwargs, response, 0, 1)
+
+        assert gen.message_completion is not None
+        assert gen.message_completion["role"] == "assistant"
+        assert gen.message_completion["content"] == "The answer is 4."
+
+    def test_token_counts(self):
+        kwargs = _make_kwargs()
+        response = _make_response()
+        gen = _build_generation(kwargs, response, 0, 1)
+
+        assert gen.input_token_count == 10
+        assert gen.output_token_count == 5
+
+    def test_provider_extraction(self):
+        kwargs = _make_kwargs(model="openai/gpt-4o")
+        response = _make_response(model="openai/gpt-4o")
+        gen = _build_generation(kwargs, response, 0, 1)
+        assert gen.provider == "openai"
+
+    def test_provider_without_slash(self):
+        kwargs = _make_kwargs(model="gpt-4o")
+        response = _make_response(model="gpt-4o")
+        gen = _build_generation(kwargs, response, 0, 1)
+        assert gen.provider == "litellm"
+
+    def test_none_response(self):
+        kwargs = _make_kwargs()
+        gen = _build_generation(kwargs, None, 0, 1)
+        assert gen.message_completion is None
+        assert gen.input_token_count is None
+
+    def test_empty_choices(self):
+        response = SimpleNamespace(choices=[], usage=None, model="test", id="x")
+        kwargs = _make_kwargs()
+        gen = _build_generation(kwargs, response, 0, 1)
+        assert gen.message_completion is None
+
+    def test_settings_strip_none(self):
+        kwargs = {"model": "test", "messages": [], "stream": False}
+        response = _make_response()
+        gen = _build_generation(kwargs, response, 0, 1)
+        assert "temperature" not in gen.settings
+        assert "top_p" not in gen.settings
+
+
+class TestInstrumentLiteLLM:
+    def test_registers_callback(self):
+        import litellm
+
+        original_count = len(litellm.callbacks)
+        instrument_litellm()
+        assert len(litellm.callbacks) >= original_count + 1
+
+    def test_idempotent(self):
+        instrument_litellm()
+        import litellm
+
+        count_after = len(litellm.callbacks)
+        instrument_litellm()
+        assert len(litellm.callbacks) == count_after
+
+    def test_callback_creates_step(self, mock_chainlit_context):
+        import litellm
+
+        instrument_litellm()
+        logger = [
+            cb for cb in litellm.callbacks if type(cb).__name__ == "ChainlitLogger"
+        ][-1]
+
+        kwargs = _make_kwargs()
+        response = _make_response()
+
+        async def run():
+            async with mock_chainlit_context:
+                with patch.object(Step, "send", new_callable=AsyncMock):
+                    logger.log_success_event(kwargs, response, time.time(), time.time())
+
+        import asyncio
+
+        asyncio.get_event_loop().run_until_complete(run())
+
+    def test_callback_handles_no_context(self):
+        import litellm
+
+        instrument_litellm()
+        logger = [
+            cb for cb in litellm.callbacks if type(cb).__name__ == "ChainlitLogger"
+        ][-1]
+
+        kwargs = _make_kwargs()
+        response = _make_response()
+
+        logger.log_success_event(kwargs, response, time.time(), time.time())