organize anthropic test

proxy/routes/anthropic.py

@@ -71,7 +71,7 @@ async def anthropic_proxy(
     request_body_json = json.loads(request_body)
 
     anthropic_url = f"https://api.anthropic.com/{endpoint}"
-    client = httpx.AsyncClient()
+    client = httpx.AsyncClient(timeout=httpx.Timeout(60))
     
     anthropic_request = client.build_request(
         "POST", 
@@ -81,17 +81,19 @@ async def anthropic_proxy(
     )
     invariant_authorization = request.headers.get("invariant-authorization")
 
-    response = await client.send(anthropic_request)
-
+    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:
-        response = await client.send(anthropic_request)
-        if response.status_code != 200:
+        print("anthropic_request:", anthropic_request)
+        try:    
+            response = await client.send(anthropic_request)
+        except httpx.HTTPStatusError as e:
             raise HTTPException(
                 status_code=response.status_code,
-                detail=f"Failed to fetch response: {response.text}",
+                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
         )
@@ -108,6 +110,7 @@ async def push_to_explorer(
     # Combine the messages from the request body and Anthropic response
     messages = request_body.get("messages", [])
     messages += [merged_response]
+
     if reformat:
         messages = anthropic_to_invariant_messages(messages)
     _ = await push_trace(
@@ -168,7 +171,6 @@ async def handle_streaming_response(
                     formatted_invariant_response[-1],
                     json.loads(anthropic_request.content),
                     invariant_authorization,
-                    reformat = False
                 )
     
     generator = event_generator()
@@ -182,12 +184,14 @@ def process_chunk_text(chunk, formatted_invariant_response):
     Example of chunk list can be find in:
     ../../resources/streaming_chunk_text/anthropic.txt
     """
-
     text_decode = chunk.decode().strip()
     for text_block in text_decode.split("\n\n"):
-        text_data = text_block.split("\ndata:")[1]
-        text_json = json.loads(text_data)
-        update_formatted_invariant_response(text_json, formatted_invariant_response)
+        # might be empty block
+      
+        if len(text_block.split("\ndata:"))>1:
+            text_data = text_block.split("\ndata:")[1]
+            text_json = json.loads(text_data)
+            update_formatted_invariant_response(text_json, formatted_invariant_response)
 
 def update_formatted_invariant_response(text_json, formatted_invariant_response):
     if text_json.get("type") == MESSAGE_START:
@@ -200,8 +204,20 @@ def update_formatted_invariant_response(text_json, formatted_invariant_response)
             "stop_reason": message.get("stop_reason"),
             "stop_sequence": message.get("stop_sequence"),
         })
+    elif text_json.get("type") == CONTENT_BLOCK_START and text_json.get("content_block").get("type")=="tool_use":
+        content_block = text_json.get("content_block")
+        formatted_invariant_response.append(
+            {
+                "role": "tool",
+                "tool_id": content_block.get("id"),
+                "content": "",
+            }
+        )
     elif text_json.get("type") == CONTENT_BLOCK_DELTA:
-        formatted_invariant_response[-1]["content"] += text_json.get("delta").get("text")
+        if formatted_invariant_response[-1]["role"]=="assistant":
+            formatted_invariant_response[-1]["content"] += text_json.get("delta").get("text")
+        elif formatted_invariant_response[-1]["role"]=="tool":
+            formatted_invariant_response[-1]["content"] += text_json.get("delta").get("partial_json")
     elif text_json.get("type") == MESSGAE_DELTA:
         formatted_invariant_response[-1]["stop_reason"] = text_json.get("delta").get("stop_reason")
 
@@ -253,24 +269,27 @@ def handle_user_message(message, keep_empty_tool_response):
 
 def handle_assistant_message(message):
     output = []
-    for sub_message in message["content"]:
-        if sub_message["type"] == "text":
-            output.append({"role": "assistant", "content": sub_message.get("text")})
-        elif sub_message["type"] == "tool_use":
-            output.append(
-                {
-                    "role": "assistant",
-                    "content": None,
-                    "tool_calls": [
-                        {
-                            "tool_id": sub_message.get("id"),
-                            "type": "function",
-                            "function": {
-                                "name": sub_message.get("name"),
-                                "arguments": sub_message.get("input"),
-                            },
-                        }
-                    ],
-                }
-            )
+    if isinstance(message["content"], list):
+        for sub_message in message["content"]:
+            if sub_message["type"] == "text":
+                output.append({"role": "assistant", "content": sub_message.get("text")})
+            elif sub_message["type"] == "tool_use":
+                output.append(
+                    {
+                        "role": "assistant",
+                        "content": None,
+                        "tool_calls": [
+                            {
+                                "tool_id": sub_message.get("id"),
+                                "type": "function",
+                                "function": {
+                                    "name": sub_message.get("name"),
+                                    "arguments": sub_message.get("input"),
+                                },
+                            }
+                        ],
+                    }
+                )
+    else:
+        output.append({"role": "assistant", "content": message["content"]})
     return output

resources/streaming_chunk_text/anthropic.txt

@@ -3,4 +3,4 @@ anthropic_chunk_list = [b'event: message_start\ndata: {"type":"message_start","m
     , b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" and leisure."}   }\n\n'
     , b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n'
     , b'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":301}      }\n\n'
-    , b'event: message_stop\ndata: {"type":"message_stop"       }\n\n']
+    , b'event: message_stop\ndata: {"type":"message_stop"       }\n\n']

tests/anthropic/test_chat_with_tool_call.py

@@ -0,0 +1,219 @@
+import datetime
+import os
+from typing import Dict
+import json
+import anthropic
+import pytest
+from httpx import Client
+import sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from util import *  # needed for pytest fixtures
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+class WeatherAgent:
+    def __init__(self,proxy_url):
+        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")
+        self.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",
+        )
+        self.get_weather_function = {
+            "name": "get_weather",
+            "description": "Get the current weather in a given location",
+            "input_schema": {
+                "type": "object",
+                "properties": {
+                    "location": {
+                        "type": "string",
+                        "description": "The city and state, e.g. San Francisco, CA",
+                    },
+                    "unit": {
+                        "type": "string",
+                        "enum": ["celsius", "fahrenheit"],
+                        "description": 'The unit of temperature, either "celsius" or "fahrenheit"',
+                    },
+                },
+                "required": ["location"],
+            },
+        }
+
+        # self.system_prompt = """You are an assistant that can perform weather searches using function calls.
+        #     When a user asks for weather information, respond with a JSON object specifying
+        #     the function name `get_weather` and the arguments latitude and longitude are needed."""
+
+    def get_response(self, user_query: str) -> Dict:
+        """
+        Get the response from the agent for a given user query for weather.
+        """
+        messages = [{"role": "user", "content": user_query}]
+        response_list = []
+        while True:
+            response = self.client.messages.create(
+                # system=self.system_prompt,
+                tools=[self.get_weather_function],
+                model="claude-3-5-sonnet-20241022",
+                max_tokens=1024,
+                messages=messages
+            )
+            response_list.append(response)
+            # If there's tool call, Extract the tool call parameters from the response
+            if len(response.content) > 1 and response.content[1].type == "tool_use":
+                tool_call_params = response.content[1].input
+                tool_call_result = self.get_weather(tool_call_params["location"])
+                tool_call_id = response.content[1].id
+                messages.append({"role": response.role, "content": response.content})
+                messages.append(
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "tool_result",
+                                "tool_use_id": tool_call_id,
+                                "content": tool_call_result,
+                            }
+                        ],
+                    }
+                )
+            else:
+                return response_list
+            
+    def get_streaming_response(self, user_query: str) -> Dict:
+        messages = [{"role": "user", "content": user_query}]
+        response_list = []
+
+        def clean_quotes(text):
+            # Convert \' to '
+            text = text.replace("\'", "'")
+            # Convert \" to "
+            text = text.replace('\"', '"')
+            text = text.replace("\n", " ")
+            return text
+        
+        while True:
+            json_data = ""
+            content = []
+            response = []
+            with self.client.messages.stream(
+                tools=[self.get_weather_function],
+                model="claude-3-5-sonnet-20241022",
+                max_tokens=1024,
+                messages=messages,
+            ) as stream:
+                for event in stream:
+                    if isinstance(event, anthropic.types.RawContentBlockStartEvent):
+                        # Start a new block
+                        current_block = event.content_block
+                        current_text = ""
+                    elif isinstance(event, anthropic.types.RawContentBlockDeltaEvent):
+                        if hasattr(event.delta, 'text'):
+                            # Accumulate text for TextBlock
+                            current_text += clean_quotes(event.delta.text)
+                        elif hasattr(event.delta, 'partial_json'):
+                            # Accumulate JSON for ToolUseBlock
+                            json_data += clean_quotes(event.delta.partial_json)
+                            current_text += clean_quotes(event.delta.partial_json)
+                    elif isinstance(event, anthropic.types.RawContentBlockStopEvent):
+                        # Block is complete, add it to content
+                        if current_block.type == 'text':
+                            content.append(anthropic.types.TextBlock(citations=None, text=current_text, type="text"))
+                        elif current_block.type == 'tool_use':
+                            content.append(
+                                anthropic.types.ToolUseBlock(id=current_block.id,
+                                input=json.loads(current_text), 
+                                name=current_block.name, 
+                                type="tool_use")
+                            )
+            response_list.append(content)    
+            if isinstance(event, anthropic.types.RawMessageStopEvent) and event.message.stop_reason == "tool_use":
+                tool_call_params = json.loads(json_data)
+                tool_call_result = self.get_weather(tool_call_params["location"])
+                messages.append({"role": "assistant", "content": content})
+                messages.append(
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "tool_result",
+                                "tool_use_id": content[-1].id,
+                                "content": tool_call_result,
+                            }
+                        ],
+                    }
+                )
+            else:
+                return response_list
+
+    def get_weather(self, location: str):
+        """Get the current weather in a given location using latitude and longitude."""
+        response = f'''Weather in {location}:
+                    Good morning! Expect overcast skies with intermittent showers throughout the day. 
+                    Temperatures will range from a cool 15°C in the early hours to around 19°C by mid-afternoon.
+                    Light winds from the northeast at about 10 km/h will keep conditions mild. 
+                    It might be a good idea to carry an umbrella if you’re heading out. Stay dry and have a great day!
+                    '''
+        return response
+
+@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
+):
+    """Test the chat completion without streaming for the weather agent."""
+    
+    weather_agent = WeatherAgent(proxy_url)
+
+    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"]
+    # Process each query
+    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"
+        assert response[0].content[0].type == "text"
+        assert response[0].content[1].type == "tool_use"
+        assert cities[index] in response[0].content[1].input["location"].lower()
+
+        assert response[1].role == "assistant"
+        assert response[1].stop_reason == "end_turn"
+        assert cities[index] in response[1].content[0].text.lower()
+
+
+@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="No ANTHROPIC_API_KEY set")
+async def test_chat_completion_with_streaming(
+    context, explorer_api_url, proxy_url
+):
+    """Test the chat completion with streaming for the weather agent."""
+    weather_agent = WeatherAgent(proxy_url)
+
+    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"]
+
+
+    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"
+        assert response[0][1].name == "get_weather"
+        assert cities[index] in response[0][1].input["location"].lower()
+        
+        assert response[1][0].type == "text"
+        assert cities[index] in response[1][0].text.lower()

tests/anthropic/test_chat_without_tool_call.py

@@ -0,0 +1,82 @@
+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):
+    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",
+        )
+    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?",
+        "Tell me the history of New York within 100 words?",
+        "How's the weather in London next week?"
+    ]
+
+    # Process each query
+    for query in queries:
+        response = client.messages.create(
+                # system=self.system_prompt,
+                model="claude-3-5-sonnet-20241022",
+                max_tokens=1024,
+                messages=[{"role": "user", "content": query}],
+            )
+        response_text = response.content[0].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):    
+
+    cities = ["zurich", "new york", "london"]
+    queries = [
+        "Can you introduce Zurich city within 200 words?",
+        "Tell me the history of New York within 100 words?",
+        "How's the weather in London next week?"
+    ]
+    # Process each query
+    for index,query in enumerate(queries):
+        messages = [
+            {
+                "role": "user",
+                "content": query
+            }
+        ]
+        response_text = ""
+        attempt = 0
+  
+        with client.messages.stream(
+            model="claude-3-5-sonnet-20241022",
+            max_tokens=1024,
+            messages=messages,
+        ) as response:
+            for reply in response.text_stream:
+                response_text += reply
+            assert cities[index] in response_text.lower()
+    
+        assert response_text is not None
+        assert cities[queries.index(query)] in response_text.lower()

tests/anthropic/test_claude_agent_streaming.py

@@ -1,63 +0,0 @@
-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_plugins = ("pytest_asyncio")
-@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="No ANTHROPIC_API_KEY set")
-def test_streaming_response_without_toolcall(proxy_url):    
-    # Example queries
-    dataset_name = "claude_streaming_agent_test" + str(datetime.datetime.now().strftime("%Y%m%d%H%M%S"))
-    invariant_api_key = os.environ.get("INVARIANT_API_KEY")
-
-    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?",
-        "Tell me the history of New York within 100 words?",
-        "How's the weather in London next week?"
-    ]
-    # Process each query
-    for index,query in enumerate(queries):
-        messages = [
-        {
-            "role": "user",
-            "content": query
-        }
-        ]
-        response_text = ""
-        attempt = 0
-        while attempt<3:
-            try:
-                with client.messages.stream(
-                    model="claude-3-5-sonnet-20241022",
-                    max_tokens=1024,
-                    messages=messages,
-                    # stream = True
-                ) as response:
-                    for reply in response.text_stream:
-                        response_text += reply
-                    assert cities[index] in response_text.lower()
-                break
-            except httpx.RemoteProtocolError as e:
-                attempt += 1
-                print(f"Streming error on attempt {attempt}: {e}")
-        else:
-            print("Max retries reached. Exiting.")

tests/anthropic/test_claude_weather_agent.py

@@ -1,114 +0,0 @@
-import datetime
-import os
-from typing import Dict
-
-import anthropic
-import pytest
-from httpx import Client
-import sys
-sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-from util import *  # needed for pytest fixtures
-
-pytest_plugins = ("pytest_asyncio",)
-
-
-class WeatherAgent:
-    def __init__(self,proxy_url):
-        # proxy_url = "http://localhost"
-        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")
-        self.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",
-        )
-        self.get_weather_function = {
-            "name": "get_weather",
-            "description": "Get the current weather in a given location",
-            "input_schema": {
-                "type": "object",
-                "properties": {
-                    "location": {
-                        "type": "string",
-                        "description": "The city and state, e.g. San Francisco, CA",
-                    },
-                    "unit": {
-                        "type": "string",
-                        "enum": ["celsius", "fahrenheit"],
-                        "description": 'The unit of temperature, either "celsius" or "fahrenheit"',
-                    },
-                },
-                "required": ["location"],
-            },
-        }
-
-        # self.system_prompt = """You are an assistant that can perform weather searches using function calls.
-        #     When a user asks for weather information, respond with a JSON object specifying
-        #     the function name `get_weather` and the arguments latitude and longitude are needed."""
-
-    def get_response(self, user_query: str) -> Dict:
-        """
-        Get the response from the agent for a given user query for weather.
-        """
-        messages = [{"role": "user", "content": user_query}]
-        while True:
-            response = self.client.messages.create(
-                # system=self.system_prompt,
-                tools=[self.get_weather_function],
-                model="claude-3-5-sonnet-20241022",
-                max_tokens=1024,
-                messages=messages
-            )
-
-
-            # If there's tool call, Extract the tool call parameters from the response
-            if len(response.content) > 1 and response.content[1].type == "tool_use":
-                tool_call_params = response.content[1].input
-                tool_call_result = self.get_weather(tool_call_params["location"])
-                tool_call_id = response.content[1].id
-                messages.append({"role": response.role, "content": response.content})
-                messages.append(
-                    {
-                        "role": "user",
-                        "content": [
-                            {
-                                "type": "tool_result",
-                                "tool_use_id": tool_call_id,
-                                "content": tool_call_result,
-                            }
-                        ],
-                    }
-                )
-            else:
-                return response.content[0].text
-
-    def get_weather(self, location: str):
-        """Get the current weather in a given location using latitude and longitude."""
-        response = f'''Weather in {location}:
-                    Good morning! Expect overcast skies with intermittent showers throughout the day. Temperatures will range from a cool 15°C in the early hours to around 19°C by mid-afternoon. Light winds from the northeast at about 10 km/h will keep conditions mild. It might be a good idea to carry an umbrella if you’re heading out. Stay dry and have a great day!
-                    '''
-        return response
-
-
-@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="No ANTHROPIC_API_KEY set")
-def test_proxy_response(proxy_url):
-    """Test the proxy response for the weather agent."""
-    # Initialize agent with Anthropic API key
-    weather_agent = WeatherAgent(proxy_url)
-
-    # Example queries
-    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"]
-    # Process each query
-    for index, query in enumerate(queries):
-        response = weather_agent.get_response(query)
-        assert response is not None
-        assert cities[index] in response.lower()

tests/open_ai/test_chat_with_tool_call.py

@@ -7,7 +7,6 @@ import uuid
 
 import pytest
 from httpx import Client
-
 # add tests folder (parent) to sys.path
 from openai import OpenAI
 
@@ -206,17 +205,11 @@ async def test_chat_completion_with_tool_call_with_streaming(
     )
     final_response = {"role": "assistant", "content": ""}
     attempt = 0
-    while attempt<3:
-        try:
-            for chunk in chat_response_final:
-                if chunk.choices and chunk.choices[0].delta.content:
-                    final_response["content"] += chunk.choices[0].delta.content
-            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_final:
+        if chunk.choices and chunk.choices[0].delta.content:
+            final_response["content"] += chunk.choices[0].delta.content
+
     # Fetch the trace ids for the dataset
     traces_response = await context.request.get(
         f"{explorer_api_url}/api/v1/dataset/byuser/developer/{dataset_name}/traces"