diff --git a/agentic_security/llm_providers/litellm_provider.py b/agentic_security/llm_providers/litellm_provider.py index 6dfab92..46f720a 100644 --- a/agentic_security/llm_providers/litellm_provider.py +++ b/agentic_security/llm_providers/litellm_provider.py @@ -2,7 +2,10 @@ from typing import Any -import litellm +try: + import litellm +except ImportError: + litellm = None from agentic_security.llm_providers.base import ( BaseLLMProvider, @@ -29,6 +32,10 @@ class LiteLLMProvider(BaseLLMProvider): api_base: str | None = None, **kwargs: Any, ) -> None: + if litellm is None: + raise LLMProviderError( + "litellm is not installed. Install it with: pip install litellm" + ) super().__init__(model, **kwargs) self._api_key = api_key self._api_base = api_base diff --git a/tests/unit/llm_providers/test_litellm_provider.py b/tests/unit/llm_providers/test_litellm_provider.py index de43879..4e947f1 100644 --- a/tests/unit/llm_providers/test_litellm_provider.py +++ b/tests/unit/llm_providers/test_litellm_provider.py @@ -4,6 +4,8 @@ import pytest from inline_snapshot import snapshot from unittest.mock import MagicMock, AsyncMock, patch +pytest.importorskip("litellm") + from agentic_security.llm_providers.litellm_provider import LiteLLMProvider from agentic_security.llm_providers.base import ( LLMMessage,