feat: add litellm provider

pyproject.toml

@@ -71,6 +71,7 @@ postgres = ["pgvector>=0.3.4", "sqlalchemy[postgresql-psycopgbinary]>=2.0.36"]
 langgraph = ["langgraph>=0.0.10", "langchain-core>=0.1.0"]
 claude = ["claude-agent-sdk>=0.1.24"]
 lazyllm = ["lazyllm>=0.7.3"]
+litellm = ["litellm>=1.80.0,<1.87.0"]
 # Rich document ingestion (PDF, Word, PowerPoint, Excel, ...) via MarkItDown.
 document = ["markitdown[docx,pptx,xlsx,xls,pdf]>=0.1.0"]
 
@@ -81,7 +82,7 @@ document = ["markitdown[docx,pptx,xlsx,xls,pdf]>=0.1.0"]
 
 [tool.deptry.per_rule_ignores]
 # Optional dependencies used in examples/ or imported lazily behind extras.
-DEP002 = ["claude-agent-sdk", "markitdown"]
+DEP002 = ["claude-agent-sdk", "litellm", "markitdown"]
 
 [tool.mypy]
 files = ["src", "tests"]

src/memu/app/settings.py

@@ -146,6 +146,13 @@ def set_provider_defaults(self) -> "LLMConfig":
                 self.api_key = api_key
             if self.chat_model == "gpt-5.4-mini":
                 self.chat_model = chat_model
+        elif self.provider == "litellm":
+            if self.client_backend == "sdk":
+                self.client_backend = "litellm"
+            if self.base_url == "https://api.openai.com/v1":
+                self.base_url = "http://localhost:4000"
+            if self.api_key == "OPENAI_API_KEY":
+                self.api_key = "LITELLM_API_KEY"
         return self
 
 

src/memu/embedding/gateway.py

@@ -47,6 +47,18 @@ def _build_lazyllm_client(cfg: EmbeddingConfig) -> Any:
     )
 
 
+def _build_litellm_client(cfg: EmbeddingConfig) -> Any:
+    from memu.llm.litellm_sdk import LiteLLMSDKClient
+
+    return LiteLLMSDKClient(
+        chat_model="",
+        embed_model=cfg.embed_model,
+        api_key=cfg.api_key,
+        api_base=cfg.base_url if cfg.base_url != "https://api.openai.com/v1" else None,
+        embed_batch_size=cfg.embed_batch_size,
+    )
+
+
 def _build_anthropic_client(cfg: EmbeddingConfig) -> Any:
     msg = (
         "Anthropic does not provide an embeddings API. Configure an embedding "
@@ -59,6 +71,7 @@ def _build_anthropic_client(cfg: EmbeddingConfig) -> Any:
 EMBEDDING_CLIENT_BUILDERS: dict[str, Callable[[EmbeddingConfig], Any]] = {
     "sdk": _build_sdk_client,
     "httpx": _build_httpx_client,
+    "litellm": _build_litellm_client,
     "lazyllm_backend": _build_lazyllm_client,
     "anthropic": _build_anthropic_client,
 }

src/memu/llm/backends/__init__.py

@@ -4,6 +4,7 @@
 from memu.llm.backends.doubao import DoubaoLLMBackend
 from memu.llm.backends.grok import GrokBackend
 from memu.llm.backends.kimi import KimiLLMBackend
+from memu.llm.backends.litellm import LiteLLMBackend
 from memu.llm.backends.minimax import MiniMaxLLMBackend
 from memu.llm.backends.openai import OpenAILLMBackend
 from memu.llm.backends.openrouter import OpenRouterLLMBackend
@@ -15,6 +16,7 @@
     "GrokBackend",
     "KimiLLMBackend",
     "LLMBackend",
+    "LiteLLMBackend",
     "MiniMaxLLMBackend",
     "OpenAILLMBackend",
     "OpenRouterLLMBackend",

src/memu/llm/backends/litellm.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from memu.llm.backends.openai import OpenAILLMBackend
+
+
+class LiteLLMBackend(OpenAILLMBackend):
+    """Backend for LiteLLM AI gateway proxy (OpenAI-compatible)."""
+
+    name = "litellm"

src/memu/llm/gateway.py

@@ -50,6 +50,18 @@ def _build_httpx_client(cfg: LLMConfig) -> Any:
     )
 
 
+def _build_litellm_client(cfg: LLMConfig) -> Any:
+    from memu.llm.litellm_sdk import LiteLLMSDKClient
+
+    return LiteLLMSDKClient(
+        chat_model=cfg.chat_model,
+        embed_model=cfg.embed_model,
+        api_key=cfg.api_key,
+        api_base=cfg.base_url if cfg.base_url != "http://localhost:4000" else None,
+        embed_batch_size=cfg.embed_batch_size,
+    )
+
+
 def _build_lazyllm_client(cfg: LLMConfig) -> Any:
     from memu.llm.lazyllm_client import LazyLLMClient
 
@@ -72,6 +84,7 @@ def _build_lazyllm_client(cfg: LLMConfig) -> Any:
     "sdk": _build_sdk_client,
     "anthropic": _build_anthropic_client,
     "httpx": _build_httpx_client,
+    "litellm": _build_litellm_client,
     "lazyllm_backend": _build_lazyllm_client,
 }
 

src/memu/llm/http_client.py

@@ -15,6 +15,7 @@
 from memu.llm.backends.doubao import DoubaoLLMBackend
 from memu.llm.backends.grok import GrokBackend
 from memu.llm.backends.kimi import KimiLLMBackend
+from memu.llm.backends.litellm import LiteLLMBackend
 from memu.llm.backends.minimax import MiniMaxLLMBackend
 from memu.llm.backends.openai import OpenAILLMBackend
 from memu.llm.backends.openrouter import OpenRouterLLMBackend
@@ -30,6 +31,7 @@ def _load_proxy() -> str | None:
     OpenAILLMBackend.name: OpenAILLMBackend,
     ClaudeLLMBackend.name: ClaudeLLMBackend,
     GrokBackend.name: GrokBackend,
+    LiteLLMBackend.name: LiteLLMBackend,
     DeepSeekLLMBackend.name: DeepSeekLLMBackend,
     KimiLLMBackend.name: KimiLLMBackend,
     MiniMaxLLMBackend.name: MiniMaxLLMBackend,

src/memu/llm/litellm_sdk.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+import base64
+import logging
+from pathlib import Path
+from typing import Any, cast
+
+logger = logging.getLogger(__name__)
+
+
+class LiteLLMSDKClient:
+    """LLM client using the LiteLLM Python SDK for 100+ provider support."""
+
+    def __init__(
+        self,
+        *,
+        chat_model: str,
+        embed_model: str,
+        api_key: str | None = None,
+        api_base: str | None = None,
+        embed_batch_size: int = 1,
+    ):
+        self.chat_model = chat_model
+        self.embed_model = embed_model
+        self.api_key = api_key or None
+        self.api_base = api_base or None
+        self.embed_batch_size = embed_batch_size
+
+    async def chat(
+        self,
+        prompt: str,
+        *,
+        max_tokens: int | None = None,
+        system_prompt: str | None = None,
+        temperature: float = 0.2,
+    ) -> tuple[str, dict[str, Any]]:
+        import litellm
+
+        messages: list[dict[str, str]] = []
+        if system_prompt is not None:
+            messages.append({"role": "system", "content": system_prompt})
+        messages.append({"role": "user", "content": prompt})
+
+        kwargs: dict[str, Any] = {
+            "model": self.chat_model,
+            "messages": messages,
+            "temperature": temperature,
+            "drop_params": True,
+        }
+        if max_tokens is not None:
+            kwargs["max_tokens"] = max_tokens
+        if self.api_key:
+            kwargs["api_key"] = self.api_key
+        if self.api_base:
+            kwargs["api_base"] = self.api_base
+
+        response = await litellm.acompletion(**kwargs)
+        data = response.model_dump()
+        content = data["choices"][0]["message"]["content"] or ""
+        logger.debug("LiteLLM chat response: %s", data)
+        return content, data
+
+    async def summarize(
+        self,
+        text: str,
+        *,
+        max_tokens: int | None = None,
+        system_prompt: str | None = None,
+    ) -> tuple[str, dict[str, Any]]:
+        prompt = system_prompt or "Summarize the text in one short paragraph."
+        return await self.chat(
+            text,
+            max_tokens=max_tokens,
+            system_prompt=prompt,
+            temperature=0.2,
+        )
+
+    async def vision(
+        self,
+        prompt: str,
+        image_path: str,
+        *,
+        max_tokens: int | None = None,
+        system_prompt: str | None = None,
+    ) -> tuple[str, dict[str, Any]]:
+        import litellm
+
+        image_data = Path(image_path).read_bytes()
+        base64_image = base64.b64encode(image_data).decode("utf-8")
+
+        suffix = Path(image_path).suffix.lower()
+        mime_type = {
+            ".jpg": "image/jpeg",
+            ".jpeg": "image/jpeg",
+            ".png": "image/png",
+            ".gif": "image/gif",
+            ".webp": "image/webp",
+        }.get(suffix, "image/jpeg")
+
+        messages: list[dict[str, Any]] = []
+        if system_prompt:
+            messages.append({"role": "system", "content": system_prompt})
+
+        messages.append({
+            "role": "user",
+            "content": [
+                {"type": "text", "text": prompt},
+                {"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{base64_image}"}},
+            ],
+        })
+
+        kwargs: dict[str, Any] = {
+            "model": self.chat_model,
+            "messages": messages,
+            "temperature": 0.2,
+            "drop_params": True,
+        }
+        if max_tokens is not None:
+            kwargs["max_tokens"] = max_tokens
+        if self.api_key:
+            kwargs["api_key"] = self.api_key
+        if self.api_base:
+            kwargs["api_base"] = self.api_base
+
+        response = await litellm.acompletion(**kwargs)
+        data = response.model_dump()
+        content = data["choices"][0]["message"]["content"] or ""
+        logger.debug("LiteLLM vision response: %s", data)
+        return content, data
+
+    async def embed(self, inputs: list[str]) -> tuple[list[list[float]], dict[str, Any] | None]:
+        import litellm
+
+        kwargs: dict[str, Any] = {"model": self.embed_model, "drop_params": True}
+        if self.api_key:
+            kwargs["api_key"] = self.api_key
+        if self.api_base:
+            kwargs["api_base"] = self.api_base
+
+        if len(inputs) <= self.embed_batch_size:
+            response = await litellm.aembedding(input=inputs, **kwargs)
+            data = response.model_dump()
+            return [cast(list[float], d["embedding"]) for d in data["data"]], data
+
+        all_embeddings: list[list[float]] = []
+        last_data: dict[str, Any] | None = None
+        for idx in range(0, len(inputs), self.embed_batch_size):
+            batch = inputs[idx : idx + self.embed_batch_size]
+            response = await litellm.aembedding(input=batch, **kwargs)
+            data = response.model_dump()
+            all_embeddings.extend([cast(list[float], d["embedding"]) for d in data["data"]])
+            last_data = data
+
+        return all_embeddings, last_data