feat: add LiteLLM generator

goldenverba/components/generation/LiteLLMGenerator.py

@@ -0,0 +1,151 @@
+import os
+
+from dotenv import load_dotenv
+from wasabi import msg
+
+from goldenverba.components.interfaces import Generator
+from goldenverba.components.types import InputConfig
+
+load_dotenv()
+
+
+class LiteLLMGenerator(Generator):
+    """LiteLLM Generator.
+
+    Routes chat completions through ``litellm.acompletion()`` to 100+ providers
+    (OpenAI, Anthropic, Bedrock, Vertex, Gemini, Ollama, OpenRouter, Groq, DeepSeek,
+    etc.) using provider-native API keys. The user picks any LiteLLM-supported
+    model by its prefixed name (``anthropic/claude-3-5-sonnet-20241022``,
+    ``gemini/gemini-1.5-pro``, ``bedrock/anthropic.claude-3-sonnet-20240229-v1:0``,
+    ``ollama/llama3``, ...).
+
+    See https://docs.litellm.ai/docs/providers for the full list.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.name = "LiteLLM"
+        self.description = (
+            "Using LiteLLM to route to 100+ LLM providers via a unified interface"
+        )
+        self.context_window = 10000
+        # LiteLLM is an optional dep; surface this in the UI via the standard
+        # requires_library availability check.
+        self.requires_library = ["litellm"]
+
+        default_model = os.getenv("LITELLM_MODEL", "openai/gpt-4o-mini")
+        self.config["Model"] = InputConfig(
+            type="text",
+            value=default_model,
+            description=(
+                "LiteLLM-style model name, e.g. 'anthropic/claude-3-5-sonnet-20241022', "
+                "'gemini/gemini-1.5-pro', 'bedrock/anthropic.claude-3-sonnet-20240229-v1:0', "
+                "'ollama/llama3'. See https://docs.litellm.ai/docs/providers."
+            ),
+            values=[],
+        )
+        if os.getenv("LITELLM_API_KEY") is None:
+            self.config["API Key"] = InputConfig(
+                type="password",
+                value="",
+                description=(
+                    "Optional provider API key. If left blank, LiteLLM falls back to "
+                    "the provider-specific env var (OPENAI_API_KEY, ANTHROPIC_API_KEY, "
+                    "GEMINI_API_KEY, AWS_*, GROQ_API_KEY, ...) based on the selected model."
+                ),
+                values=[],
+            )
+        if os.getenv("LITELLM_BASE_URL") is None:
+            self.config["URL"] = InputConfig(
+                type="text",
+                value="",
+                description=(
+                    "Optional custom base URL, forwarded to LiteLLM as 'api_base'. "
+                    "Leave blank to use the provider's default endpoint."
+                ),
+                values=[],
+            )
+
+    async def generate_stream(
+        self,
+        config: dict,
+        query: str,
+        context: str,
+        conversation: list[dict] = [],
+    ):
+        try:
+            import litellm  # lazy import; optional dep
+        except ImportError as err:
+            msg.warn(
+                "LiteLLM is not installed. Install with: pip install 'goldenverba[litellm]'"
+            )
+            raise ImportError(
+                "LiteLLM is not installed. Install with: pip install 'goldenverba[litellm]'"
+            ) from err
+
+        system_message = config.get("System Message").value
+        model = config.get("Model").value
+
+        # Optional credentials — read directly so a missing value doesn't raise.
+        api_key = self._resolve_optional(config, "API Key", "LITELLM_API_KEY")
+        api_base = self._resolve_optional(config, "URL", "LITELLM_BASE_URL")
+
+        messages = self.prepare_messages(query, context, conversation, system_message)
+
+        kwargs: dict = {
+            "model": model,
+            "messages": messages,
+            "stream": True,
+        }
+        if api_key:
+            kwargs["api_key"] = api_key
+        if api_base:
+            kwargs["api_base"] = api_base
+
+        response = await litellm.acompletion(**kwargs)
+        async for chunk in response:
+            if not getattr(chunk, "choices", None):
+                continue
+            choice = chunk.choices[0]
+            delta = getattr(choice, "delta", None)
+            content = getattr(delta, "content", None) if delta else None
+            finish_reason = getattr(choice, "finish_reason", None)
+
+            if content:
+                yield {"message": content, "finish_reason": finish_reason}
+            elif finish_reason:
+                yield {"message": "", "finish_reason": finish_reason}
+
+    @staticmethod
+    def _resolve_optional(config: dict, config_key: str, env_key: str) -> str | None:
+        """Read an optional value from the user config, falling back to an env var.
+
+        Returns ``None`` when neither is set so the caller can skip forwarding the
+        kwarg and let LiteLLM pick up provider-specific env vars on its own.
+        """
+        if config_key in config:
+            value = config[config_key].value
+            if value:
+                return value
+        env_value = os.getenv(env_key)
+        return env_value or None
+
+    def prepare_messages(
+        self,
+        query: str,
+        context: str,
+        conversation: list[dict],
+        system_message: str,
+    ) -> list[dict]:
+        messages = [{"role": "system", "content": system_message}]
+
+        for message in conversation:
+            messages.append({"role": message.type, "content": message.content})
+
+        messages.append(
+            {
+                "role": "user",
+                "content": f"Answer this query: '{query}' with this provided context: {context}",
+            }
+        )
+        return messages

goldenverba/components/managers.py

@@ -71,6 +71,7 @@
 from goldenverba.components.generation.GroqGenerator import GroqGenerator
 from goldenverba.components.generation.NovitaGenerator import NovitaGenerator
 from goldenverba.components.generation.UpstageGenerator import UpstageGenerator
+from goldenverba.components.generation.LiteLLMGenerator import LiteLLMGenerator
 
 try:
     import tiktoken
@@ -118,6 +119,7 @@
         GroqGenerator(),
         NovitaGenerator(),
         UpstageGenerator(),
+        LiteLLMGenerator(),
     ]
 else:
     readers = [
@@ -152,6 +154,7 @@
         AnthropicGenerator(),
         CohereGenerator(),
         UpstageGenerator(),
+        LiteLLMGenerator(),
     ]
 
 

goldenverba/tests/test_litellm_generator.py

@@ -0,0 +1,156 @@
+"""Unit tests for LiteLLMGenerator."""
+
+from __future__ import annotations
+
+import sys
+import types
+from types import SimpleNamespace
+from unittest import mock
+
+import pytest
+
+from goldenverba.components.generation.LiteLLMGenerator import LiteLLMGenerator
+
+
+def _install_litellm_stub(acompletion_return):
+    """Register a fake ``litellm`` module so ``import litellm`` resolves."""
+    fake = types.ModuleType("litellm")
+
+    async def _acompletion(**kwargs):
+        # Record the call for assertions on the pytest side.
+        _acompletion.call_args = kwargs  # type: ignore[attr-defined]
+        return acompletion_return
+
+    fake.acompletion = _acompletion
+    sys.modules["litellm"] = fake
+    return _acompletion
+
+
+def _make_chunk(
+    content: str | None, finish_reason: str | None = None
+) -> SimpleNamespace:
+    delta = SimpleNamespace(content=content)
+    choice = SimpleNamespace(delta=delta, finish_reason=finish_reason)
+    return SimpleNamespace(choices=[choice])
+
+
+class _AsyncIter:
+    def __init__(self, items):
+        self._items = items
+
+    def __aiter__(self):
+        return self._iter()
+
+    async def _iter(self):
+        for item in self._items:
+            yield item
+
+
+def test_litellm_generator_metadata():
+    gen = LiteLLMGenerator()
+    assert gen.name == "LiteLLM"
+    assert "LiteLLM" in gen.description
+    assert "litellm" in gen.requires_library
+
+
+def test_litellm_generator_config_surface():
+    gen = LiteLLMGenerator()
+    assert "Model" in gen.config
+    # System Message is set by the parent Generator __init__
+    assert "System Message" in gen.config
+    # API Key / URL are conditionally added when the corresponding env var is unset.
+
+
+def test_prepare_messages_builds_system_user_pair():
+    gen = LiteLLMGenerator()
+    messages = gen.prepare_messages(
+        query="What is Verba?",
+        context="Verba is a Weaviate-based RAG stack.",
+        conversation=[],
+        system_message="You are Verba.",
+    )
+    assert messages[0] == {"role": "system", "content": "You are Verba."}
+    assert messages[-1]["role"] == "user"
+    assert "What is Verba?" in messages[-1]["content"]
+    assert "Verba is a Weaviate-based RAG stack." in messages[-1]["content"]
+
+
+@pytest.mark.asyncio
+async def test_generate_stream_yields_chunks_and_finish():
+    """generate_stream must yield {'message', 'finish_reason'} dicts and forward
+    the configured model + optional api_key/api_base to litellm.acompletion."""
+    chunks = [
+        _make_chunk("Hello "),
+        _make_chunk("world", finish_reason=None),
+        _make_chunk(None, finish_reason="stop"),
+    ]
+    acompletion = _install_litellm_stub(_AsyncIter(chunks))
+
+    gen = LiteLLMGenerator()
+    config = {
+        "Model": SimpleNamespace(value="anthropic/claude-3-5-sonnet-20241022"),
+        "System Message": SimpleNamespace(value="You are Verba."),
+        "API Key": SimpleNamespace(value="sk-test"),
+        "URL": SimpleNamespace(value="https://proxy.example.com/v1"),
+    }
+
+    outputs = []
+    async for item in gen.generate_stream(
+        config=config, query="hi", context="ctx", conversation=[]
+    ):
+        outputs.append(item)
+
+    assert outputs[0] == {"message": "Hello ", "finish_reason": None}
+    assert outputs[1] == {"message": "world", "finish_reason": None}
+    assert outputs[2] == {"message": "", "finish_reason": "stop"}
+
+    kwargs = acompletion.call_args  # type: ignore[attr-defined]
+    assert kwargs["model"] == "anthropic/claude-3-5-sonnet-20241022"
+    assert kwargs["stream"] is True
+    assert kwargs["api_key"] == "sk-test"
+    assert kwargs["api_base"] == "https://proxy.example.com/v1"
+
+
+@pytest.mark.asyncio
+async def test_generate_stream_omits_api_key_when_blank():
+    """When neither config API Key nor LITELLM_API_KEY is set, the kwarg must be
+    omitted so LiteLLM can fall back to provider-specific env vars."""
+    acompletion = _install_litellm_stub(
+        _AsyncIter([_make_chunk("ok", finish_reason="stop")])
+    )
+
+    gen = LiteLLMGenerator()
+    config = {
+        "Model": SimpleNamespace(value="openai/gpt-4o-mini"),
+        "System Message": SimpleNamespace(value="You are Verba."),
+        "API Key": SimpleNamespace(value=""),
+        "URL": SimpleNamespace(value=""),
+    }
+
+    async for _ in gen.generate_stream(
+        config=config, query="hi", context="ctx", conversation=[]
+    ):
+        pass
+
+    kwargs = acompletion.call_args  # type: ignore[attr-defined]
+    assert "api_key" not in kwargs
+    assert "api_base" not in kwargs
+
+
+@pytest.mark.asyncio
+async def test_generate_stream_raises_import_error_without_litellm():
+    """If ``litellm`` isn't installed, generate_stream should raise ImportError
+    with an install hint matching our ``extras_require`` entry."""
+    sys.modules.pop("litellm", None)
+    # Also ensure it can't be found by the import machinery during the test.
+    with mock.patch.dict(sys.modules, {"litellm": None}):
+        gen = LiteLLMGenerator()
+        config = {
+            "Model": SimpleNamespace(value="openai/gpt-4o-mini"),
+            "System Message": SimpleNamespace(value="You are Verba."),
+        }
+        with pytest.raises(ImportError, match="goldenverba\\[litellm\\]"):
+            async for _ in gen.generate_stream(
+                config=config, query="hi", context="ctx", conversation=[]
+            ):
+                pass

setup.py

@@ -56,5 +56,8 @@
         "huggingface": [
             "sentence-transformers==3.0.1",
         ],
+        "litellm": [
+            "litellm>=1.60,<1.85",
+        ],
     },
 )