fix(make litellm optional import):

This commit is contained in:
Alexander Myasoedov committed 2026-06-03 15:08:23 +03:00
1 parent a193ef9c2c
commit 816c8c6bc7
2 files changed
+10 -1

No files matched your search

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