complete anthropic test

proxy/routes/anthropic.py

@@ -68,11 +68,9 @@ async def anthropic_proxy(
     )
     invariant_authorization = request.headers.get("invariant-authorization")
 
-    print("is stream:", request_body_json.get("stream"))    
     if request_body_json.get("stream"):
         return await handle_streaming_response(client, anthropic_request, dataset_name, invariant_authorization)
     else:
-        print("anthropic_request:", anthropic_request)
         try:    
             response = await client.send(anthropic_request)
         except httpx.HTTPStatusError as e:
@@ -80,7 +78,6 @@ async def anthropic_proxy(
                 status_code=response.status_code,
                 detail=f"Failed to fetch response: {response.text}, got error{e}",
             )
-        print("response:", response.status_code)
         await handle_non_streaming_response(
             response, dataset_name, request_body_json, invariant_authorization
         )

proxy/utils/explorer.py

@@ -8,7 +8,6 @@ from invariant_sdk.types.push_traces import PushTracesRequest, PushTracesRespons
 
 DEFAULT_API_URL = "https://explorer.invariantlabs.ai"
 
-
 async def push_trace(
     messages: List[List[Dict[str, Any]]],
     dataset_name: str,

tests/anthropic/test_anthropic_with_tool_call.py

@@ -15,7 +15,7 @@ pytest_plugins = ("pytest_asyncio",)
 
 class WeatherAgent:
     def __init__(self,proxy_url):
-        dataset_name = "claude_weather_agent_test" + str(
+        self.dataset_name = "claude_weather_agent_test" + str(
             datetime.datetime.now().strftime("%Y%m%d%H%M%S")
         )
         invariant_api_key = os.environ.get("INVARIANT_API_KEY","None")
@@ -23,7 +23,7 @@ class WeatherAgent:
             http_client=Client(
                 headers={"Invariant-Authorization": f"Bearer {invariant_api_key}"},
             ),
-            base_url=f"{proxy_url}/api/v1/proxy/{dataset_name}/anthropic",
+            base_url=f"{proxy_url}/api/v1/proxy/{self.dataset_name}/anthropic",
         )
         self.get_weather_function = {
             "name": "get_weather",
@@ -171,13 +171,12 @@ async def test_chat_completion_without_streaming(
     queries = [
         "What's the weather like in Zurich city?",
         "Tell me the weather for New York",
-        "How's the weather in London next week?",
     ]
-    cities = ["zurich", "new york", "london"]
+    cities = ["zurich", "new york"]
     # Process each query
+    responses = []
     for index, query in enumerate(queries):
         response = weather_agent.get_response(query)
-        print("non streaming response:",response)
         assert response is not None
         assert response[0].role == "assistant"
         assert response[0].stop_reason == "tool_use"
@@ -188,7 +187,33 @@ async def test_chat_completion_without_streaming(
         assert response[1].role == "assistant"
         assert response[1].stop_reason == "end_turn"
         assert cities[index] in response[1].content[0].text.lower()
-    
+        responses.append(response)
+
+    traces_response = await context.request.get(
+    f"{explorer_api_url}/api/v1/dataset/byuser/developer/{weather_agent.dataset_name}/traces"
+    )
+    traces = await traces_response.json()
+    assert len(traces) == len(queries)
+
+    for index,trace in enumerate(traces): 
+        trace_id = trace["id"]
+        # Fetch the trace
+        trace_response = await context.request.get(
+            f"{explorer_api_url}/api/v1/trace/{trace_id}"
+        )
+        trace = await trace_response.json()
+        trace_messages = trace["messages"]
+
+        assert trace_messages[0]["role"] == "user"
+        assert trace_messages[0]["content"] == queries[index]
+        assert trace_messages[1]["role"] == "assistant"
+        assert cities[index] in trace_messages[1]["content"].lower()
+        assert trace_messages[2]["role"] == "assistant"
+        assert trace_messages[2]["tool_calls"][0]["function"]["name"] == "get_weather"
+        assert cities[index] in trace_messages[2]["tool_calls"][0]["function"]["arguments"]["location"].lower()
+        assert trace_messages[3]["role"] == "tool"
+        assert trace_messages[4]["role"] == "assistant"
+        assert cities[index] in trace_messages[4]["content"].lower()
 
 
 
@@ -202,14 +227,12 @@ async def test_chat_completion_with_streaming(
     queries = [
         "What's the weather like in Zurich city?",
         "Tell me the weather for New York",
-        "How's the weather in London next week?",
     ]
-    cities = ["zurich", "new york", "london"]
+    cities = ["zurich", "new york"]
 
 
     for index, query in enumerate(queries):
         response = weather_agent.get_streaming_response(query)
-        print("streaming response:",response)
         assert response is not None
         assert response[0][0].type == "text"
         assert response[0][1].type == "tool_use"
@@ -218,3 +241,29 @@ async def test_chat_completion_with_streaming(
         
         assert response[1][0].type == "text"
         assert cities[index] in response[1][0].text.lower()
+    
+    traces_response = await context.request.get(
+    f"{explorer_api_url}/api/v1/dataset/byuser/developer/{weather_agent.dataset_name}/traces"
+    )
+    traces = await traces_response.json()
+    assert len(traces) == len(queries)
+
+    for index,trace in enumerate(traces): 
+        trace_id = trace["id"]
+        # Fetch the trace
+        trace_response = await context.request.get(
+            f"{explorer_api_url}/api/v1/trace/{trace_id}"
+        )
+        trace = await trace_response.json()
+        trace_messages = trace["messages"]
+        assert trace_messages[0]["role"] == "user"
+        assert trace_messages[0]["content"] == queries[index]
+        assert trace_messages[1]["role"] == "assistant"
+        assert cities[index] in trace_messages[1]["content"].lower()
+        assert trace_messages[2]["role"] == "assistant"
+        assert trace_messages[2]["tool_calls"][0]["function"]["name"] == "get_weather"
+        assert cities[index] in trace_messages[2]["tool_calls"][0]["function"]["arguments"]["location"].lower()
+        assert trace_messages[3]["role"] == "tool"
+        assert trace_messages[4]["role"] == "assistant"
+        assert cities[index] in trace_messages[4]["content"].lower()
+

tests/anthropic/test_anthropic_without_tool_call.py

@@ -1,19 +1,19 @@
 import anthropic
 import os
-import anthropic
 from httpx import Client
-import os
 # from invariant import testing
 import datetime
 import pytest
 import sys
-import httpx
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 from util import *  # needed for pytest fixtures
 
-@pytest.fixture
-def client(proxy_url):
+pytest_plugins = ("pytest_asyncio")
+@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="No ANTHROPIC_API_KEY set")
+async def test_chat_completion_without_streaming(
+    context, explorer_api_url,proxy_url
+):
     dataset_name = "claude_streaming_response_without_toolcall_test" + str(datetime.datetime.now().strftime("%Y%m%d%H%M%S"))
     invariant_api_key = os.environ.get("INVARIANT_API_KEY","None")
 
@@ -23,13 +23,7 @@ def client(proxy_url):
             ),
             base_url=f"{proxy_url}/api/v1/proxy/{dataset_name}/anthropic",
         )
-    return client
-
-
-pytest_plugins = ("pytest_asyncio")
-@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="No ANTHROPIC_API_KEY set")
-async def test_chat_completion_without_streaming(client):
-
+    
     cities = ["zurich", "new york", "london"]
     queries = [
         "Can you introduce Zurich city within 200 words?",
@@ -38,6 +32,7 @@ async def test_chat_completion_without_streaming(client):
     ]
 
     # Process each query
+    responses = []
     for query in queries:
         response = client.messages.create(
                 # system=self.system_prompt,
@@ -46,12 +41,50 @@ async def test_chat_completion_without_streaming(client):
                 messages=[{"role": "user", "content": query}],
             )
         response_text = response.content[0].text
+        responses.append(response_text)
         assert response_text is not None
         assert cities[queries.index(query)] in response_text.lower()
 
-@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="No ANTHROPIC_API_KEY set")
-def test_streaming_response_without_toolcall(proxy_url,client):    
+    traces_response = await context.request.get(
+        f"{explorer_api_url}/api/v1/dataset/byuser/developer/{dataset_name}/traces"
+    )
+    traces = await traces_response.json()
+    assert len(traces) == len(queries)
 
+    for index,trace in enumerate(traces): 
+        trace_id = trace["id"]
+        # Fetch the trace
+        trace_response = await context.request.get(
+            f"{explorer_api_url}/api/v1/trace/{trace_id}"
+        )
+        trace = await trace_response.json()
+        assert trace["messages"] == [
+            {
+                "role": "user",
+                "content": queries[index]
+            },
+            {
+                "role": "assistant",
+                "content": responses[index]
+            }
+        ]
+
+@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="No ANTHROPIC_API_KEY set")
+async def test_streaming_response_without_toolcall(
+    context,
+    explorer_api_url,
+    proxy_url):    
+
+    dataset_name = "claude_streaming_response_without_toolcall_test" + str(datetime.datetime.now().strftime("%Y%m%d%H%M%S"))
+    invariant_api_key = os.environ.get("INVARIANT_API_KEY","None")
+
+    client = anthropic.Anthropic(
+            http_client=Client(
+                headers={"Invariant-Authorization": f"Bearer {invariant_api_key}"},
+            ),
+            base_url=f"{proxy_url}/api/v1/proxy/{dataset_name}/anthropic",
+        )
+    
     cities = ["zurich", "new york", "london"]
     queries = [
         "Can you introduce Zurich city within 200 words?",
@@ -59,6 +92,7 @@ def test_streaming_response_without_toolcall(proxy_url,client):
         "How's the weather in London next week?"
     ]
     # Process each query
+    responses = []
     for index,query in enumerate(queries):
         messages = [
             {
@@ -67,7 +101,6 @@ def test_streaming_response_without_toolcall(proxy_url,client):
             }
         ]
         response_text = ""
-        attempt = 0
   
         with client.messages.stream(
             model="claude-3-5-sonnet-20241022",
@@ -77,6 +110,30 @@ def test_streaming_response_without_toolcall(proxy_url,client):
             for reply in response.text_stream:
                 response_text += reply
             assert cities[index] in response_text.lower()
-    
+        responses.append(response_text)
         assert response_text is not None
         assert cities[queries.index(query)] in response_text.lower()
+
+    traces_response = await context.request.get(
+        f"{explorer_api_url}/api/v1/dataset/byuser/developer/{dataset_name}/traces"
+    )
+    traces = await traces_response.json()
+    assert len(traces) == len(queries)
+
+    for index,trace in enumerate(traces): 
+        trace_id = trace["id"]
+        # Fetch the trace
+        trace_response = await context.request.get(
+            f"{explorer_api_url}/api/v1/trace/{trace_id}"
+        )
+        trace = await trace_response.json()
+        assert trace["messages"] == [
+            {
+                "role": "user",
+                "content": queries[index]
+            },
+            {
+                "role": "assistant",
+                "content": responses[index]
+            }
+        ]

tests/open_ai/test_chat_without_tool_calls.py

@@ -44,19 +44,10 @@ async def test_chat_completion(context, explorer_api_url, proxy_url, do_stream):
         expected_assistant_message = chat_response.choices[0].message.content
     else:
         full_response = ""
-        attempt = 0
-        while attempt<3:
-            try:
-                for chunk in chat_response:
-                    if chunk.choices and chunk.choices[0].delta.content:
-                        full_response += chunk.choices[0].delta.content
-                assert "PARIS" in full_response.upper()
-                break
-            except httpx.RemoteProtocolError as e:
-                attempt += 1
-                print(f"Streming error on attempt {attempt}: {e}")
-        else:
-            print("Max retries reached. Exiting.")
+        for chunk in chat_response:
+            if chunk.choices and chunk.choices[0].delta.content:
+                full_response += chunk.choices[0].delta.content
+        assert "PARIS" in full_response.upper()
         expected_assistant_message = full_response
 
     # Fetch the trace ids for the dataset