Merge branch 'main' into mistralai-ai-usage

Author: bitterpanda63

.github/workflows/end2end.yml

@@ -20,6 +20,7 @@ jobs:
   end2end-test:
     runs-on: ubuntu-latest
     continue-on-error: true
+    timeout-minutes: 15
     needs: build
     strategy:
       matrix:

.github/workflows/unit-test.yml

@@ -6,6 +6,7 @@ jobs:
   test:
     runs-on: ubuntu-latest
     continue-on-error: true
+    timeout-minutes: 15
     strategy:
       # Don't cancel jobs if one fails
       fail-fast: false

README.md

@@ -56,7 +56,9 @@ Zen for Python 3 is compatible with:
 
 ### AI SDKs
 * ✅ [`openai`](https://pypi.org/project/openai) ^1.0
+* ✅ [`anthropic`](https://pypi.org/project/anthropic/)
 * ✅ [`mistralai`](https://pypi.org/project/mistralai) ^1.0.0
+
 ## Reporting to your Aikido Security dashboard
 
 > Aikido is your no nonsense application security platform. One central system that scans your source code & cloud, shows you what vulnerabilities matter, and how to fix them - fast. So you can get back to building.

aikido_zen/__init__.py

@@ -2,6 +2,8 @@
 Aggregates from the different modules
 """
 
+import os
+
 # Re-export functions :
 from aikido_zen.context.users import set_user
 from aikido_zen.middleware import should_block_request
@@ -17,7 +19,7 @@
 from aikido_zen.helpers.aikido_disabled_flag_active import aikido_disabled_flag_active
 
 
-def protect(mode="daemon"):
+def protect(mode="daemon", token=""):
     """
     Mode can be set to :
     - daemon : Default, imports sinks/sources and starts background_process
@@ -28,6 +30,9 @@ def protect(mode="daemon"):
     if aikido_disabled_flag_active():
         # Do not run any aikido code when the disabled flag is on
         return
+    if token:
+        os.environ["AIKIDO_TOKEN"] = token
+
     if mode in ("daemon", "daemon_only"):
         start_background_process()
     if mode == "daemon_only":
@@ -68,6 +73,7 @@ def protect(mode="daemon"):
 
     # Import AI sinks
     import aikido_zen.sinks.openai
+    import aikido_zen.sinks.anthropic
     import aikido_zen.sinks.mistralai
 
     logger.info("Zen by Aikido v%s starting.", PKG_VERSION)

aikido_zen/init_test.py

@@ -1,6 +1,9 @@
 import pytest
+
+import aikido_zen
 from aikido_zen import protect
 from aikido_zen.background_process import get_comms, reset_comms
+from aikido_zen.helpers.token import get_token_from_env
 
 
 def test_protect_with_django(monkeypatch, caplog):
@@ -9,3 +12,8 @@ def test_protect_with_django(monkeypatch, caplog):
     assert "starting" in caplog.text
     reset_comms()
     assert get_comms() == None
+
+
+def test_protect_sets_token():
+    aikido_zen.protect(token="MY_TOKEN_1")
+    assert get_token_from_env().token == "MY_TOKEN_1"

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/builtins_import.py

@@ -14,10 +14,11 @@ def _import(func, instance, args, kwargs, return_value):
         return
     name = getattr(return_value, "__package__")
 
-    if not name or "." in name:
-        # Make sure the name exists and that it's not a submodule
+    if not name:
+        # Make sure the name exists
         return
-    if name == "importlib_metadata":
+    name = name.split(".")[0]  # Remove submodules
+    if name == "importlib" or name == "importlib_metadata":
         # Avoid circular dependencies
         return
 

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

end2end/django_mysql_test.py

@@ -88,5 +88,6 @@ def test_initial_heartbeat():
             "method": "POST",
             "path": "/app/create"
         }], 
-        {"aborted":0,"attacksDetected":{"blocked":2,"total":2},"total":3}
+        {"aborted":0,"attacksDetected":{"blocked":2,"total":2},"total":3},
+                       {'asgiref', 'regex', 'mysqlclient', 'sqlparse', 'aikido_zen', 'django'}
     )

end2end/server/check_events_from_mock.py

@@ -21,8 +21,13 @@ def validate_started_event(event, stack, dry_mode=False, serverless=False, os_na
     # if stack is not None:
     #    assert set(event["agent"]["stack"]) == set(stack)
 
-def validate_heartbeat(event, routes, req_stats):
+def validate_heartbeat(event, routes=None, req_stats=None, packages=None):
     assert event["type"] == "heartbeat", f"Expected event type 'heartbeat', but got '{event['type']}'"
-    assert event["routes"] == routes, f"Expected routes '{routes}', but got '{event['routes']}'"
-    assert event["stats"]["requests"] == req_stats, f"Expected request stats '{req_stats}', but got '{event['stats']['requests']}'"
+    if packages:
+        package_names = set(map(lambda x: x["name"], event["packages"]))
+        assert package_names == packages, f"Expected {packages} but got {package_names}"
+    if routes:
+        assert event["routes"] == routes, f"Expected routes '{routes}', but got '{event['routes']}'"
+    if req_stats:
+        assert event["stats"]["requests"] == req_stats, f"Expected request stats '{req_stats}', but got '{event['stats']['requests']}'"