From 568230a143f07fff6e1deec9bdcb278bf37e0637 Mon Sep 17 00:00:00 2001 From: RheagalFire Date: Sun, 7 Jun 2026 02:11:39 +0530 Subject: [PATCH] feat: add LiteLLM as AI gateway provider --- app/ai/providers/base.py | 19 +- app/ai/providers/litellm_provider.py | 253 +++++++++++++++++++++++++++ app/ai/registry.py | 22 +++ pyproject.toml | 1 + 4 files changed, 288 insertions(+), 7 deletions(-) create mode 100644 app/ai/providers/litellm_provider.py diff --git a/app/ai/providers/base.py b/app/ai/providers/base.py index 62f937bd..f1d5d7d6 100644 --- a/app/ai/providers/base.py +++ b/app/ai/providers/base.py @@ -18,6 +18,7 @@ # Provider enum — add new providers here # --------------------------------------------------------------------------- + class ProviderType(str, Enum): GOOGLE = "google" OPENAI = "openai" @@ -25,12 +26,14 @@ class ProviderType(str, Enum): OLLAMA = "ollama" VOYAGE = "voyage" COHERE = "cohere" + LITELLM = "litellm" # --------------------------------------------------------------------------- # Runtime config loaded from DB # --------------------------------------------------------------------------- + @dataclass class ProviderConfig: """ @@ -42,19 +45,21 @@ class ProviderConfig: need capability metadata (context window, supports_tools, ...) read it from `config.spec` rather than hard-coding per model_id. """ + provider: ProviderType api_key: str = "" model_id: str = "" - base_url: Optional[str] = None # For Ollama, Azure, proxies - dimensions: Optional[int] = None # Embedding output dimensions + base_url: Optional[str] = None # For Ollama, Azure, proxies + dimensions: Optional[int] = None # Embedding output dimensions extra: dict = field(default_factory=dict) # Provider-specific params - spec: Optional[object] = None # LLMModelSpec | VisionModelSpec | EmbeddingModelSpec + spec: Optional[object] = None # LLMModelSpec | VisionModelSpec | EmbeddingModelSpec # --------------------------------------------------------------------------- # Embedding # --------------------------------------------------------------------------- + class EmbeddingProvider(ABC): """Generate vector embeddings for text.""" @@ -91,6 +96,7 @@ def dimensions(self) -> int: # LLM (text generation) # --------------------------------------------------------------------------- + class LLMProvider(ABC): """Generate text — used for summarization, webhook gateway, etc.""" @@ -127,14 +133,14 @@ async def generate_with_tools( ) @abstractmethod - async def test_connection(self) -> tuple[bool, str]: - ... + async def test_connection(self) -> tuple[bool, str]: ... # --------------------------------------------------------------------------- # Vision (image analysis) # --------------------------------------------------------------------------- + class VisionProvider(ABC): """Analyze images — used during document ingestion for image captioning.""" @@ -152,5 +158,4 @@ async def analyze_image( ... @abstractmethod - async def test_connection(self) -> tuple[bool, str]: - ... + async def test_connection(self) -> tuple[bool, str]: ... diff --git a/app/ai/providers/litellm_provider.py b/app/ai/providers/litellm_provider.py new file mode 100644 index 00000000..f1a8821a --- /dev/null +++ b/app/ai/providers/litellm_provider.py @@ -0,0 +1,253 @@ +""" +LiteLLM provider — embedding, LLM, and vision via the LiteLLM SDK. + +Routes through LiteLLM's unified Python SDK to 100+ LLM providers +(Anthropic, OpenAI, Google, Bedrock, Azure, Ollama, etc.) using +provider-prefixed model IDs. ``drop_params=True`` is always passed so +provider-unsupported kwargs are silently dropped. + +Works in two modes: + - Direct SDK: set api_key per provider (ANTHROPIC_API_KEY, etc.) + - Via proxy: set base_url to a running LiteLLM proxy +""" + +import asyncio +import base64 +import json +from typing import Optional + +from loguru import logger + +from app.ai.agent_protocol import ( + AssistantTurn, + ToolCall, +) +from app.ai.providers.base import ( + EmbeddingProvider, + LLMProvider, + VisionProvider, +) + + +def _import_litellm(): + try: + import litellm + + return litellm + except ImportError as exc: + raise ImportError( + "litellm is not installed. Install with: pip install 'arkon[litellm]'" + ) from exc + + +class LiteLLMEmbedding(EmbeddingProvider): + """LiteLLM embedding provider.""" + + async def embed(self, text: str) -> list[float]: + litellm = _import_litellm() + + kwargs: dict = { + "model": self.config.model_id, + "input": [text], + "drop_params": True, + } + if self.config.api_key: + kwargs["api_key"] = self.config.api_key + if self.config.base_url: + kwargs["api_base"] = self.config.base_url + + response = await litellm.aembedding(**kwargs) + embedding = response.data[0].embedding + + if self.config.dimensions and len(embedding) > self.config.dimensions: + embedding = embedding[: self.config.dimensions] + return embedding + + async def embed_batch( + self, texts: list[str], concurrency: int = 5 + ) -> list[list[float]]: + semaphore = asyncio.Semaphore(concurrency) + + async def _embed_one(text: str) -> list[float]: + async with semaphore: + return await self.embed(text) + + results = await asyncio.gather(*[_embed_one(t) for t in texts]) + return list(results) + + async def test_connection(self) -> tuple[bool, str]: + try: + result = await self.embed("test connection") + dim = len(result) + return True, f"OK — model={self.config.model_id}, dimensions={dim}" + except Exception as e: + return False, f"LiteLLM embedding error: {e}" + + +class LiteLLMChat(LLMProvider): + """LiteLLM LLM provider.""" + + async def generate( + self, + prompt: str, + system: Optional[str] = None, + max_tokens: Optional[int] = None, + temperature: float = 0.7, + ) -> str: + litellm = _import_litellm() + + messages: list[dict] = [] + if system: + messages.append({"role": "system", "content": system}) + messages.append({"role": "user", "content": prompt}) + + kwargs: dict = { + "model": self.config.model_id, + "messages": messages, + "temperature": temperature, + "drop_params": True, + } + if self.config.api_key: + kwargs["api_key"] = self.config.api_key + if self.config.base_url: + kwargs["api_base"] = self.config.base_url + if max_tokens is not None: + kwargs["max_tokens"] = max_tokens + + response = await litellm.acompletion(**kwargs) + return response.choices[0].message.content or "" + + async def generate_with_tools( + self, + messages: list[dict], + tools: list[dict], + system: Optional[str] = None, + max_tokens: Optional[int] = None, + temperature: float = 0.2, + ) -> AssistantTurn: + litellm = _import_litellm() + + llm_messages: list[dict] = [] + if system: + llm_messages.append({"role": "system", "content": system}) + llm_messages.extend(messages) + + kwargs: dict = { + "model": self.config.model_id, + "messages": llm_messages, + "tools": tools, + "temperature": temperature, + "drop_params": True, + } + if self.config.api_key: + kwargs["api_key"] = self.config.api_key + if self.config.base_url: + kwargs["api_base"] = self.config.base_url + if max_tokens is not None: + kwargs["max_tokens"] = max_tokens + + response = await litellm.acompletion(**kwargs) + + choice = response.choices[0] + message = choice.message + text = message.content + tool_calls: list[ToolCall] = [] + + if hasattr(message, "tool_calls") and message.tool_calls: + for tc in message.tool_calls: + args: dict = {} + if tc.function.arguments: + try: + args = json.loads(tc.function.arguments) + except Exception: + pass + tool_calls.append( + ToolCall(id=tc.id, name=tc.function.name, arguments=args) + ) + + reason_map = { + "stop": "end_turn", + "tool_calls": "tool_use", + "length": "max_tokens", + } + finish_reason = reason_map.get(choice.finish_reason or "stop", "end_turn") + + return AssistantTurn( + text=text or None, + tool_calls=tool_calls, + finish_reason=finish_reason, + ) + + async def test_connection(self) -> tuple[bool, str]: + try: + result = await self.generate("Say 'OK'", max_tokens=10, temperature=0) + return True, f"OK — model={self.config.model_id}, response='{result[:50]}'" + except Exception as e: + return False, f"LiteLLM LLM error: {e}" + + +class LiteLLMVision(VisionProvider): + """LiteLLM vision provider.""" + + async def analyze_image( + self, + image_data: bytes, + mime_type: str = "image/jpeg", + prompt: Optional[str] = None, + ) -> str: + litellm = _import_litellm() + + if not prompt: + prompt = ( + "Describe this image in detail. " + "If it's a diagram, flowchart, or table, explain the meaning and steps. " + "If it's a regular image, provide a concise description." + ) + + b64_image = base64.b64encode(image_data).decode("utf-8") + data_url = f"data:{mime_type};base64,{b64_image}" + + kwargs: dict = { + "model": self.config.model_id, + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + { + "type": "image_url", + "image_url": {"url": data_url, "detail": "low"}, + }, + ], + } + ], + "temperature": 0.2, + "drop_params": True, + } + if self.config.api_key: + kwargs["api_key"] = self.config.api_key + if self.config.base_url: + kwargs["api_base"] = self.config.base_url + + for attempt in range(3): + try: + response = await litellm.acompletion(**kwargs) + return response.choices[0].message.content or "" + except Exception as e: + logger.warning(f"LiteLLM Vision attempt {attempt + 1} failed: {e}") + if attempt < 2: + await asyncio.sleep(2) + return "" + + async def test_connection(self) -> tuple[bool, str]: + try: + tiny_png = ( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01" + b"\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde\x00" + b"\x00\x00\x0cIDATx\x9cc\xf8\x0f\x00\x00\x01\x01\x00" + b"\x05\x18\xd8N\x00\x00\x00\x00IEND\xaeB`\x82" + ) + await self.analyze_image(tiny_png, "image/png", "What is this?") + return True, f"OK — model={self.config.model_id}" + except Exception as e: + return False, f"LiteLLM Vision error: {e}" diff --git a/app/ai/registry.py b/app/ai/registry.py index e100dc40..0f36f5bc 100644 --- a/app/ai/registry.py +++ b/app/ai/registry.py @@ -40,36 +40,56 @@ # Provider class mappings — add new providers here # --------------------------------------------------------------------------- + def _get_embedding_class(provider: ProviderType) -> type[EmbeddingProvider]: if provider == ProviderType.GOOGLE: from app.ai.providers.google import GoogleEmbedding + return GoogleEmbedding elif provider == ProviderType.OPENAI: from app.ai.providers.openai_provider import OpenAIEmbedding + return OpenAIEmbedding + elif provider == ProviderType.LITELLM: + from app.ai.providers.litellm_provider import LiteLLMEmbedding + + return LiteLLMEmbedding raise ValueError(f"Unsupported embedding provider: {provider}") def _get_llm_class(provider: ProviderType) -> type[LLMProvider]: if provider == ProviderType.GOOGLE: from app.ai.providers.google import GoogleLLM + return GoogleLLM elif provider == ProviderType.OPENAI: from app.ai.providers.openai_provider import OpenAILLM + return OpenAILLM elif provider == ProviderType.ANTHROPIC: from app.ai.providers.anthropic_provider import AnthropicLLM + return AnthropicLLM + elif provider == ProviderType.LITELLM: + from app.ai.providers.litellm_provider import LiteLLMChat + + return LiteLLMChat raise ValueError(f"Unsupported LLM provider: {provider}") def _get_vision_class(provider: ProviderType) -> type[VisionProvider]: if provider == ProviderType.GOOGLE: from app.ai.providers.google import GoogleVision + return GoogleVision elif provider == ProviderType.OPENAI: from app.ai.providers.openai_provider import OpenAIVision + return OpenAIVision + elif provider == ProviderType.LITELLM: + from app.ai.providers.litellm_provider import LiteLLMVision + + return LiteLLMVision raise ValueError(f"Unsupported vision provider: {provider}") @@ -77,6 +97,7 @@ def _get_vision_class(provider: ProviderType) -> type[VisionProvider]: # Registry # --------------------------------------------------------------------------- + class ProviderRegistry: """ Resolves provider configs from DB and returns the correct implementation. @@ -313,6 +334,7 @@ async def _load_vision_config(self) -> ProviderConfig: "ollama": "Ollama", "voyage": "Voyage AI", "cohere": "Cohere", + "litellm": "LiteLLM", } diff --git a/pyproject.toml b/pyproject.toml index f5d9d794..f7d66ba8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,7 @@ dependencies = [ ] [project.optional-dependencies] +litellm = ["litellm>=1.80.0,<2.0"] dev = [ "pytest>=9.0.3", "pytest-asyncio>=1.2.0",