Update anthropic test

Author: antonpirker

test-anthropic/main.py

@@ -3,12 +3,14 @@
 from anthropic import Anthropic
 
 import sentry_sdk
-from sentry_sdk.integrations.anthropic import AnthropicIntegration
 from sentry_sdk.consts import SPANTEMPLATE
+from sentry_sdk.integrations.anthropic import AnthropicIntegration
 
 
 @sentry_sdk.trace(name="Custom AI Agent", template=SPANTEMPLATE.AI_AGENT)
-def my_pipeline(client):
+def my_custom_agent(client):
+    print("~~~ Starting my_custom_agent ~~~")
+
     # Sync create message
     message = client.messages.create(
         messages=[
@@ -20,8 +22,8 @@ def my_pipeline(client):
         model="claude-3-5-haiku-latest",
         max_tokens=1024,
     )
-    print("Message:")
-    print(message.dict())
+    print("~~~ First result (sync message): ~~~")
+    print(message.model_dump())
 
     # Sync create streaming message
     stream = client.messages.create(
@@ -35,9 +37,11 @@ def my_pipeline(client):
         max_tokens=1024,
         stream=True,
     )
-    print("Message (Stream):")
+    print("~~~ Second result (sync streaming message): ~~~")
     for event in stream:
-        print(event.dict())
+        print(event.model_dump())
+
+    print("~~~ Done ~~~")
 
 
 def main():
@@ -56,8 +60,8 @@ def main():
         api_key=os.environ.get("ANTHROPIC_API_KEY"),
     )
 
-    with sentry_sdk.start_transaction(name="anthropic-sync"):
-        my_pipeline(client)
+    # with sentry_sdk.start_transaction(name="anthropic-sync"):
+    my_custom_agent(client)
 
 
 if __name__ == "__main__":

test-anthropic/main_async.py

@@ -9,7 +9,9 @@
 
 
 @sentry_sdk.trace(name="Custom AI Agent", template=SPANTEMPLATE.AI_AGENT)
-async def my_pipeline(client):
+async def my_custom_agent(client):
+    print("~~~ Starting my_custom_agent ~~~")
+
     # Async create message
     message = await client.messages.create(
         messages=[
@@ -21,7 +23,8 @@ async def my_pipeline(client):
         model="claude-3-haiku-20240307",
         max_tokens=1024,
     )
-    print(message.dict())
+    print("~~~ First result (async message): ~~~")
+    print(message.model_dump())
 
     # Async create streaming message
     stream = await client.messages.create(
@@ -35,8 +38,11 @@ async def my_pipeline(client):
         max_tokens=1024,
         stream=True,
     )
+    print("~~~ Second result (async streaming message): ~~~")
     async for event in stream:
-        print(event.dict())
+        print(event.model_dump())
+
+    print("~~~ Done ~~~")
 
 
 async def main():
@@ -56,7 +62,7 @@ async def main():
     )
 
     with sentry_sdk.start_transaction(name="anthropic-async"):
-        await my_pipeline(client)
+        await my_custom_agent(client)
 
 
 asyncio.run(main())

test-anthropic/main_tool.py

@@ -3,35 +3,35 @@
 from anthropic import Anthropic
 
 import sentry_sdk
-from sentry_sdk.integrations.anthropic import AnthropicIntegration
 from sentry_sdk.consts import SPANTEMPLATE
+from sentry_sdk.integrations.anthropic import AnthropicIntegration
 
 
+@sentry_sdk.trace(name="get_weather", template=SPANTEMPLATE.AI_TOOL)
 def get_weather(location):
-    return f"It is sunny with a high of 23°C."
+    return "It is sunny with a high of 23°C."
 
 
-# Define the tool
-tools = [
-      {
-        "name": "get_weather",
-        "description": "Get the current weather in a given location",
-        "input_schema": {
-          "type": "object",
-          "properties": {
+tools = [{
+    "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"
+                "type": "string",
+                "description": "The city and state, e.g. San Francisco, CA"
             }
-          },
-          "required": ["location"]
-        }
-      }
-    ]
+        },
+        "required": ["location"]
+    }
+}]
 
 
 @sentry_sdk.trace(name="Custom AI Agent", template=SPANTEMPLATE.AI_AGENT)
-def my_pipeline(client):
+def my_custom_agent(client):
+    print("~~~ Starting my_custom_agent ~~~")
+
     # Sync create message with tools
     message = client.messages.create(
         messages=[
@@ -45,19 +45,21 @@ def my_pipeline(client):
         max_tokens=1024,
         temperature=0,
     )
-    print("Message:")
-    print(message.dict())
+    print("~~~ First result: ~~~")
+    print(message.model_dump())
+
     # If model wants to run a tool, run it.
     if message.stop_reason == "tool_use":
+        print("~~~ Tool use detected ~~~")
         tool_use_block = message.content[1]
         function_name = tool_use_block.name
         tool_args = tool_use_block.input
 
         if function_name in globals() and callable(globals()[function_name]):
             tool_result = globals()[function_name](**tool_args)
-            print(f"Tool result: {tool_result}")
+            print(f"~~~ Tool result: {tool_result}")
 
-            # Continue the conversation with the tool result
+            # Continue the conversation including the tool result
             response = client.messages.create(
                 messages=[
                     {
@@ -85,8 +87,9 @@ def my_pipeline(client):
                 temperature=0,
             )
 
-            print("Final response:")
+            print("~~~ Final result: ~~~")
             print(response.content[0].text)
+            print("~~~ Done ~~~")
 
 
 def main():
@@ -106,7 +109,7 @@ def main():
     )
 
     with sentry_sdk.start_transaction(name="anthropic-sync-tool"):
-        my_pipeline(client)
+        my_custom_agent(client)
 
 
 if __name__ == "__main__":