diff --git a/.github/workflows/python-testing.yml b/.github/workflows/python-testing.yml index 357624a55..b446ecd9e 100644 --- a/.github/workflows/python-testing.yml +++ b/.github/workflows/python-testing.yml @@ -24,6 +24,13 @@ jobs: runs-on: ubuntu-latest 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 }} strategy: matrix: @@ -42,3 +49,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 }} diff --git a/agentops/__init__.py b/agentops/__init__.py index 4150f839a..7981ecd0b 100755 --- a/agentops/__init__.py +++ b/agentops/__init__.py @@ -21,6 +21,47 @@ 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() diff --git a/agentops/llms/providers/ai21.py b/agentops/llms/providers/ai21.py index 8c907d525..de3ae8e09 100644 --- a/agentops/llms/providers/ai21.py +++ b/agentops/llms/providers/ai21.py @@ -145,14 +145,12 @@ 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 @@ -160,17 +158,17 @@ def patched_function(*args, **kwargs): 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 @@ -178,65 +176,18 @@ async def patched_function(*args, **kwargs): 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, - ) + if self.original_create is not None and self.original_create_async is not None: + 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 diff --git a/agentops/llms/providers/groq.py b/agentops/llms/providers/groq.py index 226a9123e..120cb05ce 100644 --- a/agentops/llms/providers/groq.py +++ b/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 diff --git a/agentops/llms/providers/litellm.py b/agentops/llms/providers/litellm.py index dff40765c..15c172291 100644 --- a/agentops/llms/providers/litellm.py +++ b/agentops/llms/providers/litellm.py @@ -1,3 +1,4 @@ +import inspect import pprint from typing import Optional @@ -113,13 +114,42 @@ 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): diff --git a/agentops/llms/providers/mistral.py b/agentops/llms/providers/mistral.py index 1754cae52..7a60dbf28 100644 --- a/agentops/llms/providers/mistral.py +++ b/agentops/llms/providers/mistral.py @@ -1,4 +1,5 @@ import inspect +import os import pprint import sys from typing import Optional @@ -7,23 +8,33 @@ 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 original_stream = None original_stream_async = None - def __init__(self, client): + def __init__(self, client=None): + from mistralai import Mistral + + if client is None: + if os.getenv("MISTRAL_API_KEY") is None: + raise ValueError("MISTRAL_API_KEY environment variable is required") + client = Mistral(api_key=os.getenv("MISTRAL_API_KEY")) super().__init__(client) self._provider_name = "Mistral" 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_completion import ChatCompletionResponse, ChatCompletionStreamResponse + + llm_event = LLMEvent(init_timestamp=init_timestamp, params=kwargs) if session is not None: @@ -33,52 +44,42 @@ def handle_stream_chunk(chunk: dict): # NOTE: prompt/completion usage not returned in response when streaming # We take the first ChatCompletionChunk and accumulate the deltas from all subsequent chunks to build one full chat completion if llm_event.returns is None: - llm_event.returns = chunk.data - - try: - accumulated_delta = llm_event.returns.choices[0].delta - llm_event.agent_id = check_call_stack_for_agent_id() - llm_event.model = "mistral/" + chunk.data.model - llm_event.prompt = kwargs["messages"] - - # NOTE: We assume for completion only choices[0] is relevant - choice = chunk.data.choices[0] - - if choice.delta.content: - accumulated_delta.content += choice.delta.content - - 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: - accumulated_delta.tool_calls = choice.delta.tool_calls - - if choice.finish_reason: - # Streaming is done. Record LLMEvent - llm_event.returns.choices[0].finish_reason = choice.finish_reason - llm_event.completion = { - "role": accumulated_delta.role, - "content": accumulated_delta.content, - "tool_calls": accumulated_delta.tool_calls, - } - llm_event.prompt_tokens = chunk.data.usage.prompt_tokens - llm_event.completion_tokens = chunk.data.usage.completion_tokens - llm_event.end_timestamp = get_ISO_time() - self._safe_record(session, llm_event) - - except Exception as e: - self._safe_record(session, ErrorEvent(trigger_event=llm_event, exception=e)) - - kwargs_str = pprint.pformat(kwargs) - chunk = pprint.pformat(chunk) - logger.warning( - f"Unable to parse a chunk for LLM call. Skipping upload to AgentOps\n" - f"chunk:\n {chunk}\n" - f"kwargs:\n {kwargs_str}\n" - ) + llm_event.returns = chunk + + accumulated_delta = llm_event.returns.choices[0].delta + llm_event.agent_id = check_call_stack_for_agent_id() + llm_event.model = "mistral/" + chunk.model + llm_event.prompt = kwargs["messages"] + + # NOTE: We assume for completion only choices[0] is relevant + choice = chunk.choices[0] + + if hasattr(choice.delta, "content") and choice.delta.content: + if not hasattr(accumulated_delta, "content"): + accumulated_delta.content = "" + accumulated_delta.content += choice.delta.content + + if hasattr(choice.delta, "role") and choice.delta.role: + accumulated_delta.role = choice.delta.role + + # 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 + llm_event.returns.choices[0].finish_reason = choice.finish_reason + llm_event.completion = { + "role": accumulated_delta.role, + "content": accumulated_delta.content, + "tool_calls": accumulated_delta.tool_calls, + } + llm_event.prompt_tokens = chunk.usage.prompt_tokens + llm_event.completion_tokens = chunk.usage.completion_tokens + llm_event.end_timestamp = get_ISO_time() + self._safe_record(session, llm_event) # if the response is a generator, decorate the generator if inspect.isgenerator(response): @@ -99,34 +100,23 @@ async def async_generator(): return async_generator() - try: - llm_event.returns = response - llm_event.agent_id = check_call_stack_for_agent_id() - llm_event.model = "mistral/" + response.model - llm_event.prompt = kwargs["messages"] - llm_event.prompt_tokens = response.usage.prompt_tokens - llm_event.completion = response.choices[0].message.model_dump() - llm_event.completion_tokens = response.usage.completion_tokens - llm_event.end_timestamp = get_ISO_time() - - self._safe_record(session, llm_event) - except Exception as e: - self._safe_record(session, ErrorEvent(trigger_event=llm_event, exception=e)) - kwargs_str = pprint.pformat(kwargs) - response = pprint.pformat(response) - logger.warning( - f"Unable to parse response for LLM call. Skipping upload to AgentOps\n" - f"response:\n {response}\n" - f"kwargs:\n {kwargs_str}\n" - ) + llm_event.returns = response + llm_event.agent_id = check_call_stack_for_agent_id() + llm_event.model = "mistral/" + response.model + llm_event.prompt = kwargs["messages"] + llm_event.prompt_tokens = response.usage.prompt_tokens + llm_event.completion = response.choices[0].message.model_dump() + llm_event.completion_tokens = response.usage.completion_tokens + llm_event.end_timestamp = get_ISO_time() + + self._safe_record(session, llm_event) 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 @@ -134,17 +124,16 @@ def patched_function(*args, **kwargs): 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 @@ -152,17 +141,16 @@ async def patched_function(*args, **kwargs): 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 @@ -170,17 +158,16 @@ def patched_function(*args, **kwargs): 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 @@ -188,11 +175,11 @@ async def patched_function(*args, **kwargs): 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() @@ -209,7 +196,7 @@ def undo_override(self): ): from mistralai import Chat - Chat.complete = self.original_complete - Chat.complete_async = self.original_complete_async - Chat.stream = self.original_stream - Chat.stream_async = self.original_stream_async + self.client.chat.complete = self.original_complete + self.client.chat.complete_async = self.original_complete_async + self.client.chat.stream = self.original_stream + self.client.chat.stream_async = self.original_stream_async diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 000000000..f23b4ab72 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +markers = + integration: marks tests as integration tests diff --git a/tests/ai21_handlers/test_ai21_integration.py b/tests/ai21_handlers/test_ai21_integration.py new file mode 100644 index 000000000..5a6621019 --- /dev/null +++ b/tests/ai21_handlers/test_ai21_integration.py @@ -0,0 +1,80 @@ +import os +import pytest +import agentops +import asyncio +import ai21 # Import module to trigger provider initialization +from ai21 import AI21Client, AsyncAI21Client +from ai21.models.chat import ChatMessage + + +@pytest.mark.integration +def test_ai21_integration(): + """Integration test demonstrating all four AI21 call patterns: + 1. Sync (non-streaming) + 2. Sync (streaming) + 3. Async (non-streaming) + 4. Async (streaming) + + Verifies that AgentOps correctly tracks all LLM calls via analytics. + """ + print("AGENTOPS_API_KEY present:", bool(os.getenv("AGENTOPS_API_KEY"))) + print("AI21_API_KEY present:", bool(os.getenv("AI21_API_KEY"))) + + # Initialize AgentOps without auto-starting session + agentops.init(auto_start_session=False) + session = agentops.start_session() + + def sync_no_stream(): + client = AI21Client(api_key=os.getenv("AI21_API_KEY")) + messages = [ChatMessage(content="Hello from sync no stream", role="user")] + client.chat.completions.create( + messages=messages, + model="jamba-1.5-large", + max_tokens=20, + ) + + def sync_stream(): + client = AI21Client(api_key=os.getenv("AI21_API_KEY")) + messages = [ChatMessage(content="Hello from sync streaming", role="user")] + response = client.chat.completions.create( + messages=messages, + model="jamba-1.5-large", + max_tokens=20, + stream=True, + ) + for chunk in response: + if hasattr(chunk, "choices") and chunk.choices[0].delta.content: + pass + + async def async_no_stream(): + client = AsyncAI21Client(api_key=os.getenv("AI21_API_KEY")) + messages = [ChatMessage(content="Hello from async no stream", role="user")] + await client.chat.completions.create( + messages=messages, + model="jamba-1.5-large", + max_tokens=20, + ) + + async def async_stream(): + client = AsyncAI21Client(api_key=os.getenv("AI21_API_KEY")) + messages = [ChatMessage(content="Hello from async streaming", role="user")] + response = await client.chat.completions.create( + messages=messages, + model="jamba-1.5-large", + max_tokens=20, + stream=True, + ) + async for chunk in response: + if hasattr(chunk, "choices") and chunk.choices[0].delta.content: + pass + + # Call each function + sync_no_stream() + sync_stream() + asyncio.run(async_no_stream()) + asyncio.run(async_stream()) + session.end_session("Success") + analytics = session.get_analytics() + print("Final analytics:", analytics) + # Verify that all LLM calls were tracked + assert analytics["LLM calls"] >= 4, f"Expected at least 4 LLM calls, but got {analytics['LLM calls']}" diff --git a/tests/anthropic_handlers/test_anthropic_integration.py b/tests/anthropic_handlers/test_anthropic_integration.py new file mode 100644 index 000000000..8e7d62917 --- /dev/null +++ b/tests/anthropic_handlers/test_anthropic_integration.py @@ -0,0 +1,72 @@ +import os +import pytest +import agentops +import asyncio +import anthropic # Import module to trigger provider initialization +from anthropic import Anthropic, AsyncAnthropic + + +@pytest.mark.integration +def test_anthropic_integration(): + """Integration test demonstrating all four Anthropic call patterns: + 1. Sync (non-streaming) + 2. Sync (streaming) + 3. Async (non-streaming) + 4. Async (streaming) + + Verifies that AgentOps correctly tracks all LLM calls via analytics. + """ + print("AGENTOPS_API_KEY present:", bool(os.getenv("AGENTOPS_API_KEY"))) + print("ANTHROPIC_API_KEY present:", bool(os.getenv("ANTHROPIC_API_KEY"))) + # Initialize AgentOps without auto-starting session + agentops.init(auto_start_session=False) + session = agentops.start_session() + + def sync_no_stream(): + client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY")) + client.messages.create( + model="claude-3-5-sonnet-20240620", + messages=[{"role": "user", "content": "Hello from sync no stream"}], + max_tokens=20, + ) + + def sync_stream(): + client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY")) + stream_result = client.messages.create( + model="claude-3-5-sonnet-20240620", + messages=[{"role": "user", "content": "Hello from sync streaming"}], + max_tokens=20, + stream=True, + ) + for _ in stream_result: + pass + + async def async_no_stream(): + client = AsyncAnthropic(api_key=os.getenv("ANTHROPIC_API_KEY")) + await client.messages.create( + model="claude-3-5-sonnet-20240620", + messages=[{"role": "user", "content": "Hello from async no stream"}], + max_tokens=20, + ) + + async def async_stream(): + client = AsyncAnthropic(api_key=os.getenv("ANTHROPIC_API_KEY")) + async_stream_result = await client.messages.create( + model="claude-3-5-sonnet-20240620", + messages=[{"role": "user", "content": "Hello from async streaming"}], + max_tokens=20, + stream=True, + ) + async for _ in async_stream_result: + pass + + # Call each function + sync_no_stream() + sync_stream() + asyncio.run(async_no_stream()) + asyncio.run(async_stream()) + session.end_session("Success") + analytics = session.get_analytics() + print(analytics) + # Verify that all LLM calls were tracked + assert analytics["LLM calls"] >= 4, f"Expected at least 4 LLM calls, but got {analytics['LLM calls']}" diff --git a/tests/cohere_handlers/test_cohere_integration.py b/tests/cohere_handlers/test_cohere_integration.py new file mode 100644 index 000000000..f916b2199 --- /dev/null +++ b/tests/cohere_handlers/test_cohere_integration.py @@ -0,0 +1,85 @@ +import os +import pytest +import agentops +import asyncio +import cohere +from cohere.types.chat_text_generation_event import ChatTextGenerationEvent +from cohere.types.chat_stream_start_event import ChatStreamStartEvent +from cohere.types.chat_stream_end_event import ChatStreamEndEvent + + +@pytest.mark.integration +def test_cohere_integration(): + """Integration test demonstrating all four Cohere call patterns: + 1. Sync (non-streaming) + 2. Sync (streaming) + 3. Async (non-streaming) + 4. Async (streaming) + + Verifies that AgentOps correctly tracks all LLM calls via analytics. + """ + print("AGENTOPS_API_KEY present:", bool(os.getenv("AGENTOPS_API_KEY"))) + print("COHERE_API_KEY present:", bool(os.getenv("COHERE_API_KEY"))) + + # Initialize AgentOps without auto-starting session + agentops.init(auto_start_session=False) + session = agentops.start_session() + + def sync_no_stream(): + client = cohere.Client(api_key=os.getenv("COHERE_API_KEY")) + client.chat( + message="Hello from sync no stream", + model="command", + max_tokens=100, + ) + + def sync_stream(): + client = cohere.Client(api_key=os.getenv("COHERE_API_KEY")) + stream_result = client.chat( + message="Hello from sync streaming", + model="command", + max_tokens=100, + stream=True, + ) + for chunk in stream_result: + if isinstance(chunk, ChatTextGenerationEvent): + continue + elif isinstance(chunk, ChatStreamStartEvent): + continue + elif isinstance(chunk, ChatStreamEndEvent): + break + + async def async_no_stream(): + client = cohere.AsyncClient(api_key=os.getenv("COHERE_API_KEY")) + await client.chat( + message="Hello from async no stream", + model="command", + max_tokens=100, + ) + + async def async_stream(): + client = cohere.AsyncClient(api_key=os.getenv("COHERE_API_KEY")) + async_stream_result = await client.chat( + message="Hello from async streaming", + model="command", + max_tokens=100, + stream=True, + ) + async for chunk in async_stream_result: + if isinstance(chunk, ChatTextGenerationEvent): + continue + elif isinstance(chunk, ChatStreamStartEvent): + continue + elif isinstance(chunk, ChatStreamEndEvent): + break + + # Call each function + sync_no_stream() + sync_stream() + asyncio.run(async_no_stream()) + asyncio.run(async_stream()) + session.end_session("Success") + analytics = session.get_analytics() + print(analytics) + # Verify that all LLM calls were tracked + assert analytics["LLM calls"] >= 4, f"Expected at least 4 LLM calls, but got {analytics['LLM calls']}" diff --git a/tests/groq_handlers/test_groq_integration.py b/tests/groq_handlers/test_groq_integration.py new file mode 100644 index 000000000..3195dfb72 --- /dev/null +++ b/tests/groq_handlers/test_groq_integration.py @@ -0,0 +1,71 @@ +import os +import pytest +import agentops +import asyncio +import groq # Import module to trigger provider initialization +from groq import Groq +from groq.resources.chat import AsyncCompletions + + +@pytest.mark.integration +def test_groq_integration(): + """Integration test demonstrating all four Groq call patterns: + 1. Sync (non-streaming) + 2. Sync (streaming) + 3. Async (non-streaming) + 4. Async (streaming) + + Verifies that AgentOps correctly tracks all LLM calls via analytics. + """ + print("AGENTOPS_API_KEY present:", bool(os.getenv("AGENTOPS_API_KEY"))) + print("GROQ_API_KEY present:", bool(os.getenv("GROQ_API_KEY"))) + + # Initialize AgentOps without auto-starting session + agentops.init(auto_start_session=False) + session = agentops.start_session() + + def sync_no_stream(): + client = Groq(api_key=os.getenv("GROQ_API_KEY")) + client.chat.completions.create( + messages=[{"role": "user", "content": "Hello from sync no stream"}], + model="mixtral-8x7b-32768", + ) + + def sync_stream(): + client = Groq(api_key=os.getenv("GROQ_API_KEY")) + stream_result = client.chat.completions.create( + messages=[{"role": "user", "content": "Hello from sync streaming"}], + model="mixtral-8x7b-32768", + stream=True, + ) + for _ in stream_result: + pass + + async def async_no_stream(): + client = Groq(api_key=os.getenv("GROQ_API_KEY")) + result = client.chat.completions.create( + messages=[{"role": "user", "content": "Hello from async no stream"}], + model="mixtral-8x7b-32768", + ) + return result + + async def async_stream(): + client = Groq(api_key=os.getenv("GROQ_API_KEY")) + async_stream_result = client.chat.completions.create( + messages=[{"role": "user", "content": "Hello from async streaming"}], + model="mixtral-8x7b-32768", + stream=True, + ) + for _ in async_stream_result: + pass + + # Call each function + sync_no_stream() + sync_stream() + asyncio.run(async_no_stream()) + asyncio.run(async_stream()) + session.end_session("Success") + analytics = session.get_analytics() + print(analytics) + # Verify that all LLM calls were tracked + assert analytics["LLM calls"] >= 4, f"Expected at least 4 LLM calls, but got {analytics['LLM calls']}" diff --git a/tests/litellm_handlers/test_litellm_integration.py b/tests/litellm_handlers/test_litellm_integration.py new file mode 100644 index 000000000..f41b8c84c --- /dev/null +++ b/tests/litellm_handlers/test_litellm_integration.py @@ -0,0 +1,74 @@ +import os +import pytest +import agentops +import asyncio +import litellm + + +@pytest.mark.integration +def test_litellm_integration(): + """Integration test demonstrating all four LiteLLM call patterns: + 1. Sync (non-streaming) + 2. Sync (streaming) + 3. Async (non-streaming) + 4. Async (streaming) + + Verifies that AgentOps correctly tracks all LLM calls via analytics. + Uses Anthropic's Claude model as the backend provider. + """ + print("AGENTOPS_API_KEY present:", bool(os.getenv("AGENTOPS_API_KEY"))) + print("ANTHROPIC_API_KEY present:", bool(os.getenv("ANTHROPIC_API_KEY"))) # LiteLLM uses Anthropic + + # Initialize AgentOps without auto-starting session + agentops.init(auto_start_session=False) + session = agentops.start_session() + + # Set API key once at the start + litellm.api_key = os.getenv("ANTHROPIC_API_KEY") + + async def run_all_tests(): + # Sync non-streaming (using acompletion for consistency) + await litellm.acompletion( + model="anthropic/claude-2", + messages=[{"role": "user", "content": "Hello from sync no stream"}], + max_tokens=100, + ) + + # Sync streaming + response = await litellm.acompletion( + model="anthropic/claude-2", + messages=[{"role": "user", "content": "Hello from sync streaming"}], + stream=True, + max_tokens=100, + ) + async for chunk in response: + if hasattr(chunk, "choices") and chunk.choices[0].delta.content: + pass + + # Async non-streaming + await litellm.acompletion( + model="anthropic/claude-2", + messages=[{"role": "user", "content": "Hello from async no stream"}], + max_tokens=100, + ) + + # Async streaming + async_stream_result = await litellm.acompletion( + model="anthropic/claude-2", + messages=[{"role": "user", "content": "Hello from async streaming"}], + stream=True, + max_tokens=100, + ) + async for chunk in async_stream_result: + if hasattr(chunk, "choices") and chunk.choices[0].delta.content: + pass + + # Run all tests in a single event loop + asyncio.run(run_all_tests()) + + # End session and verify analytics + session.end_session("Success") + analytics = session.get_analytics() + print(analytics) + # Verify that all LLM calls were tracked + assert analytics["LLM calls"] >= 4, f"Expected at least 4 LLM calls, but got {analytics['LLM calls']}" diff --git a/tests/mistral_handlers/test_mistral_integration.py b/tests/mistral_handlers/test_mistral_integration.py new file mode 100644 index 000000000..9497d1fd7 --- /dev/null +++ b/tests/mistral_handlers/test_mistral_integration.py @@ -0,0 +1,69 @@ +import os +import pytest +import agentops +import asyncio +import mistralai # Import module to trigger provider initialization +from mistralai import Mistral + + +@pytest.mark.integration +def test_mistral_integration(): + """Integration test demonstrating all four Mistral call patterns: + 1. Sync (non-streaming) + 2. Sync (streaming) + 3. Async (non-streaming) + 4. Async (streaming) + + Verifies that AgentOps correctly tracks all LLM calls via analytics. + """ + print("AGENTOPS_API_KEY present:", bool(os.getenv("AGENTOPS_API_KEY"))) + print("MISTRAL_API_KEY present:", bool(os.getenv("MISTRAL_API_KEY"))) + + # Initialize AgentOps without auto-starting session + agentops.init(auto_start_session=False) + session = agentops.start_session() + + def sync_no_stream(): + client = Mistral(api_key=os.getenv("MISTRAL_API_KEY")) + client.chat.complete( + model="mistral-large-latest", + messages=[{"role": "user", "content": "Hello from sync no stream"}], + ) + + def sync_stream(): + client = Mistral(api_key=os.getenv("MISTRAL_API_KEY")) + stream_result = client.chat.stream( + model="mistral-large-latest", + messages=[{"role": "user", "content": "Hello from sync streaming"}], + ) + for chunk in stream_result: + if chunk.choices[0].delta.content: + pass + + async def async_no_stream(): + client = Mistral(api_key=os.getenv("MISTRAL_API_KEY")) + await client.chat.complete_async( + model="mistral-large-latest", + messages=[{"role": "user", "content": "Hello from async no stream"}], + ) + + async def async_stream(): + client = Mistral(api_key=os.getenv("MISTRAL_API_KEY")) + async_stream_result = await client.chat.stream_async( + model="mistral-large-latest", + messages=[{"role": "user", "content": "Hello from async streaming"}], + ) + async for chunk in async_stream_result: + if chunk.choices[0].delta.content: + pass + + # Call each function + sync_no_stream() + sync_stream() + asyncio.run(async_no_stream()) + asyncio.run(async_stream()) + session.end_session("Success") + analytics = session.get_analytics() + print(analytics) + # Verify that all LLM calls were tracked + assert analytics["LLM calls"] >= 4, f"Expected at least 4 LLM calls, but got {analytics['LLM calls']}" diff --git a/tox.ini b/tox.ini index b4167ae3e..96d0d8632 100644 --- a/tox.ini +++ b/tox.ini @@ -32,6 +32,12 @@ deps = langchain termcolor python-dotenv + anthropic + cohere + groq + mistralai + ai21 + litellm -e . commands = coverage run --source . -m pytest @@ -41,6 +47,22 @@ commands = passenv = OPENAI_API_KEY AGENTOPS_API_KEY + ANTHROPIC_API_KEY + COHERE_API_KEY + GROQ_API_KEY + LITELLM_API_KEY + MISTRAL_API_KEY + AI21_API_KEY +setenv = + PYTHONPATH = {toxinidir} + OPENAI_API_KEY = {env:OPENAI_API_KEY} + AGENTOPS_API_KEY = {env:AGENTOPS_API_KEY} + ANTHROPIC_API_KEY = {env:ANTHROPIC_API_KEY} + COHERE_API_KEY = {env:COHERE_API_KEY} + GROQ_API_KEY = {env:GROQ_API_KEY} + LITELLM_API_KEY = {env:LITELLM_API_KEY} + MISTRAL_API_KEY = {env:MISTRAL_API_KEY} + AI21_API_KEY = {env:AI21_API_KEY} [coverage:run] branch = True