diff --git a/README.md b/README.md index 947f444d..13e23bab 100644 --- a/README.md +++ b/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 diff --git a/aikido_zen/__init__.py b/aikido_zen/__init__.py index 168b4489..fcbe9b7a 100644 --- a/aikido_zen/__init__.py +++ b/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) diff --git a/aikido_zen/sinks/anthropic.py b/aikido_zen/sinks/anthropic.py new file mode 100644 index 00000000..690de6a3 --- /dev/null +++ b/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) diff --git a/aikido_zen/sinks/tests/anthropic_test.py b/aikido_zen/sinks/tests/anthropic_test.py new file mode 100644 index 00000000..2b11449a --- /dev/null +++ b/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 diff --git a/poetry.lock b/poetry.lock index a5e953a0..dde5a197 100644 --- a/poetry.lock +++ b/poetry.lock @@ -15,6 +15,31 @@ files = [ [package.dependencies] typing-extensions = {version = ">=4.0.0", markers = "python_version < \"3.9\""} +[[package]] +name = "anthropic" +version = "0.54.0" +description = "The official Python library for the anthropic API" +optional = false +python-versions = ">=3.8" +groups = ["dev"] +files = [ + {file = "anthropic-0.54.0-py3-none-any.whl", hash = "sha256:c1062a0a905daeec17ca9c06c401e4b3f24cb0495841d29d752568a1d4018d56"}, + {file = "anthropic-0.54.0.tar.gz", hash = "sha256:5e6f997d97ce8e70eac603c3ec2e7f23addeff953fbbb76b19430562bb6ba815"}, +] + +[package.dependencies] +anyio = ">=3.5.0,<5" +distro = ">=1.7.0,<2" +httpx = ">=0.25.0,<1" +jiter = ">=0.4.0,<1" +pydantic = ">=1.9.0,<3" +sniffio = "*" +typing-extensions = ">=4.10,<5" + +[package.extras] +bedrock = ["boto3 (>=1.28.57)", "botocore (>=1.31.57)"] +vertex = ["google-auth[requests] (>=2,<3)"] + [[package]] name = "anyio" version = "4.5.2" @@ -2261,4 +2286,4 @@ type = ["pytest-mypy"] [metadata] lock-version = "2.1" python-versions = "^3.8" -content-hash = "4c6f6438f25158b1a1f8401c2e55f3e7a2ce4f9f4fa542f349edbb41169fed2f" +content-hash = "b5a222795eac86eb7a904f9c4cf7b2b0690ab820cf1e67094760458499bb07ec" diff --git a/pyproject.toml b/pyproject.toml index 7e81d277..195d2ef4 100644 --- a/pyproject.toml +++ b/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"]