Add async messages wrapping for anthropic

aikido_zen/sinks/__init__.py

@@ -114,3 +114,24 @@ def decorator(func, instance, args, kwargs):
         return return_value
 
     return decorator
+
+
+def after_async(wrapper):
+    """
+    Surrounds an async patch with try-except, calls the original function and gives the return value to the patch
+    """
+
+    async def decorator(func, instance, args, kwargs):
+        return_value = await func(*args, **kwargs)  # Call the original function
+        try:
+            await wrapper(func, instance, args, kwargs, return_value)  # Call the patch
+        except AikidoException as e:
+            raise e  # Re-raise AikidoException
+        except Exception as e:
+            logger.debug(
+                "%s:%s wrapping-after error: %s", func.__module__, func.__name__, e
+            )
+
+        return return_value
+
+    return decorator

aikido_zen/sinks/anthropic.py

@@ -1,6 +1,6 @@
 from aikido_zen.helpers.on_ai_call import on_ai_call
 from aikido_zen.helpers.register_call import register_call
-from aikido_zen.sinks import on_import, patch_function, after
+from aikido_zen.sinks import on_import, patch_function, after, after_async
 
 
 @after
@@ -16,6 +16,20 @@ def _messages_create(func, instance, args, kwargs, return_value):
     )
 
 
+@after_async
+async def _messages_create_async(func, instance, args, kwargs, return_value):
+    op = f"anthropic.resources.messages.messages.Messages.create"
+    register_call(op, "ai_op")
+
+    on_ai_call(
+        provider="anthropic",
+        model=return_value.model,
+        input_tokens=return_value.usage.input_tokens,
+        output_tokens=return_value.usage.output_tokens,
+    )
+
+
 @on_import("anthropic.resources.messages")
 def patch(m):
     patch_function(m, "messages.Messages.create", _messages_create)
+    patch_function(m, "messages.AsyncMessages.create", _messages_create_async)

aikido_zen/sinks/tests/anthropic_test.py

@@ -3,6 +3,7 @@
 import pytest
 import aikido_zen.sinks.anthropic
 import anthropic
+import asyncio
 
 from aikido_zen.thread.thread_cache import get_cache
 
@@ -27,7 +28,7 @@ def get_ai_stats():
 def test_anthropic_messages_create():
     client = anthropic.Anthropic()
     response = client.messages.create(
-        model="claude-3-opus-20240229",
+        model="claude-sonnet-4-20250514",
         max_tokens=20,
         messages=[
             {
@@ -38,7 +39,31 @@ def test_anthropic_messages_create():
     )
     print(response)
 
-    assert get_ai_stats()[0]["model"] == "claude-3-opus-20240229"
+    assert get_ai_stats()[0]["model"] == "claude-sonnet-4-20250514"
+    assert get_ai_stats()[0]["calls"] == 1
+    assert get_ai_stats()[0]["provider"] == "anthropic"
+    assert get_ai_stats()[0]["tokens"]["input"] == 21
+    assert get_ai_stats()[0]["tokens"]["output"] == 20
+    assert get_ai_stats()[0]["tokens"]["total"] == 41
+
+
+@skip_no_api_key
+@pytest.mark.asyncio
+async def test_anthropic_messages_create_async():
+    client = anthropic.AsyncAnthropic()
+    response = await client.messages.create(
+        model="claude-sonnet-4-20250514",
+        max_tokens=20,
+        messages=[
+            {
+                "role": "user",
+                "content": "Write the longest response possible, just as I am writing a long content",
+            }
+        ],
+    )
+    print(response)
+
+    assert get_ai_stats()[0]["model"] == "claude-sonnet-4-20250514"
     assert get_ai_stats()[0]["calls"] == 1
     assert get_ai_stats()[0]["provider"] == "anthropic"
     assert get_ai_stats()[0]["tokens"]["input"] == 21