Add AI wrapping for openai

README.md

@@ -54,6 +54,9 @@ Zen for Python 3 is compatible with:
 * ✅ [`motor`](https://pypi.org/project/motor/) (See `pymongo` version)
 * ✅ [`clickhouse-driver`](https://pypi.org/project/clickhouse-driver)
 
+### AI SDKs
+* ✅ [`openai`](https://pypi.org/project/openai) ^1.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

@@ -64,4 +64,7 @@ def protect(mode="daemon"):
     import aikido_zen.sinks.os_system
     import aikido_zen.sinks.subprocess
 
+    # Import AI sinks
+    import aikido_zen.sinks.openai
+
     logger.info("Zen by Aikido v%s starting.", PKG_VERSION)

aikido_zen/background_process/cloud_connection_manager/__init__.py

@@ -10,6 +10,7 @@
 from aikido_zen.storage.users import Users
 from aikido_zen.storage.hostnames import Hostnames
 from ..realtime.start_polling_for_changes import start_polling_for_changes
+from ...storage.ai_statistics import AIStatistics
 from ...storage.statistics import Statistics
 
 # Import functions :
@@ -46,6 +47,7 @@
         self.users = Users(1000)
         self.packages = {}
         self.statistics = Statistics()
+        self.ai_stats = AIStatistics()
         self.middleware_installed = False
 
         if isinstance(serverless, str) and len(serverless) == 0:
@@ -69,7 +71,11 @@
         This is run 1m after startup, and checks if we should send out
         a preliminary heartbeat with some stats.
         """
-        data_present = not self.statistics.empty() or len(self.routes.routes) > 0
+        data_present = (
+            not self.statistics.empty()
+            or len(self.routes.routes) > 0
+            or not self.ai_stats.empty()
+        )
         should_report_initial_stats = data_present and not self.conf.received_any_stats
         if should_report_initial_stats:
             self.send_heartbeat()

aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py

@@ -15,18 +15,21 @@
     users = connection_manager.users.as_array()
     routes = list(connection_manager.routes)
     outgoing_domains = connection_manager.hostnames.as_array()
+    ai_stats = connection_manager.ai_stats.get_stats()
 
     connection_manager.statistics.clear()
     connection_manager.users.clear()
     connection_manager.routes.clear()
     connection_manager.hostnames.clear()
+    connection_manager.ai_stats.clear()
     res = connection_manager.api.report(
         connection_manager.token,
         {
             "type": "heartbeat",
             "time": get_unixtime_ms(),
             "agent": connection_manager.get_manager_info(),
             "stats": stats,
+            "ai": ai_stats,
             "hostnames": outgoing_domains,
             "routes": routes,
             "users": users,

aikido_zen/background_process/commands/sync_data.py

@@ -46,6 +46,9 @@ def process_sync_data(connection_manager, data, conn, queue=None):
     # Sync stats
     connection_manager.statistics.import_from_record(data.get("stats", {}))
 
+    # Sync ai stats
+    connection_manager.ai_stats.import_list(data.get("ai_stats", []))
+
     if connection_manager.conf.last_updated_at > 0:
         # Only report data if the config has been fetched.
         return {

aikido_zen/helpers/on_ai_call.py

@@ -0,0 +1,7 @@
+from aikido_zen.thread.thread_cache import get_cache
+
+
+def on_ai_call(provider, model, input_tokens, output_tokens):
+    cache = get_cache()
+    if cache:
+        cache.ai_stats.on_ai_call(provider, model, input_tokens, output_tokens)

aikido_zen/sinks/openai.py

@@ -0,0 +1,47 @@
+"""
+patching module openai
+- patches function create(...) on Responses class, to inspect response
+- patches function create(...) on Completions class, to inspect response
+"""
+
+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 _create_responses(func, instance, args, kwargs, return_value):
+    op = f"openai.resources.responses.responses.Responses.{func.__name__}"
+    register_call(op, "ai_op")
+
+    on_ai_call(
+        provider="openai",
+        model=return_value.model,
+        input_tokens=return_value.usage.input_tokens,
+        output_tokens=return_value.usage.output_tokens,
+    )
+
+
+@after
+def _create_completions(func, instance, args, kwargs, return_value):
+    op = f"openai.resources.chat.completions.completions.Completions.{func.__name__}"
+    register_call(op, "ai_op")
+
+    on_ai_call(
+        provider="openai",
+        model=return_value.model,
+        input_tokens=return_value.usage.prompt_tokens,
+        output_tokens=return_value.usage.completion_tokens,
+    )
+
+
+@on_import("openai.resources.responses.responses", "openai", "1.0")
+def patch_responses(m):
+    patch_function(m, "Responses.create", _create_responses)
+    patch_function(m, "Responses.parse", _create_responses)
+
+
+@on_import("openai.resources.chat.completions.completions", "openai", "1.0")
+def patch_chat_completions(m):
+    patch_function(m, "Completions.create", _create_completions)
+    patch_function(m, "Completions.update", _create_completions)

aikido_zen/storage/ai_statistics.py

@@ -0,0 +1,56 @@
+import copy
+
+
+class AIStatistics:
+    def __init__(self):
+        self.calls = {}
+
+    def ensure_provider_stats(self, provider, model):
+        key = get_provider_key(provider, model)
+
+        if key not in self.calls:
+            self.calls[key] = {
+                "provider": provider,
+                "model": model,
+                "calls": 0,
+                "tokens": {
+                    "input": 0,
+                    "output": 0,
+                    "total": 0,
+                },
+            }
+
+        return self.calls[key]
+
+    def on_ai_call(self, provider, model, input_tokens, output_tokens):
+        if not provider or not model:
+            return
+
+        provider_stats = self.ensure_provider_stats(provider, model)
+        provider_stats["calls"] += 1
+        provider_stats["tokens"]["input"] += input_tokens
+        provider_stats["tokens"]["output"] += output_tokens
+        provider_stats["tokens"]["total"] += input_tokens + output_tokens
+
+    def get_stats(self):
+        return [copy.deepcopy(stats) for stats in self.calls.values()]
+
+    def import_list(self, ai_stats_list):
+        for new_entry in ai_stats_list:
+            existing_entry = self.ensure_provider_stats(
+                new_entry["provider"], new_entry["model"]
+            )
+            existing_entry["calls"] += new_entry["calls"]
+            existing_entry["tokens"]["input"] += new_entry["tokens"]["input"]
+            existing_entry["tokens"]["output"] += new_entry["tokens"]["output"]
+            existing_entry["tokens"]["total"] += new_entry["tokens"]["total"]
+
+    def clear(self):
+        self.calls.clear()
+
+    def empty(self):
+        return len(self.calls) == 0
+
+
+def get_provider_key(provider, model):
+    return f"{provider}:{model}"