Merge pull request #53 from DARPA-ASKEM/4841-extract-interventions-fr…

…om-document

Add new interventions from paper method

gollm/entities.py

@@ -9,6 +9,11 @@ class ConfigureModelDocument(BaseModel):
     amr: str  # expects AMR in a stringified JSON object
 
 
+class InterventionsFromDocument(BaseModel):
+    research_paper: str
+    amr: str  # expects AMR in a stringified JSON object
+
+
 class ConfigureModelDataset(BaseModel):
     dataset: List[str]
     amr: str  # expects AMR in a stringified JSON object

gollm/openai/prompts/interventions_from_document.py

@@ -0,0 +1,36 @@
+INTERVENTIONS_FROM_DOCUMENT_PROMPT = """
+You are a helpful agent designed to find intervention policies for a given AMR model described in a research paper.
+For context, intervention policies can include multiple interventions that include only static interventions or dynamic interventions.
+Static interventions are applied at a specific point in time and permanently change the value of a specific parameter or state.
+Dynamic interventions try to keep a specific parameter from going above or below a threshold value.
+
+Use the following AMR model JSON file as a reference:
+
+---START AMR MODEL JSON---
+{amr}
+---END AMR MODEL JSON---
+
+Use the following user-provided text as the research paper to answer the query:
+
+---START USER-PROVIDED TEXT---
+{research_paper}
+---END USER-PROVIDED TEXT---
+
+Assume that the user-provided text describes multiple intervention policies to apply to the model.
+For each intervention policy, create a list of interventions depending on what the text describes.
+Be sure to use a meaningful descriptive name for the intervention policy, instead of 'intervention_1' and 'intervention_2'.
+
+For each policy described in the paper, create an interventionPolicy object. To do this, follow the instructions below:
+1.	Create a value for `name` and `description` from the user-provided text.
+2.	For the description, provide a long-form description of the condition. If the description cannot be created from the user-provided text, set it to an empty string.
+3.	`model_id` id a UUID. If the AMR model has an id, you can use it. Otherwise, you can set as the nil UUID "00000000-0000-0000-0000-000000000000".
+4.	For each intervention specified in the policy create an intervention object with the following rules.
+    a.	Create a value for `name` from the user-provided text.
+    b.	`appliedTo` should reference the id of the parameter or initial state of the AMR Model. If you cannot find an initial state or parameter that matches the intervention, do not create an intervention object.
+    c.	`type` should be either "state" or "parameter" depending on what the intervention is applied to.
+    d.	create a list of either static or dynamic interventions, but not both.
+
+Do not respond in full sentences; only create a JSON object that satisfies the JSON schema specified in the response format.
+
+Answer:
+"""

gollm/openai/schemas/configuration.json

@@ -3,6 +3,7 @@
   "$defs": {
     "modelConfiguration": {
       "type": "object",
+      "description": "state and parameter configurations for a model",
       "properties": {
         "description": {
           "type": [

gollm/openai/schemas/intervention_policy.json

@@ -0,0 +1,125 @@
+{
+  "$schema": "https://json-schema.org/draft/2020-12/schema",
+  "$defs": {
+    "intervention": {
+      "type": "object",
+      "description": "A list of interventions that can be applied to the model.",
+      "properties": {
+        "name": {
+          "type": "string"
+        },
+        "appliedTo": {
+          "type": "string"
+        },
+        "type": {
+          "type": "string",
+          "enum": [
+            "state",
+            "parameter"
+          ]
+        },
+        "staticInterventions": {
+          "type": "array",
+          "description": "A list of static interventions that can be applied to the model.",
+          "items": {
+            "type": "object",
+            "properties": {
+              "timestep": {
+                "type": "number",
+                "description": "The timestep at which the intervention is applied"
+              },
+              "value": {
+                "type": "number",
+                "description": "The new value of the state or parameter"
+              }
+            },
+            "required": [
+              "timestep",
+              "value"
+            ],
+            "additionalProperties": false
+          }
+        },
+        "dynamicInterventions": {
+          "type": "array",
+          "description": "A list of dynamic interventions that can be applied to the model.",
+          "items": {
+            "type": "object",
+            "properties": {
+              "parameter": {
+                "type": "string",
+                "description": "The parameter to which the intervention is applied"
+              },
+              "threshold": {
+                "type": "string",
+                "description": "The threshold value for the which the parameter can't go above or below"
+              },
+              "value": {
+                "type": "string",
+                "description": "The new value of the state or parameter"
+              },
+              "isGreaterThan": {
+                "type": "string",
+                "description": "Whether the parameter should be greater than or less than the threshold"
+              }
+            },
+            "required": [
+              "parameter",
+              "threshold",
+              "value",
+              "isGreaterThan"
+            ],
+            "additionalProperties": false
+          }
+        }
+      },
+      "required": [
+        "name",
+        "appliedTo",
+        "type",
+        "staticInterventions",
+        "dynamicInterventions"
+      ],
+      "additionalProperties": false
+    }
+  },
+  "type": "object",
+  "properties": {
+    "interventionPolicies": {
+      "type": "array",
+      "items": {
+        "type": "object",
+        "properties": {
+          "modelId": {
+            "type": "string",
+            "description": "The ID of the AMR model to which the intervention policy applies."
+          },
+          "name": {
+            "type": "string"
+          },
+          "description": {
+            "type": "string"
+          },
+          "interventions": {
+            "type": "array",
+            "items": {
+              "$ref": "#/$defs/intervention"
+            },
+            "additionalProperties": false
+          }
+        },
+        "required": [
+          "modelId",
+          "name",
+          "description",
+          "interventions"
+        ],
+        "additionalProperties": false
+      }
+    }
+  },
+  "required": [
+    "interventionPolicies"
+  ],
+  "additionalProperties": false
+}

gollm/openai/tool_utils.py

@@ -26,6 +26,7 @@
 from gollm.openai.prompts.model_card import INSTRUCTIONS
 from gollm.openai.prompts.model_meta_compare import MODEL_METADATA_COMPARE_PROMPT
 from gollm.openai.prompts.config_from_document import CONFIGURE_FROM_DOCUMENT_PROMPT
+from gollm.openai.prompts.interventions_from_document import INTERVENTIONS_FROM_DOCUMENT_PROMPT
 
 SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
 
@@ -37,6 +38,52 @@ def escape_curly_braces(text: str):
     return text.replace("{", "{{").replace("}", "}}")
 
 
+def interventions_from_document(research_paper: str, amr: str) -> dict:
+    print("Extracting and formatting research paper...")
+    research_paper = normalize_greek_alphabet(research_paper)
+
+    print("Uploading and validating intervention policy schema...")
+    config_path = os.path.join(SCRIPT_DIR, 'schemas', 'intervention_policy.json')
+    with open(config_path, 'r') as config_file:
+        response_schema = json.load(config_file)
+    validate_schema(response_schema)
+
+    print("Building prompt to extract model configurations from a reasearch paper...")
+    prompt = INTERVENTIONS_FROM_DOCUMENT_PROMPT.format(
+        amr=escape_curly_braces(amr),
+        research_paper=escape_curly_braces(research_paper)
+    )
+
+    print("Sending request to OpenAI API...")
+    client = OpenAI()
+    output = client.chat.completions.create(
+        model="gpt-4o-2024-08-06",
+        frequency_penalty=0,
+        max_tokens=4000,
+        presence_penalty=0,
+        seed=905,
+        temperature=0,
+        top_p=1,
+        response_format={
+            "type": "json_schema",
+            "json_schema": {
+                "name": "intervention_policy",
+                "schema": response_schema
+            }
+        },
+        messages=[
+            {"role": "user", "content": prompt},
+        ]
+    )
+
+    print("Received response from OpenAI API. Formatting response to work with HMI...")
+    output_json = json.loads(output.choices[0].message.content)
+
+    print("There are ", len(output_json["interventionPolicies"]), "intervention policies identified from the text.")
+
+    return output_json
+
+
 def model_config_from_document(research_paper: str, amr: str) -> dict:
     print("Extracting and formatting research paper...")
     research_paper = normalize_greek_alphabet(research_paper)