Refactor openai model

Author: mbaluda

python/ql/lib/semmle/python/frameworks/OpenAI.qll

@@ -7,14 +7,57 @@ private import python
 private import semmle.python.ApiGraphs
 
 /**
- * Provides models for Agent (instances of the `agents.Agent` class).
+ * Provides models for Agent (instances of the `openai.OpenAI` class).
  *
- * See https://github.com/openai/openai-agents-python.
+ * See https://github.com/openai/openai-python.
  */
-module Agent {
-  /** Gets a reference to the `agents.Agent` class. */
-  API::Node classRef() { result = API::moduleImport("agents").getMember("Agent") }
+module OpenAI {
+  /** Gets a reference to the `openai.OpenAI` class. */
+  API::Node classRef() {
+    result =
+      API::moduleImport("openai").getMember(["OpenAI", "AsyncOpenAI", "AzureOpenAI"]).getReturn()
+  }
 
-  /** Gets a reference to a potential property of `agents.Agent` called instructions which refers to the system prompt. */
-  API::Node sink() { result = classRef().getACall().getKeywordParameter("instructions") }
+  /** Gets a reference to a potential property of `openai.OpenAI` called instructions which refers to the system prompt. */
+  API::Node getContentNode() {
+    exists(API::Node content |
+      content =
+        classRef()
+            .getMember("responses")
+            .getMember("create")
+            .getKeywordParameter(["input", "instructions"]) or
+      content =
+        classRef()
+            .getMember("responses")
+            .getMember("create")
+            .getKeywordParameter(["input", "instructions"])
+            .getASubscript()
+            .getSubscript("content") or
+      content =
+        classRef()
+            .getMember("realtime")
+            .getMember("connect")
+            .getReturn()
+            .getMember("conversation")
+            .getMember("item")
+            .getMember("create")
+            .getKeywordParameter("item")
+            .getSubscript("content") or
+      content =
+        classRef()
+            .getMember("chat")
+            .getMember("completions")
+            .getMember("create")
+            .getKeywordParameter("messages")
+            .getASubscript()
+            .getSubscript("content")
+    |
+      // content
+      if not exists(content.getASubscript())
+      then result = content
+      else
+        // content.text
+        result = content.getASubscript().getSubscript("text")
+    )
+  }
 }

python/ql/lib/semmle/python/security/dataflow/PromptInjectionCustomizations.qll

@@ -10,7 +10,7 @@ private import semmle.python.Concepts
 private import semmle.python.dataflow.new.RemoteFlowSources
 private import semmle.python.dataflow.new.BarrierGuards
 private import semmle.python.frameworks.data.ModelsAsData
-private import semmle.python.ApiGraphs
+private import semmle.python.frameworks.OpenAI
 
 /**
  * Provides default sources, sinks and sanitizers for detecting
@@ -51,48 +51,7 @@ module PromptInjection {
 
   private class PromptContentSink extends Sink {
     PromptContentSink() {
-      exists(API::Node openai, API::Node content |
-        openai =
-          API::moduleImport("openai")
-              .getMember(["OpenAI", "AsyncOpenAI", "AzureOpenAI"])
-              .getReturn() and
-        content =
-          [
-            openai
-                .getMember("responses")
-                .getMember("create")
-                .getKeywordParameter(["input", "instructions"]),
-            openai
-                .getMember("responses")
-                .getMember("create")
-                .getKeywordParameter(["input", "instructions"])
-                .getASubscript()
-                .getSubscript("content"),
-            openai
-                .getMember("realtime")
-                .getMember("connect")
-                .getReturn()
-                .getMember("conversation")
-                .getMember("item")
-                .getMember("create")
-                .getKeywordParameter("item")
-                .getSubscript("content"),
-            openai
-                .getMember("chat")
-                .getMember("completions")
-                .getMember("create")
-                .getKeywordParameter("messages")
-                .getASubscript()
-                .getSubscript("content")
-          ]
-      |
-        // content
-        if not exists(content.getASubscript())
-        then this = content.asSink()
-        else
-          // content.text
-          this = content.getASubscript().getSubscript("text").asSink()
-      )
+      this = OpenAI::getContentNode().asSink()
     }
   }