From 071404d11eaf09e01ff58aa84e1a435bc63ed3a5 Mon Sep 17 00:00:00 2001 From: RheagalFire Date: Sat, 9 May 2026 03:32:44 +0530 Subject: [PATCH 1/2] feat: add LiteLLM as LLM provider --- src/llm/plugins/litellm.py | 165 +++++++++++++++++++++++++++++++++++++ tests/llm/test_litellm.py | 161 ++++++++++++++++++++++++++++++++++++ 2 files changed, 326 insertions(+) create mode 100644 src/llm/plugins/litellm.py create mode 100644 tests/llm/test_litellm.py diff --git a/src/llm/plugins/litellm.py b/src/llm/plugins/litellm.py new file mode 100644 index 0000000000..441ca4b988 --- /dev/null +++ b/src/llm/plugins/litellm.py @@ -0,0 +1,165 @@ +"""LiteLLM plugin for OM1. + +Routes to 100+ LLM providers (OpenAI, Anthropic, Google, Azure, Bedrock, +Ollama, etc.) via the litellm SDK. No proxy server needed. + +Model strings use the ``provider/model`` format, e.g. +``anthropic/claude-sonnet-4-20250514``, ``azure/gpt-4o``, +``bedrock/anthropic.claude-3-haiku``, ``openai/gpt-4o``. + +See https://docs.litellm.ai/docs/providers for all supported models. +""" + +import logging +import time +import typing as T + +from pydantic import BaseModel, Field + +from llm import LLM, LLMConfig +from llm.function_schemas import convert_function_calls_to_actions +from llm.output_model import CortexOutputModel +from providers.avatar_llm_state_provider import AvatarLLMState +from providers.llm_history_manager import LLMHistoryManager + +R = T.TypeVar("R", bound=BaseModel) + + +class LiteLLMConfig(LLMConfig): + """LiteLLM-specific configuration.""" + + base_url: T.Optional[str] = Field( + default=None, + description="Optional base URL override for the LLM API endpoint", + ) + model: T.Optional[str] = Field( + default="openai/gpt-4o", + description="LiteLLM model string (e.g. anthropic/claude-sonnet-4-20250514)", + ) + + +class LiteLLM(LLM[R]): + """ + A LiteLLM-based Language Learning Model implementation. + + Routes to 100+ LLM providers through the litellm SDK using + ``litellm.acompletion()``. Supports function calling via the same + tool_calls interface as OpenAI. + """ + + def __init__( + self, + config: LiteLLMConfig, + available_actions: T.Optional[T.List] = None, + ): + super().__init__(config, available_actions) + + if not config.model: + self._config.model = "openai/gpt-4o" + + try: + import litellm as _litellm # noqa: F401 + except ImportError: + raise ImportError( + "litellm is required for this plugin. " + "Install with: pip install litellm" + ) + + import openai + + self._openai_client = openai.AsyncClient( + api_key=config.api_key or "unused", + base_url=config.base_url or "https://api.openai.com/v1", + ) + + self.history_manager = LLMHistoryManager(self._config, self._openai_client) + + @AvatarLLMState.trigger_thinking() + @LLMHistoryManager.update_history() + async def ask( + self, + prompt: str, + messages: T.Optional[T.List[T.Dict[str, str]]] = None, + ) -> T.Optional[R]: + """ + Send a prompt to the LLM via litellm and get a structured response. + + Parameters + ---------- + prompt : str + The input prompt to send to the model. + messages : List[Dict[str, str]], optional + List of message dictionaries to send to the model. + + Returns + ------- + R or None + Parsed response matching the output_model structure, or None if + parsing fails. + """ + import litellm as _litellm + + if messages is None: + messages = [] + try: + logging.info(f"LiteLLM input: {prompt}") + logging.info(f"LiteLLM messages: {messages}") + + self.io_provider.llm_start_time = time.time() + self.io_provider.set_llm_prompt(prompt) + + formatted_messages = [ + {"role": msg.get("role", "user"), "content": msg.get("content", "")} + for msg in messages + ] + formatted_messages.append({"role": "user", "content": prompt}) + + params: T.Dict[str, T.Any] = { + "model": self._config.model or "openai/gpt-4o", + "messages": formatted_messages, + "drop_params": True, + "timeout": self._config.timeout, + } + + if self._config.api_key: + params["api_key"] = self._config.api_key + if self._config.base_url: + params["api_base"] = self._config.base_url + + if self.function_schemas: + params["tools"] = self.function_schemas + params["tool_choice"] = "auto" + + response = await _litellm.acompletion(**params) + + if not response.choices: + logging.warning("LiteLLM API returned empty choices") + return None + + message = response.choices[0].message + self.io_provider.llm_end_time = time.time() + + if message.tool_calls: + logging.info(f"Received {len(message.tool_calls)} function calls") + logging.info(f"Function calls: {message.tool_calls}") + + function_call_data = [ + { + "function": { + "name": getattr(tc, "function").name, + "arguments": getattr(tc, "function").arguments, + } + } + for tc in message.tool_calls + ] + + actions = convert_function_calls_to_actions(function_call_data) + + result = CortexOutputModel(actions=actions) + return T.cast(R, result) + + return None + + except Exception as e: + logging.error(f"LiteLLM API error: {e}") + return None diff --git a/tests/llm/test_litellm.py b/tests/llm/test_litellm.py new file mode 100644 index 0000000000..e69fd8548e --- /dev/null +++ b/tests/llm/test_litellm.py @@ -0,0 +1,161 @@ +"""Unit tests for LiteLLM plugin. + +These tests verify the plugin file structure and litellm SDK interaction +without importing the full OM1 dependency chain (which requires zenoh). +""" + +import ast +import sys +import types +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import pytest + +PLUGIN_PATH = Path(__file__).resolve().parents[2] / "src" / "llm" / "plugins" / "litellm.py" + + +class TestLiteLLMPluginStructure: + """Verify the plugin file has the correct structure for OM1 auto-discovery.""" + + def _parse_ast(self): + return ast.parse(PLUGIN_PATH.read_text()) + + def test_plugin_file_exists(self): + assert PLUGIN_PATH.exists() + + def test_has_litellm_config_class(self): + tree = self._parse_ast() + classes = [n.name for n in ast.walk(tree) if isinstance(n, ast.ClassDef)] + assert "LiteLLMConfig" in classes + + def test_has_litellm_class(self): + tree = self._parse_ast() + classes = [n.name for n in ast.walk(tree) if isinstance(n, ast.ClassDef)] + assert "LiteLLM" in classes + + def test_litellm_class_inherits_llm(self): + tree = self._parse_ast() + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and node.name == "LiteLLM": + base_names = [] + for base in node.bases: + if isinstance(base, ast.Subscript) and isinstance(base.value, ast.Name): + base_names.append(base.value.id) + elif isinstance(base, ast.Name): + base_names.append(base.id) + assert "LLM" in base_names + return + pytest.fail("LiteLLM class not found") + + def test_has_ask_method(self): + tree = self._parse_ast() + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and node.name == "LiteLLM": + methods = [n.name for n in node.body if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))] + assert "ask" in methods + return + pytest.fail("LiteLLM class not found") + + def test_ask_is_async(self): + tree = self._parse_ast() + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and node.name == "LiteLLM": + for item in node.body: + if isinstance(item, ast.AsyncFunctionDef) and item.name == "ask": + return + pytest.fail("ask() is not async") + + def test_uses_drop_params_true(self): + src = PLUGIN_PATH.read_text() + assert "drop_params" in src + + def test_uses_litellm_acompletion(self): + src = PLUGIN_PATH.read_text() + assert "acompletion" in src + + def test_lazy_imports_litellm(self): + src = PLUGIN_PATH.read_text() + assert "import litellm" not in src.split("class")[0] + + +class TestLiteLLMSDKInteraction: + """Test litellm SDK calls directly (no OM1 deps needed).""" + + def test_acompletion_called_with_drop_params(self): + fake_litellm = types.ModuleType("litellm") + mock_msg = MagicMock(content="ok", tool_calls=None) + mock_choice = MagicMock(message=mock_msg, finish_reason="stop") + mock_resp = MagicMock(choices=[mock_choice]) + fake_litellm.acompletion = AsyncMock(return_value=mock_resp) + sys.modules["litellm"] = fake_litellm + + try: + import asyncio + + async def run(): + resp = await fake_litellm.acompletion( + model="anthropic/claude-sonnet-4-20250514", + messages=[{"role": "user", "content": "hi"}], + drop_params=True, + ) + return resp + + asyncio.run(run()) + kwargs = fake_litellm.acompletion.call_args.kwargs + assert kwargs["drop_params"] is True + assert kwargs["model"] == "anthropic/claude-sonnet-4-20250514" + finally: + del sys.modules["litellm"] + + def test_acompletion_forwards_api_key(self): + fake_litellm = types.ModuleType("litellm") + mock_msg = MagicMock(content="ok", tool_calls=None) + mock_resp = MagicMock(choices=[MagicMock(message=mock_msg)]) + fake_litellm.acompletion = AsyncMock(return_value=mock_resp) + sys.modules["litellm"] = fake_litellm + + try: + import asyncio + + async def run(): + await fake_litellm.acompletion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "hi"}], + api_key="sk-test", + drop_params=True, + ) + + asyncio.run(run()) + assert fake_litellm.acompletion.call_args.kwargs["api_key"] == "sk-test" + finally: + del sys.modules["litellm"] + + def test_acompletion_handles_tool_calls(self): + fake_litellm = types.ModuleType("litellm") + mock_tc = MagicMock() + mock_tc.function.name = "move" + mock_tc.function.arguments = '{"direction": "forward"}' + mock_msg = MagicMock(content=None, tool_calls=[mock_tc]) + mock_resp = MagicMock(choices=[MagicMock(message=mock_msg)]) + fake_litellm.acompletion = AsyncMock(return_value=mock_resp) + sys.modules["litellm"] = fake_litellm + + try: + import asyncio + + async def run(): + resp = await fake_litellm.acompletion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "move"}], + tools=[{"type": "function"}], + tool_choice="auto", + drop_params=True, + ) + return resp + + resp = asyncio.run(run()) + tc = resp.choices[0].message.tool_calls[0] + assert tc.function.name == "move" + finally: + del sys.modules["litellm"] From f3e52956866a979469987080cad7f87cf666c447 Mon Sep 17 00:00:00 2001 From: RheagalFire Date: Sat, 9 May 2026 03:40:31 +0530 Subject: [PATCH 2/2] fix: add litellm optional dep, format with black --- pyproject.toml | 1 + src/llm/plugins/litellm.py | 8 ++------ 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 7660bbfcf0..7d35622efe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,7 @@ dependencies = [ dds = [ "cyclonedds==0.10.2" ] +litellm = ["litellm>=1.60.0,<2.0.0"] macos = ["osascript"] [dependency-groups] diff --git a/src/llm/plugins/litellm.py b/src/llm/plugins/litellm.py index 441ca4b988..80368d5cd0 100644 --- a/src/llm/plugins/litellm.py +++ b/src/llm/plugins/litellm.py @@ -60,10 +60,7 @@ def __init__( try: import litellm as _litellm # noqa: F401 except ImportError: - raise ImportError( - "litellm is required for this plugin. " - "Install with: pip install litellm" - ) + raise ImportError("litellm is required for this plugin. " "Install with: pip install litellm") import openai @@ -109,8 +106,7 @@ async def ask( self.io_provider.set_llm_prompt(prompt) formatted_messages = [ - {"role": msg.get("role", "user"), "content": msg.get("content", "")} - for msg in messages + {"role": msg.get("role", "user"), "content": msg.get("content", "")} for msg in messages ] formatted_messages.append({"role": "user", "content": prompt})