feat: Added tool_use support

Author: 7HR4IZ3

examples/anthrophic.py

@@ -1,14 +1,10 @@
 import os
 import asyncio
-from dotenv import load_dotenv
 from anthropic import Anthropic, AsyncAnthropic
 
 from lunary.anthrophic import monitor
 
-load_dotenv()
-
-
-def test_sync_non_streaming():
+def sync_non_streaming():
     client = Anthropic()
     monitor(client)
 
@@ -23,7 +19,7 @@ def test_sync_non_streaming():
     print(message.content)
 
 
-async def test_async_non_streaming():
+async def async_non_streaming():
     client = monitor(AsyncAnthropic())
 
     message = await client.messages.create(
@@ -37,7 +33,7 @@ async def test_async_non_streaming():
     print(message.content)
 
 
-def test_sync_streaming():
+def sync_streaming():
     client = monitor(Anthropic())
 
     stream = client.messages.create(
@@ -53,7 +49,7 @@ def test_sync_streaming():
         print(event)
 
 
-async def test_async_streaming():
+async def async_streaming():
     client = monitor(AsyncAnthropic())
 
     stream = await client.messages.create(
@@ -69,7 +65,7 @@ async def test_async_streaming():
         print(event)
 
 
-def test_sync_stream_helper():
+def sync_stream_helper():
     client = Anthropic()
     monitor(client)
 
@@ -84,7 +80,7 @@ def test_sync_stream_helper():
         for event in stream:
             print(event)
 
-async def test_async_stream_helper():
+async def async_stream_helper():
     client = monitor(AsyncAnthropic())
 
     async with client.messages.stream(
@@ -97,14 +93,15 @@ async def test_async_stream_helper():
         ],
         model="claude-3-opus-20240229",
     ) as stream:
-        async for event in stream:
-            print(event)
+        async for text in stream.text_stream:
+            print(text, end="", flush=True)
+        print()
 
     message = await stream.get_final_message()
     print(message.to_json())
 
 
-def test_extra_arguments():
+def extra_arguments():
     client = Anthropic()
     monitor(client)
 
@@ -129,13 +126,120 @@ def test_extra_arguments():
     print(message.content)
 
 
-# test_sync_non_streaming()
-# test_asyncio.run(async_non_streaming())
+def anthrophic_bedrock():
+    from anthropic import AnthropicBedrock
+
+    client = AnthropicBedrock()
+
+    message = client.messages.create(
+        max_tokens=1024,
+        messages=[
+            {
+                "role": "user",
+                "content": "Hello!",
+            }
+        ],
+        model="anthropic.claude-3-sonnet-20240229-v1:0",
+    )
+    print(message)
+
+def tool_calls():
+    from anthropic import Anthropic
+    from anthropic.types import ToolParam, MessageParam
+
+    client = monitor(Anthropic())
+
+    user_message: MessageParam = {
+        "role": "user",
+        "content": "What is the weather in San Francisco, California?",
+    }
+    tools: list[ToolParam] = [
+        {
+            "name": "get_weather",
+            "description": "Get the weather for a specific location",
+            "input_schema": {
+                "type": "object",
+                "properties": {"location": {"type": "string"}},
+            },
+        }
+    ]
+
+    message = client.messages.create(
+        model="claude-3-opus-20240229",
+        max_tokens=1024,
+        messages=[user_message],
+        tools=tools,
+    )
+    print(f"Initial response: {message.model_dump_json(indent=2)}")
+
+    assert message.stop_reason == "tool_use"
+
+    tool = next(c for c in message.content if c.type == "tool_use")
+    response = client.messages.create(
+        model="claude-3-opus-20240229",
+        max_tokens=1024,
+        messages=[
+            user_message,
+            {"role": message.role, "content": message.content},
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "tool_result",
+                        "tool_use_id": tool.id,
+                        "content": [{"type": "text", "text": "The weather is 73f"}],
+                    }
+                ],
+            },
+        ],
+        tools=tools,
+    )
+    print(f"\nFinal response: {response.model_dump_json(indent=2)}")
+
+
+async def async_tool_calls():
+    client = monitor(AsyncAnthropic())
+    async with client.messages.stream(
+        max_tokens=1024,
+        model="claude-3-haiku-20240307",
+        tools=[
+            {
+                "name": "get_weather",
+                "description": "Get the weather at a specific 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": "Unit for the output",
+                        },
+                    },
+                    "required": ["location"],
+                },
+            }
+        ],
+        messages=[{"role": "user", "content": "What is the weather in SF?"}],
+    ) as stream:
+        async for event in stream:
+            if event.type == "input_json":
+                print(f"delta: {repr(event.partial_json)}")
+                print(f"snapshot: {event.snapshot}")
+
+
+# sync_non_streaming()
+# asyncio.run(async_non_streaming())
+
+# sync_streaming()
+# asyncio.run(async_streaming())
+
+# extra_arguments()
 
-# test_sync_streaming()
-# test_asyncio.run(async_streaming())
+# sync_stream_helper()
+# asyncio.run(async_stream_helper())
 
-# test_extra_arguments()
+# # anthrophic_bedrock()
 
-# test_sync_stream_helper()
-asyncio.run(test_async_stream_helper())
+# tool_calls()
+# asyncio.run(async_tool_calls())

lunary/__init__.py

@@ -143,7 +143,6 @@ def track_event(
             "appId": app_id
         }
 
-
         if callback_queue is not None:
             callback_queue.append(event)
         else: