diff --git a/agentops/llms/anthropic.py b/agentops/llms/anthropic.py index cbefe0544..322d21815 100644 --- 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 index 79f2cb042..68658761e 100644 --- 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 index 54e61473e..7d5f68005 100644 --- 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 index 39836c10a..053c42516 100644 --- 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 index 6b3a8f84b..bdcb21909 100644 --- 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 index fc16ab01f..0fd31a1d2 100644 --- 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