dont undo if not done (#368)

AgentOps-AI · Aug 21, 2024 · 84c8557 · 84c8557
1 parent caaacc3
commit 84c8557
Show file tree

Hide file tree

Showing 6 changed files with 38 additions and 24 deletions.
diff --git a/agentops/llms/anthropic.py b/agentops/llms/anthropic.py
@@ -283,7 +283,8 @@ async def patched_function(*args, **kwargs):
         messages.AsyncMessages.create = patched_function
 
     def undo_override(self):
-        from anthropic.resources import messages
+        if self.original_create is not None and self.original_create_async is not None:
+            from anthropic.resources import messages
 
-        messages.Messages.create = self.original_create
-        messages.AsyncMessages.create = self.original_create_async
+            messages.Messages.create = self.original_create
+            messages.AsyncMessages.create = self.original_create_async
diff --git a/agentops/llms/cohere.py b/agentops/llms/cohere.py
@@ -22,11 +22,16 @@ def override(self):
         self._override_async_chat()
 
     def undo_override(self):
-        import cohere
-
-        cohere.Client.chat = self.original_create
-        cohere.Client.chat_stream = self.original_create_stream
-        cohere.AsyncClient.chat = self.original_create_async
+        if (
+            self.original_create is not None
+            and self.original_create_async is not None
+            and self.original_create_stream is not None
+        ):
+            import cohere
+
+            cohere.Client.chat = self.original_create
+            cohere.Client.chat_stream = self.original_create_stream
+            cohere.AsyncClient.chat = self.original_create_async
 
     def __init__(self, client):
         super().__init__(client)

diff --git a/agentops/llms/groq.py b/agentops/llms/groq.py
@@ -23,10 +23,11 @@ def override(self):
         self._override_async_chat()
 
     def undo_override(self):
-        from groq.resources.chat import completions
+        if self.original_create is not None and self.original_async_create is not None:
+            from groq.resources.chat import completions
 
-        completions.Completions.create = self.original_create
-        completions.AsyncCompletions.create = self.original_create
+            completions.Completions.create = self.original_create
+            completions.AsyncCompletions.create = self.original_create
 
     def handle_response(
         self, response, kwargs, init_timestamp, session: Optional[Session] = None

diff --git a/agentops/llms/litellm.py b/agentops/llms/litellm.py
@@ -25,14 +25,20 @@ def override(self):
         self._override_completion()
 
     def undo_override(self):
-        import litellm
-        from openai.resources.chat import completions
-
-        litellm.acompletion = self.original_create_async
-        litellm.completion = self.original_create
-
-        completions.Completions.create = self.original_oai_create
-        completions.AsyncCompletions.create = self.original_oai_create_async
+        if (
+            self.original_create is not None
+            and self.original_create_async is not None
+            and self.original_oai_create is not None
+            and self.original_oai_create_async is not None
+        ):
+            import litellm
+            from openai.resources.chat import completions
+
+            litellm.acompletion = self.original_create_async
+            litellm.completion = self.original_create
+
+            completions.Completions.create = self.original_oai_create
+            completions.AsyncCompletions.create = self.original_oai_create_async
 
     def handle_response(
         self, response, kwargs, init_timestamp, session: Optional[Session] = None

diff --git a/agentops/llms/ollama.py b/agentops/llms/ollama.py
@@ -65,7 +65,7 @@ def override(self):
         self._override_chat_async_client()
 
     def undo_override(self):
-        if "ollama" in sys.modules:
+        if original_func is not None:
             import ollama
 
             ollama.chat = original_func["ollama.chat"]

diff --git a/agentops/llms/openai.py b/agentops/llms/openai.py
@@ -244,14 +244,15 @@ async def patched_function(*args, **kwargs):
             #     kwargs["messages"] = prompt_override["messages"]
 
             # Call the original function with its original arguments
-            result = await original_create_async(*args, **kwargs)
+            result = await self.original_create_async(*args, **kwargs)
             return self.handle_response(result, kwargs, init_timestamp, session=session)
 
         # Override the original method with the patched one
         completions.AsyncCompletions.create = patched_function
 
     def undo_override(self):
-        from openai.resources.chat import completions
+        if self.original_create is not None and self.original_create_async is not None:
+            from openai.resources.chat import completions
 
-        completions.AsyncCompletions.create = self.original_create_async
-        completions.Completions.create = self.original_create
+            completions.AsyncCompletions.create = self.original_create_async
+            completions.Completions.create = self.original_create