Add anthropic AI Usage

README.md

@@ -56,6 +56,7 @@ Zen for Python 3 is compatible with:
 
 ### AI SDKs
 * ✅ [`openai`](https://pypi.org/project/openai) ^1.0
+* ✅ [`anthropic`](https://pypi.org/project/anthropic/)
 
 ## Reporting to your Aikido Security dashboard
 

aikido_zen/__init__.py

@@ -68,5 +68,6 @@ def protect(mode="daemon"):
 
     # Import AI sinks
     import aikido_zen.sinks.openai
+    import aikido_zen.sinks.anthropic
 
     logger.info("Zen by Aikido v%s starting.", PKG_VERSION)

aikido_zen/sinks/anthropic.py

@@ -0,0 +1,21 @@
+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
+
+
+@after
+def _messages_create(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)

aikido_zen/sinks/tests/anthropic_test.py

@@ -0,0 +1,46 @@
+import os
+
+import pytest
+import aikido_zen.sinks.anthropic
+import anthropic
+
+from aikido_zen.thread.thread_cache import get_cache
+
+skip_no_api_key = pytest.mark.skipif(
+    "ANTHROPIC_API_KEY" not in os.environ,
+    reason="ANTHROPIC_API_KEY environment variable not set",
+)
+
+
+@pytest.fixture(autouse=True)
+def setup():
+    get_cache().reset()
+    yield
+    get_cache().reset()
+
+
+def get_ai_stats():
+    return get_cache().ai_stats.get_stats()
+
+
+@skip_no_api_key
+def test_anthropic_messages_create():
+    client = anthropic.Anthropic()
+    response = client.messages.create(
+        model="claude-3-opus-20240229",
+        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-3-opus-20240229"
+    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

pyproject.toml

@@ -81,6 +81,7 @@ pygments = "^2.18.0"
 lxml = "^5.4.0"
 clickhouse-driver = "^0.2.9"
 openai = "^1.85.0"
+anthropic = "^0.54.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]