feat: add LiteLLM as LLM provider (#21)

Author: RheagalFire

ifixai/providers/litellm.py

@@ -0,0 +1,78 @@
+"""LiteLLM provider for iFixAi.
+
+Routes to 100+ LLM providers (OpenAI, Anthropic, Google, Azure, Bedrock,
+Ollama, etc.) via the litellm SDK. No proxy server needed.
+
+Install: pip install ifixai[litellm]
+
+See https://docs.litellm.ai/docs/providers for all supported models.
+"""
+
+import litellm as _litellm
+
+from ifixai.core.types import ChatMessage, ProviderConfig
+from ifixai.providers.base import (
+    ChatProvider,
+    ProviderConnectionError,
+    ProviderResponseError,
+    ProviderTimeoutError,
+)
+
+DEFAULT_MODEL = "openai/gpt-4o"
+
+
+class LiteLLMProvider(ChatProvider):
+    def __init__(self) -> None:
+        pass
+
+    async def send_message(
+        self,
+        messages: list[ChatMessage],
+        config: ProviderConfig,
+    ) -> str:
+        model = config.model or DEFAULT_MODEL
+        formatted_messages = [{"role": m.role, "content": m.content} for m in messages]
+
+        params = {
+            "model": model,
+            "messages": formatted_messages,
+            "drop_params": True,
+            "timeout": float(config.timeout),
+        }
+        if config.api_key:
+            params["api_key"] = config.api_key
+        if config.endpoint:
+            params["api_base"] = config.endpoint
+
+        try:
+            response = await _litellm.acompletion(**params)
+
+            content = response.choices[0].message.content
+            if not content:
+                raise ProviderResponseError(
+                    provider="litellm",
+                    endpoint=config.endpoint or "default",
+                    details="Empty content in response",
+                )
+            return content
+
+        except ProviderResponseError:
+            raise
+        except TimeoutError as exc:
+            raise ProviderTimeoutError(
+                provider="litellm",
+                endpoint=config.endpoint or "default",
+                details=str(exc),
+            ) from exc
+        except ConnectionError as exc:
+            raise ProviderConnectionError(
+                provider="litellm",
+                endpoint=config.endpoint or "default",
+                details=str(exc),
+            ) from exc
+        except Exception as exc:
+            raise ProviderResponseError(
+                provider="litellm",
+                endpoint=config.endpoint or "default",
+                details=str(exc),
+            ) from exc

ifixai/providers/resolver.py

@@ -43,6 +43,11 @@
 except ImportError:
     OpenRouterProvider = None
 
+try:
+    from ifixai.providers.litellm import LiteLLMProvider
+except ImportError:
+    LiteLLMProvider = None
+
 
 REGISTERED_PROVIDERS: tuple[str, ...] = (
     "http",
@@ -55,6 +60,7 @@
     "bedrock",
     "huggingface",
     "langchain",
+    "litellm",
 )
 
 _MOCK_FIXTURE_PATH = Path(__file__).parent.parent / "fixtures" / "governance" / "mock.yaml"
@@ -71,6 +77,7 @@
         "bedrock": BedrockProvider,
         "huggingface": HuggingFaceProvider,
         "langchain": LangChainProvider,
+        "litellm": LiteLLMProvider,
     }.items()
     if cls is not None
 }
@@ -139,6 +146,7 @@ def resolve_provider(provider: Union[str, object]) -> object:
     "bedrock": ("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"),
     "huggingface": ("HUGGINGFACE_API_TOKEN", "HF_TOKEN"),
     "openrouter": ("OPENROUTER_API_KEY",),
+    "litellm": ("LITELLM_API_KEY",),
 }
 
 _JUDGE_PREFERENCE_ORDER: tuple[str, ...] = (

pyproject.toml

@@ -43,12 +43,14 @@ anthropic = ["anthropic>=0.18"]
 azure = ["openai>=1.0"]
 bedrock = ["boto3>=1.28"]
 huggingface = ["huggingface-hub>=0.20"]
+litellm = ["litellm>=1.60.0,<2.0.0"]
 all = [
     "openai>=1.0",
     "google-generativeai>=0.3",
     "anthropic>=0.18",
     "boto3>=1.28",
     "huggingface-hub>=0.20",
+    "litellm>=1.60.0,<2.0.0",
 ]
 dev = [
     "ruff>=0.4",

tests/__init__.py

@@ -0,0 +1,171 @@
+"""Tests for LiteLLM provider."""
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from ifixai.core.types import ChatMessage, ProviderConfig
+from ifixai.providers.base import (
+    ChatProvider,
+    ProviderResponseError,
+    ProviderTimeoutError,
+)
+from ifixai.providers.litellm import LiteLLMProvider
+
+
+def _make_config(**overrides):
+    defaults = {
+        "provider": "litellm",
+        "model": "openai/gpt-4o",
+        "api_key": "sk-test",
+    }
+    defaults.update(overrides)
+    return ProviderConfig(**defaults)
+
+
+class TestLiteLLMProviderInit:
+    def test_extends_chat_provider(self):
+        assert issubclass(LiteLLMProvider, ChatProvider)
+
+    def test_instantiates_without_args(self):
+        p = LiteLLMProvider()
+        assert p is not None
+
+
+class TestSendMessage:
+    @pytest.mark.asyncio
+    @patch("ifixai.providers.litellm._litellm")
+    async def test_calls_acompletion_with_drop_params(self, mock_litellm):
+        mock_msg = MagicMock(content="test response")
+        mock_litellm.acompletion = AsyncMock(
+            return_value=MagicMock(choices=[MagicMock(message=mock_msg)])
+        )
+
+        p = LiteLLMProvider()
+        config = _make_config()
+        messages = [ChatMessage(role="user", content="hello")]
+        result = await p.send_message(messages, config)
+
+        assert result == "test response"
+        kwargs = mock_litellm.acompletion.call_args.kwargs
+        assert kwargs["drop_params"] is True
+        assert kwargs["model"] == "openai/gpt-4o"
+        assert kwargs["api_key"] == "sk-test"
+
+    @pytest.mark.asyncio
+    @patch("ifixai.providers.litellm._litellm")
+    async def test_omits_api_key_when_empty(self, mock_litellm):
+        mock_msg = MagicMock(content="ok")
+        mock_litellm.acompletion = AsyncMock(
+            return_value=MagicMock(choices=[MagicMock(message=mock_msg)])
+        )
+
+        p = LiteLLMProvider()
+        config = _make_config(api_key="")
+        messages = [ChatMessage(role="user", content="hi")]
+        await p.send_message(messages, config)
+
+        kwargs = mock_litellm.acompletion.call_args.kwargs
+        assert "api_key" not in kwargs
+
+    @pytest.mark.asyncio
+    @patch("ifixai.providers.litellm._litellm")
+    async def test_formats_messages_correctly(self, mock_litellm):
+        mock_msg = MagicMock(content="ok")
+        mock_litellm.acompletion = AsyncMock(
+            return_value=MagicMock(choices=[MagicMock(message=mock_msg)])
+        )
+
+        p = LiteLLMProvider()
+        config = _make_config()
+        messages = [
+            ChatMessage(role="system", content="be helpful"),
+            ChatMessage(role="user", content="hi"),
+        ]
+        await p.send_message(messages, config)
+
+        kwargs = mock_litellm.acompletion.call_args.kwargs
+        assert kwargs["messages"] == [
+            {"role": "system", "content": "be helpful"},
+            {"role": "user", "content": "hi"},
+        ]
+
+    @pytest.mark.asyncio
+    @patch("ifixai.providers.litellm._litellm")
+    async def test_uses_default_model_when_none(self, mock_litellm):
+        mock_msg = MagicMock(content="ok")
+        mock_litellm.acompletion = AsyncMock(
+            return_value=MagicMock(choices=[MagicMock(message=mock_msg)])
+        )
+
+        p = LiteLLMProvider()
+        config = _make_config(model=None)
+        messages = [ChatMessage(role="user", content="hi")]
+        await p.send_message(messages, config)
+
+        kwargs = mock_litellm.acompletion.call_args.kwargs
+        assert kwargs["model"] == "openai/gpt-4o"
+
+
+class TestEdgeCases:
+    @pytest.mark.asyncio
+    @patch("ifixai.providers.litellm._litellm")
+    async def test_empty_content_raises_response_error(self, mock_litellm):
+        mock_msg = MagicMock(content="")
+        mock_litellm.acompletion = AsyncMock(
+            return_value=MagicMock(choices=[MagicMock(message=mock_msg)])
+        )
+
+        p = LiteLLMProvider()
+        config = _make_config()
+        messages = [ChatMessage(role="user", content="hi")]
+
+        with pytest.raises(ProviderResponseError):
+            await p.send_message(messages, config)
+
+    @pytest.mark.asyncio
+    @patch("ifixai.providers.litellm._litellm")
+    async def test_api_error_raises_response_error(self, mock_litellm):
+        mock_litellm.acompletion = AsyncMock(side_effect=Exception("401 Unauthorized"))
+
+        p = LiteLLMProvider()
+        config = _make_config()
+        messages = [ChatMessage(role="user", content="hi")]
+
+        with pytest.raises(ProviderResponseError, match="401"):
+            await p.send_message(messages, config)
+
+    @pytest.mark.asyncio
+    @patch("ifixai.providers.litellm._litellm")
+    async def test_timeout_raises_timeout_error(self, mock_litellm):
+        mock_litellm.acompletion = AsyncMock(side_effect=TimeoutError("timed out"))
+
+        p = LiteLLMProvider()
+        config = _make_config()
+        messages = [ChatMessage(role="user", content="hi")]
+
+        with pytest.raises(ProviderTimeoutError):
+            await p.send_message(messages, config)
+
+
+class TestResolver:
+    def test_litellm_in_registered_providers(self):
+        from ifixai.providers.resolver import REGISTERED_PROVIDERS
+
+        assert "litellm" in REGISTERED_PROVIDERS
+
+    def test_litellm_in_provider_map(self):
+        from ifixai.providers.resolver import _PROVIDER_MAP
+
+        assert "litellm" in _PROVIDER_MAP
+
+    def test_resolve_litellm_returns_provider(self):
+        from ifixai.providers.resolver import resolve_provider
+
+        provider = resolve_provider("litellm")
+        assert isinstance(provider, LiteLLMProvider)
+
+    def test_litellm_in_credential_env_vars(self):
+        from ifixai.providers.resolver import _PROVIDER_CREDENTIAL_ENV_VARS
+
+        assert "litellm" in _PROVIDER_CREDENTIAL_ENV_VARS