dont import litellm

AgentOps-AI · Aug 14, 2024 · be8fcb8 · be8fcb8
1 parent 7e4d4c3
commit be8fcb8
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/agentops/llms/litellm.py b/agentops/llms/litellm.py
@@ -1,8 +1,6 @@
 import pprint
 from typing import Optional
 
-import litellm.utils
-
 from ..log_config import logger
 from ..event import LLMEvent, ErrorEvent
 from ..session import Session
@@ -99,7 +97,7 @@ def generator():
             return generator()
 
         # litellm uses a CustomStreamWrapper
-        if isinstance(response, litellm.utils.CustomStreamWrapper):
+        if response.__class__.__name__ == "CustomStreamWrapper":
 
             def generator():
                 for chunk in response: