fix(make litellm optional import):

This commit is contained in:
Alexander Myasoedov
2026-06-03 15:08:23 +03:00
parent a193ef9c2c
commit 816c8c6bc7
2 changed files with 10 additions and 1 deletions
@@ -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
@@ -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,