Add support for chat completions

bitterpanda63 · bitterpanda63 · commit 0f7ec71156a7 · 2025-06-11T07:53:14.000Z
diff --git a/aikido_zen/sinks/openai.py b/aikido_zen/sinks/openai.py
@@ -1,11 +1,17 @@
+"""
+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(func, instance, args, kwargs, return_value):
-    op = "openai.resources.responses.responses.Responses.create"
+def _create_responses(func, instance, args, kwargs, return_value):
+    op = f"openai.resources.responses.responses.Responses.create"
     register_call(op, "ai_op")
 
     on_ai_call(
@@ -16,10 +22,24 @@ def _create(func, instance, args, kwargs, return_value):
     )
 
 
-@on_import("openai.resources.responses.responses", "openai", version_requirement="1.0")
-def patch(m):
-    """
-    patching module openai
-    - patches function create(...) on Responses class, to inspect response
-    """
-    patch_function(m, "Responses.create", _create)
+@after
+def _create_completions(func, instance, args, kwargs, return_value):
+    op = f"openai.resources.chat.completions.completions.Completions.create"
+    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)
+
+
+@on_import("openai.resources.chat.completions.completions", "openai", "1.0")
+def patch_chat_completions(m):
+    patch_function(m, "Completions.create", _create_completions)
diff --git a/sample-apps/flask-openai/app.py b/sample-apps/flask-openai/app.py
@@ -65,3 +65,25 @@ def ask_openai():
     answer = response.output_text
 
     return render_template('ask_openai.html', question=question, answer=answer)
+
+@app.route("/ask_openai_completions", methods=['GET'])
+def show_ask_openai_completions_form():
+    return render_template('ask_openai.html')
+
+@app.route("/ask_openai_completions", methods=['POST'])
+def ask_openai_completions():
+    question = request.form['question']
+
+    completion = client.chat.completions.create(
+        model="gpt-4o",
+        messages=[
+            {"role": "developer", "content": "Talk like a pirate."},
+            {
+                "role": "user",
+                "content": question,
+            },
+        ],
+    )
+    answer = completion.choices[0].message.content
+
+    return render_template('ask_openai.html', question=question, answer=answer)
diff --git a/sample-apps/flask-openai/templates/ask_openai.html b/sample-apps/flask-openai/templates/ask_openai.html
@@ -6,7 +6,7 @@
 </head>
 <body>
     <h1>Ask OpenAI</h1>
-    <form method="POST" action="/ask_openai">
+    <form method="POST">
         <label for="question">Question:</label><br>
         <input type="text" id="question" name="question" required><br>
         <input type="submit" value="Submit">
