diff --git a/mle/model/__init__.py b/mle/model/__init__.py index fab71d84..78272716 100644 --- a/mle/model/__init__.py +++ b/mle/model/__init__.py @@ -5,6 +5,7 @@ from .openai import * from .gemini import * from .vllm import * +from .litellm import * from mle.utils import get_config @@ -16,6 +17,7 @@ MODEL_DEEPSEEK = 'DeepSeek' MODEL_GEMINI = 'Gemini' MODEL_VLLM = 'vLLM' +MODEL_LITELLM = 'LiteLLM' class ObservableModel: @@ -74,6 +76,8 @@ def load_model(project_dir: str, model_name: str=None, observable=True): model = GeminiModel(api_key=config['api_key'], model=model_name) if config['platform'] == MODEL_VLLM: model = vLLMModel(base_url=config.get('base_url', 'http://localhost:8000/v1'), model=model_name) + if config['platform'] == MODEL_LITELLM: + model = LiteLLMModel(api_key=config.get('api_key'), base_url=config.get('base_url'), model=model_name) if observable: return ObservableModel(model) diff --git a/mle/model/litellm.py b/mle/model/litellm.py new file mode 100644 index 00000000..0a7c3685 --- /dev/null +++ b/mle/model/litellm.py @@ -0,0 +1,167 @@ +import os +import json +import importlib.util +from typing import List, Dict, Any, Optional + +from mle.model.common import Model +from mle.function import SEARCH_FUNCTIONS, get_function, process_function_name + + +class LiteLLMModel(Model): + """ + LiteLLM model implementation using OpenAI-compatible API. + + LiteLLM is an AI gateway/proxy that provides a unified OpenAI-compatible + interface to 100+ LLM providers (OpenAI, Anthropic, Azure, Bedrock, + Vertex AI, Mistral, Cohere, etc.). + + See: https://docs.litellm.ai/ + """ + + def __init__(self, api_key: Optional[str] = None, + base_url: Optional[str] = None, + model: Optional[str] = None, + temperature: float = 0.7) -> None: + """Initialize the LiteLLM model. + + Args: + api_key: The LiteLLM proxy API key (master key or virtual key). + base_url: The URL of the LiteLLM proxy server. + model: The model name (e.g., gpt-4o, claude-3-5-sonnet). + temperature: The sampling temperature. + """ + super().__init__() + + dependency = "openai" + spec = importlib.util.find_spec(dependency) + if spec is not None: + self.openai = importlib.import_module(dependency).OpenAI + else: + raise ImportError( + "OpenAI package not found. Please install it using: " + "pip install openai" + ) + + self.model = model if model else "gpt-4o" + self.model_type = 'LiteLLM' + self.temperature = temperature + self.client = self.openai( + api_key=api_key or os.getenv("LITELLM_API_KEY", ""), + base_url=base_url or os.getenv("LITELLM_BASE_URL", + "http://localhost:4000/v1"), + timeout=60.0, + max_retries=2, + ) + self.func_call_history = [] + + def _convert_functions_to_tools(self, functions): + """ + Convert OpenAI-style functions to tools format. + """ + tools = [] + for func in functions: + tool = { + "type": "function", + "function": { + "name": func["name"], + "description": func.get("description", ""), + "parameters": func["parameters"], + }, + } + tools.append(tool) + return tools + + def query(self, chat_history: List[Dict[str, Any]], **kwargs) -> str: + """Query the LLM model. + + Args: + chat_history: The context (chat history). + **kwargs: Additional parameters for the API call. + + Returns: + Model's response as string. + """ + functions = kwargs.get("functions", None) + tools = self._convert_functions_to_tools(functions) if functions else None + parameters = kwargs + completion = self.client.chat.completions.create( + model=self.model, + messages=chat_history, + temperature=self.temperature, + stream=False, + tools=tools, + **parameters, + ) + + resp = completion.choices[0].message + if resp.tool_calls: + for tool_call in resp.tool_calls: + chat_history.append({ + "role": "assistant", + "content": '', + "tool_calls": [tool_call], + "prefix": False + }) + function_name = process_function_name(tool_call.function.name) + arguments = json.loads(tool_call.function.arguments) + print("[MLE FUNC CALL]: ", function_name) + self.func_call_history.append({ + "name": function_name, + "arguments": arguments + }) + search_attempts = [ + item for item in self.func_call_history + if item['name'] in SEARCH_FUNCTIONS + ] + if len(search_attempts) > 3: + parameters['tool_choice'] = "none" + result = get_function(function_name)(**arguments) + chat_history.append({ + "role": "tool", + "content": result, + "name": function_name, + "tool_call_id": tool_call.id + }) + return self.query(chat_history, **parameters) + else: + return resp.content + + def stream(self, chat_history: List[Dict[str, Any]], **kwargs) -> str: + """Stream the output from the LLM model. + + Args: + chat_history: The context (chat history). + **kwargs: Additional parameters for the API call. + + Yields: + Chunks of the model's response. + """ + arguments = "" + function_name = "" + for chunk in self.client.chat.completions.create( + model=self.model, + messages=chat_history, + temperature=self.temperature, + stream=True, + **kwargs, + ): + if chunk.choices[0].delta.tool_calls: + tool_call = chunk.choices[0].delta.tool_calls[0] + if tool_call.function.name: + chat_history.append({ + "role": "assistant", + "content": '', + "tool_calls": [tool_call], + "prefix": False + }) + function_name = process_function_name(tool_call.function.name) + arguments = json.loads(tool_call.function.arguments) + result = get_function(function_name)(**arguments) + chat_history.append({ + "role": "tool", + "content": result, + "name": function_name + }) + yield from self.stream(chat_history, **kwargs) + else: + yield chunk.choices[0].delta.content diff --git a/tests/test_litellm_model.py b/tests/test_litellm_model.py new file mode 100644 index 00000000..290eb712 --- /dev/null +++ b/tests/test_litellm_model.py @@ -0,0 +1,239 @@ +""" +Tests for the LiteLLM model provider. +""" +import os +import sys +import types +import unittest +from unittest.mock import MagicMock, patch +from types import SimpleNamespace + +# Stub mle.function to avoid pulling pandas and other heavy deps +_fake_func = types.ModuleType("mle.function") +_fake_func.SEARCH_FUNCTIONS = [] +_fake_func.get_function = lambda x: (lambda **kw: "mock_result") +_fake_func.process_function_name = lambda x: x +sys.modules["mle.function"] = _fake_func + +from mle.model.litellm import LiteLLMModel + + +class TestLiteLLMModelInit(unittest.TestCase): + + @patch.dict(os.environ, {}, clear=True) + @patch("importlib.util.find_spec", return_value=True) + @patch("importlib.import_module") + def test_init_with_explicit_params(self, mock_import, mock_spec): + mock_openai_cls = MagicMock() + mock_import.return_value = SimpleNamespace(OpenAI=mock_openai_cls) + + model = LiteLLMModel( + api_key="sk-test-key", + base_url="http://myproxy:4000/v1", + model="claude-3-5-sonnet", + temperature=0.3 + ) + + mock_openai_cls.assert_called_once_with( + api_key="sk-test-key", + base_url="http://myproxy:4000/v1", + timeout=60.0, + max_retries=2, + ) + self.assertEqual(model.model, "claude-3-5-sonnet") + self.assertEqual(model.model_type, "LiteLLM") + self.assertEqual(model.temperature, 0.3) + + @patch.dict(os.environ, { + "LITELLM_API_KEY": "sk-from-env", + "LITELLM_BASE_URL": "http://env-proxy:8000/v1" + }) + @patch("importlib.util.find_spec", return_value=True) + @patch("importlib.import_module") + def test_init_from_env_vars(self, mock_import, mock_spec): + mock_openai_cls = MagicMock() + mock_import.return_value = SimpleNamespace(OpenAI=mock_openai_cls) + + model = LiteLLMModel() + + mock_openai_cls.assert_called_once_with( + api_key="sk-from-env", + base_url="http://env-proxy:8000/v1", + timeout=60.0, + max_retries=2, + ) + self.assertEqual(model.model, "gpt-4o") + + @patch.dict(os.environ, {}, clear=True) + @patch("importlib.util.find_spec", return_value=True) + @patch("importlib.import_module") + def test_init_defaults_when_no_config(self, mock_import, mock_spec): + mock_openai_cls = MagicMock() + mock_import.return_value = SimpleNamespace(OpenAI=mock_openai_cls) + + model = LiteLLMModel() + + mock_openai_cls.assert_called_once_with( + api_key="", + base_url="http://localhost:4000/v1", + timeout=60.0, + max_retries=2, + ) + + @patch("importlib.util.find_spec", return_value=None) + def test_init_raises_when_openai_not_installed(self, mock_spec): + with self.assertRaises(ImportError) as ctx: + LiteLLMModel(api_key="key") + self.assertIn("OpenAI package not found", str(ctx.exception)) + + +def _make_model(): + with patch("importlib.util.find_spec", return_value=True), \ + patch("importlib.import_module") as mock_import: + mock_openai_cls = MagicMock() + mock_import.return_value = SimpleNamespace(OpenAI=mock_openai_cls) + model = LiteLLMModel(api_key="test", base_url="http://proxy:4000/v1") + return model + + +class TestLiteLLMModelQuery(unittest.TestCase): + + def test_simple_query_returns_content(self): + model = _make_model() + mock_resp = SimpleNamespace( + content="Paris is the capital of France.", + tool_calls=None + ) + model.client.chat.completions.create = MagicMock( + return_value=SimpleNamespace(choices=[SimpleNamespace(message=mock_resp)]) + ) + + result = model.query([{"role": "user", "content": "What is the capital of France?"}]) + + self.assertEqual(result, "Paris is the capital of France.") + call_kwargs = model.client.chat.completions.create.call_args + self.assertEqual(call_kwargs.kwargs["model"], "gpt-4o") + self.assertFalse(call_kwargs.kwargs["stream"]) + + def test_query_with_none_content(self): + model = _make_model() + mock_resp = SimpleNamespace(content=None, tool_calls=None) + model.client.chat.completions.create = MagicMock( + return_value=SimpleNamespace(choices=[SimpleNamespace(message=mock_resp)]) + ) + + result = model.query([{"role": "user", "content": "test"}]) + self.assertIsNone(result) + + def test_query_auth_error_propagates(self): + model = _make_model() + from openai import AuthenticationError + model.client.chat.completions.create = MagicMock( + side_effect=AuthenticationError( + message="Invalid API key", + response=MagicMock(status_code=401), + body=None + ) + ) + + with self.assertRaises(AuthenticationError): + model.query([{"role": "user", "content": "test"}]) + + def test_query_rate_limit_error_propagates(self): + model = _make_model() + from openai import RateLimitError + model.client.chat.completions.create = MagicMock( + side_effect=RateLimitError( + message="Rate limit exceeded", + response=MagicMock(status_code=429), + body=None + ) + ) + + with self.assertRaises(RateLimitError): + model.query([{"role": "user", "content": "test"}]) + + def test_query_timeout_propagates(self): + model = _make_model() + from openai import APITimeoutError + model.client.chat.completions.create = MagicMock( + side_effect=APITimeoutError(request=MagicMock()) + ) + + with self.assertRaises(APITimeoutError): + model.query([{"role": "user", "content": "test"}]) + + def test_query_empty_choices_raises(self): + model = _make_model() + model.client.chat.completions.create = MagicMock( + return_value=SimpleNamespace(choices=[]) + ) + + with self.assertRaises(IndexError): + model.query([{"role": "user", "content": "test"}]) + + +class TestLiteLLMModelStream(unittest.TestCase): + + def test_stream_yields_content_chunks(self): + model = _make_model() + chunks = [ + SimpleNamespace(choices=[SimpleNamespace( + delta=SimpleNamespace(content="Hello", tool_calls=None) + )]), + SimpleNamespace(choices=[SimpleNamespace( + delta=SimpleNamespace(content=" world", tool_calls=None) + )]), + ] + model.client.chat.completions.create = MagicMock(return_value=iter(chunks)) + + result = list(model.stream([{"role": "user", "content": "test"}])) + self.assertEqual(result, ["Hello", " world"]) + + def test_stream_handles_none_content_chunks(self): + model = _make_model() + chunks = [ + SimpleNamespace(choices=[SimpleNamespace( + delta=SimpleNamespace(content=None, tool_calls=None) + )]), + SimpleNamespace(choices=[SimpleNamespace( + delta=SimpleNamespace(content="data", tool_calls=None) + )]), + ] + model.client.chat.completions.create = MagicMock(return_value=iter(chunks)) + + result = list(model.stream([{"role": "user", "content": "test"}])) + self.assertEqual(result, [None, "data"]) + + +class TestLiteLLMModelToolConversion(unittest.TestCase): + + def test_convert_functions_to_tools(self): + model = _make_model() + functions = [{ + "name": "get_weather", + "description": "Get current weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"] + } + }] + tools = model._convert_functions_to_tools(functions) + + self.assertEqual(len(tools), 1) + self.assertEqual(tools[0]["type"], "function") + self.assertEqual(tools[0]["function"]["name"], "get_weather") + + def test_convert_empty_functions(self): + model = _make_model() + self.assertEqual(model._convert_functions_to_tools([]), []) + + def test_convert_function_without_description(self): + model = _make_model() + tools = model._convert_functions_to_tools([{"name": "test", "parameters": {"type": "object"}}]) + self.assertEqual(tools[0]["function"]["description"], "") + + +if __name__ == "__main__": + unittest.main()