diff --git a/goldenverba/components/generation/LiteLLMGenerator.py b/goldenverba/components/generation/LiteLLMGenerator.py new file mode 100644 index 000000000..4fe706ed1 --- /dev/null +++ b/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 diff --git a/goldenverba/components/managers.py b/goldenverba/components/managers.py index e6454d081..dc8391068 100644 --- a/goldenverba/components/managers.py +++ b/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(), ] diff --git a/goldenverba/tests/test_litellm_generator.py b/goldenverba/tests/test_litellm_generator.py new file mode 100644 index 000000000..a15e6b6fc --- /dev/null +++ b/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 diff --git a/setup.py b/setup.py index 9f20bc4c1..1f1e9589c 100644 --- a/setup.py +++ b/setup.py @@ -56,5 +56,8 @@ "huggingface": [ "sentence-transformers==3.0.1", ], + "litellm": [ + "litellm>=1.60,<1.85", + ], }, )