mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-06-23 21:59:57 +02:00
fix(make litellm optional import):
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user