Add LLM Integration Tests

.github/workflows/python-testing.yml

@@ -24,6 +24,12 @@ jobs:
     runs-on: ubuntu-latest
     env:
       OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+      ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+      COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
+      GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }}
+      LITELLM_API_KEY: ${{ secrets.LITELLM_API_KEY }}
+      MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }}
+      AI21_API_KEY: ${{ secrets.AI21_API_KEY }}
 
     strategy:
       matrix:
@@ -42,3 +48,9 @@ jobs:
         env:
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           AGENTOPS_API_KEY: ${{ secrets.AGENTOPS_API_KEY }}
+          ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+          COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
+          GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }}
+          LITELLM_API_KEY: ${{ secrets.LITELLM_API_KEY }}
+          MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }}
+          AI21_API_KEY: ${{ secrets.AI21_API_KEY }}

agentops/__init__.py

@@ -21,6 +21,41 @@
 except ModuleNotFoundError:
     pass
 
+from .llms.providers import (
+    ai21,
+    anthropic,
+    cohere,
+    groq,
+    litellm,
+    mistral,
+    openai,
+)
+
+# Initialize providers when imported
+if "ai21" in sys.modules:
+    from ai21 import AI21Client
+    ai21.AI21Provider(client=AI21Client()).override()
+
+if "anthropic" in sys.modules:
+    from anthropic import Anthropic
+    anthropic.AnthropicProvider(client=Anthropic()).override()
+
+if "cohere" in sys.modules:
+    import cohere as cohere_sdk
+    cohere.CohereProvider(client=cohere_sdk).override()
+
+if "groq" in sys.modules:
+    from groq import Groq
+    groq.GroqProvider(client=Groq()).override()
+
+if "litellm" in sys.modules:
+    import litellm as litellm_sdk
+    litellm.LiteLLMProvider(client=litellm_sdk).override()
+
+if "mistralai" in sys.modules:
+    from mistralai import Mistral
+    mistral.MistralProvider(client=Mistral()).override()
+
 if "autogen" in sys.modules:
     Client().configure(instrument_llm_calls=False)
     Client()._initialize_autogen_logger()

agentops/llms/providers/ai21.py

@@ -145,98 +145,52 @@ async def async_generator():
     def override(self):
         self._override_completion()
         self._override_completion_async()
-        self._override_answer()
-        self._override_answer_async()
 
     def _override_completion(self):
-        from ai21.clients.studio.resources.chat import ChatCompletions
+        from ai21.clients.studio.ai21_client import AI21Client
 
-        global original_create
-        original_create = ChatCompletions.create
+        # Store original method
+        self.original_create = AI21Client.chat.completions.create
 
         def patched_function(*args, **kwargs):
             # Call the original function with its original arguments
             init_timestamp = get_ISO_time()
             session = kwargs.get("session", None)
             if "session" in kwargs.keys():
                 del kwargs["session"]
-            result = original_create(*args, **kwargs)
+            result = self.original_create(*args, **kwargs)
             return self.handle_response(result, kwargs, init_timestamp, session=session)
 
         # Override the original method with the patched one
-        ChatCompletions.create = patched_function
+        AI21Client.chat.completions.create = patched_function
 
     def _override_completion_async(self):
-        from ai21.clients.studio.resources.chat import AsyncChatCompletions
+        from ai21.clients.studio.async_ai21_client import AsyncAI21Client
 
-        global original_create_async
-        original_create_async = AsyncChatCompletions.create
+        # Store original method
+        self.original_create_async = AsyncAI21Client.chat.completions.create
 
         async def patched_function(*args, **kwargs):
             # Call the original function with its original arguments
             init_timestamp = get_ISO_time()
             session = kwargs.get("session", None)
             if "session" in kwargs.keys():
                 del kwargs["session"]
-            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
-        AsyncChatCompletions.create = patched_function
+        AsyncAI21Client.chat.completions.create = patched_function
 
-    def _override_answer(self):
-        from ai21.clients.studio.resources.studio_answer import StudioAnswer
-
-        global original_answer
-        original_answer = StudioAnswer.create
-
-        def patched_function(*args, **kwargs):
-            # Call the original function with its original arguments
-            init_timestamp = get_ISO_time()
-
-            session = kwargs.get("session", None)
-            if "session" in kwargs.keys():
-                del kwargs["session"]
-            result = original_answer(*args, **kwargs)
-            return self.handle_response(result, kwargs, init_timestamp, session=session)
-
-        StudioAnswer.create = patched_function
-
-    def _override_answer_async(self):
-        from ai21.clients.studio.resources.studio_answer import AsyncStudioAnswer
-
-        global original_answer_async
-        original_answer_async = AsyncStudioAnswer.create
-
-        async def patched_function(*args, **kwargs):
-            # Call the original function with its original arguments
-            init_timestamp = get_ISO_time()
-
-            session = kwargs.get("session", None)
-            if "session" in kwargs.keys():
-                del kwargs["session"]
-            result = await original_answer_async(*args, **kwargs)
-            return self.handle_response(result, kwargs, init_timestamp, session=session)
-
-        AsyncStudioAnswer.create = patched_function
+    # Answer functionality removed as it's not available in current version
 
     def undo_override(self):
         if (
             self.original_create is not None
             and self.original_create_async is not None
-            and self.original_answer is not None
-            and self.original_answer_async is not None
         ):
-            from ai21.clients.studio.resources.chat import (
-                ChatCompletions,
-                AsyncChatCompletions,
-            )
-            from ai21.clients.studio.resources.studio_answer import (
-                StudioAnswer,
-                AsyncStudioAnswer,
-            )
+            from ai21.clients.studio.ai21_client import AI21Client
+            from ai21.clients.studio.async_ai21_client import AsyncAI21Client
 
-            ChatCompletions.create = self.original_create
-            AsyncChatCompletions.create = self.original_create_async
-            StudioAnswer.create = self.original_answer
-            AsyncStudioAnswer.create = self.original_answer_async
+            AI21Client.chat.completions.create = self.original_create
+            AsyncAI21Client.chat.completions.create = self.original_create_async

agentops/llms/providers/groq.py

@@ -168,8 +168,14 @@ def _override_async_chat(self):
         async def patched_function(*args, **kwargs):
             # Call the original function with its original arguments
             init_timestamp = get_ISO_time()
+            session = kwargs.get("session", None)
+            if "session" in kwargs.keys():
+                del kwargs["session"]
             result = await self.original_async_create(*args, **kwargs)
-            return self.handle_response(result, kwargs, init_timestamp)
+            # Convert the result to a coroutine if it's not already awaitable
+            if not hasattr(result, '__await__'):
+                result = completions.ChatCompletion.model_validate(result)
+            return self.handle_response(result, kwargs, init_timestamp, session=session)
 
         # Override the original method with the patched one
         completions.AsyncCompletions.create = patched_function

agentops/llms/providers/litellm.py

@@ -1,3 +1,4 @@
+import inspect
 import pprint
 from typing import Optional
 
@@ -113,13 +114,36 @@ def generator():
 
         # litellm uses a CustomStreamWrapper
         if isinstance(response, CustomStreamWrapper):
-
-            def generator():
-                for chunk in response:
-                    handle_stream_chunk(chunk)
-                    yield chunk
-
-            return generator()
+            if inspect.isasyncgen(response):
+                async def async_generator():
+                    try:
+                        async for chunk in response:
+                            handle_stream_chunk(chunk)
+                            yield chunk
+                    except Exception as e:
+                        logger.warning(f"Error in async stream: {e}")
+                        raise
+                return async_generator()
+            elif hasattr(response, '__aiter__'):
+                async def async_generator():
+                    try:
+                        async for chunk in response:
+                            handle_stream_chunk(chunk)
+                            yield chunk
+                    except Exception as e:
+                        logger.warning(f"Error in async stream: {e}")
+                        raise
+                return async_generator()
+            else:
+                def generator():
+                    try:
+                        for chunk in response:
+                            handle_stream_chunk(chunk)
+                            yield chunk
+                    except Exception as e:
+                        logger.warning(f"Error in sync stream: {e}")
+                        raise
+                return generator()
 
         # For asynchronous AsyncStream
         elif isinstance(response, AsyncStream):

agentops/llms/providers/mistral.py

@@ -7,9 +7,11 @@
 from agentops.session import Session
 from agentops.log_config import logger
 from agentops.helpers import get_ISO_time, check_call_stack_for_agent_id
+from agentops.singleton import singleton
 from .instrumented_provider import InstrumentedProvider
 
 
+@singleton
 class MistralProvider(InstrumentedProvider):
     original_complete = None
     original_complete_async = None
@@ -22,8 +24,8 @@ def __init__(self, client):
 
     def handle_response(self, response, kwargs, init_timestamp, session: Optional[Session] = None) -> dict:
         """Handle responses for Mistral"""
-        from mistralai import Chat
-        from mistralai.types import UNSET, UNSET_SENTINEL
+        from mistralai import Mistral
+        from mistralai.models.chat import ChatCompletionResponse, ChatCompletionStreamResponse
 
         llm_event = LLMEvent(init_timestamp=init_timestamp, params=kwargs)
         if session is not None:
@@ -50,11 +52,11 @@ def handle_stream_chunk(chunk: dict):
                 if choice.delta.role:
                     accumulated_delta.role = choice.delta.role
 
-                # Check if tool_calls is Unset and set to None if it is
-                if choice.delta.tool_calls in (UNSET, UNSET_SENTINEL):
-                    accumulated_delta.tool_calls = None
-                elif choice.delta.tool_calls:
+                # Handle tool calls if they exist
+                if hasattr(choice.delta, 'tool_calls'):
                     accumulated_delta.tool_calls = choice.delta.tool_calls
+                else:
+                    accumulated_delta.tool_calls = None
 
                 if choice.finish_reason:
                     # Streaming is done. Record LLMEvent
@@ -123,76 +125,72 @@ async def async_generator():
         return response
 
     def _override_complete(self):
-        from mistralai import Chat
+        from mistralai import Mistral
 
-        global original_complete
-        original_complete = Chat.complete
+        self.original_complete = self.client.chat.complete
 
         def patched_function(*args, **kwargs):
             # Call the original function with its original arguments
             init_timestamp = get_ISO_time()
             session = kwargs.get("session", None)
             if "session" in kwargs.keys():
                 del kwargs["session"]
-            result = original_complete(*args, **kwargs)
+            result = self.original_complete(*args, **kwargs)
             return self.handle_response(result, kwargs, init_timestamp, session=session)
 
         # Override the original method with the patched one
-        Chat.complete = patched_function
+        self.client.chat.complete = patched_function
 
     def _override_complete_async(self):
-        from mistralai import Chat
+        from mistralai import Mistral
 
-        global original_complete_async
-        original_complete_async = Chat.complete_async
+        self.original_complete_async = self.client.chat.complete_async
 
         async def patched_function(*args, **kwargs):
             # Call the original function with its original arguments
             init_timestamp = get_ISO_time()
             session = kwargs.get("session", None)
             if "session" in kwargs.keys():
                 del kwargs["session"]
-            result = await original_complete_async(*args, **kwargs)
+            result = await self.original_complete_async(*args, **kwargs)
             return self.handle_response(result, kwargs, init_timestamp, session=session)
 
         # Override the original method with the patched one
-        Chat.complete_async = patched_function
+        self.client.chat.complete_async = patched_function
 
     def _override_stream(self):
-        from mistralai import Chat
+        from mistralai import Mistral
 
-        global original_stream
-        original_stream = Chat.stream
+        self.original_stream = self.client.chat.stream
 
         def patched_function(*args, **kwargs):
             # Call the original function with its original arguments
             init_timestamp = get_ISO_time()
             session = kwargs.get("session", None)
             if "session" in kwargs.keys():
                 del kwargs["session"]
-            result = original_stream(*args, **kwargs)
+            result = self.original_stream(*args, **kwargs)
             return self.handle_response(result, kwargs, init_timestamp, session=session)
 
         # Override the original method with the patched one
-        Chat.stream = patched_function
+        self.client.chat.stream = patched_function
 
     def _override_stream_async(self):
-        from mistralai import Chat
+        from mistralai import Mistral
 
-        global original_stream_async
-        original_stream_async = Chat.stream_async
+        self.original_stream_async = self.client.chat.stream_async
 
         async def patched_function(*args, **kwargs):
             # Call the original function with its original arguments
             init_timestamp = get_ISO_time()
             session = kwargs.get("session", None)
             if "session" in kwargs.keys():
                 del kwargs["session"]
-            result = await original_stream_async(*args, **kwargs)
+            result = await self.original_stream_async(*args, **kwargs)
             return self.handle_response(result, kwargs, init_timestamp, session=session)
 
         # Override the original method with the patched one
-        Chat.stream_async = patched_function
+        self.client.chat.stream_async = patched_function
 
     def override(self):
         self._override_complete()