update known data

llm_api/openai_api.py

@@ -1,33 +1,49 @@
 import json
+
 from dotenv import load_dotenv
-from .utils import get_data_format, get_image_data
-from openai import OpenAI, BadRequestError
+from openai import BadRequestError, OpenAI
+
+from .constants import (
+    EXTRACT_INFO_HUMAN_MESSAGE,
+    EXTRACT_INFO_SYSTEM_MESSAGE,
+    FOLLOW_SCHEMA_HUMAN_MESSAGE,
+    FOLLOW_SCHEMA_SYSTEM_MESSAGE,
+)
 from .exceptions import RefusalError
-from .constants import EXTRACT_INFO_HUMAN_MESSAGE, EXTRACT_INFO_SYSTEM_MESSAGE,FOLLOW_SCHEMA_HUMAN_MESSAGE, FOLLOW_SCHEMA_SYSTEM_MESSAGE
+from .utils import get_data_format, get_image_data
 
 load_dotenv(override=True)
 client = OpenAI()
 
 
-def extract_info(img_paths, schema):
-    print('Extracting info via OpenAI...')
+def extract_info(img_paths, schema, known_data=None):
+    print("Extracting info via OpenAI...")
     text_content = [
         {
             "type": "text",
             "text": EXTRACT_INFO_HUMAN_MESSAGE,
         },
     ]
+
+    if known_data is not None:
+        text_content.append(
+            {
+                "type": "text",
+                "text": f'\nAlso exploit the known data: \n\n"{known_data}"',
+            }
+        )
+
     image_content = [
         {
             "type": "image_url",
             "image_url": {
-            "url": f"data:image/{get_image_data(img_path)};base64,{get_image_data(img_path)}",
+                "url": f"data:image/{get_image_data(img_path)};base64,{get_image_data(img_path)}",
             },
         }
         for img_path in img_paths
     ]
     response = client.beta.chat.completions.parse(
-    model="gpt-4o-2024-08-06",
+        model="gpt-4o-2024-08-06",
         messages=[
             {
                 "role": "system",
@@ -36,7 +52,7 @@ def extract_info(img_paths, schema):
             {
                 "role": "user",
                 "content": text_content + image_content,
-            }
+            },
         ],
         max_tokens=1000,
         response_format=schema,
@@ -46,25 +62,34 @@ def extract_info(img_paths, schema):
         # top_p=.0000000000000000000001
     )
     if response.choices[0].message.refusal:
-        raise RefusalError('OpenAI refused to respond to the request')
+        raise RefusalError("OpenAI refused to respond to the request")
 
     content = response.choices[0].message.content
-    parsed_data = json.loads(content) 
+    parsed_data = json.loads(content)
     model_data = schema.model_validate(parsed_data)
 
     return 200, model_data
 
 
-def follow_structure(json_info, schema):
-    print('Following structure via OpenAI...')
+def follow_structure(json_info, schema, known_data=None):
+    print("Following structure via OpenAI...")
     text_content = [
         {
             "type": "text",
             "text": FOLLOW_SCHEMA_HUMAN_MESSAGE.format(json_info=json_info),
         },
     ]
+
+    if known_data is not None:
+        text_content.append(
+            {
+                "type": "text",
+                "text": f'\nAlso Exploit the known data: \n\n"{known_data}"',
+            }
+        )
+
     response = client.beta.chat.completions.parse(
-    model="gpt-4o-2024-08-06",
+        model="gpt-4o-2024-08-06",
         messages=[
             {
                 "role": "system",
@@ -73,7 +98,7 @@ def follow_structure(json_info, schema):
             {
                 "role": "user",
                 "content": text_content,
-            }
+            },
         ],
         max_tokens=1000,
         response_format=schema,
@@ -84,11 +109,10 @@ def follow_structure(json_info, schema):
     )
 
     if response.choices[0].message.refusal:
-        raise RefusalError('OpenAI refused to respond to the request')
+        raise RefusalError("OpenAI refused to respond to the request")
 
     content = response.choices[0].message.content
-    parsed_data = json.loads(content) 
+    parsed_data = json.loads(content)
     model_data = schema.model_validate(parsed_data)
 
     return 200, model_data
-