Update update_docx_with_pdf.py

update_docx_with_pdf.py

@@ -1,20 +1,36 @@
-import openai
+from openai import OpenAI
 import json
 import os
 
 def update_json_with_pdf(word_json_file, pdf_txt_file, output_file):
+    """
+    word_json_file: file-like object or file path (docx extraction JSON)
+    pdf_txt_file: file-like object or file path (PDF plain text)
+    output_file: file-like object (opened for writing) or file path
+    """
     # --- Load files ---
-    with open(word_json_file, "r", encoding="utf-8") as f:
-        word_json = f.read()
-    with open(pdf_txt_file, "r", encoding="utf-8") as f:
-        pdf_txt = f.read()
+    def read_any(f):
+        if hasattr(f, "read"):
+            f.seek(0)
+            content = f.read()
+            if isinstance(content, bytes):
+                content = content.decode("utf-8")
+            return content
+        else:
+            with open(f, "r", encoding="utf-8") as fh:
+                return fh.read()
+
+    word_json = read_any(word_json_file)
+    pdf_txt = read_any(pdf_txt_file)
 
     # --- Build prompt ---
     user_prompt = f"""
 Here is a JSON template. It contains only the fields that need updating:
 {word_json}
+
 Here is the extracted text from a PDF:
 {pdf_txt}
+
 Instructions:
 - ONLY update the fields present in the JSON template, using information from the PDF text.
 - DO NOT add any extra fields, and do not change the JSON structure.
@@ -26,7 +42,8 @@ Instructions:
     api_key = os.environ.get("OPENAI_API_KEY")
     if not api_key:
         raise RuntimeError("OPENAI_API_KEY not found in environment variables!")
-    client = openai.OpenAI(api_key=api_key)
+    client = OpenAI(api_key=api_key)
+
     response = client.chat.completions.create(
         model="gpt-4o",
         messages=[
@@ -42,9 +59,13 @@ Instructions:
     # --- Try to parse as JSON ---
     try:
         parsed = json.loads(updated_json_str)
-        with open(output_file, "w", encoding="utf-8") as f:
-            json.dump(parsed, f, indent=2, ensure_ascii=False)
-        print("✅ JSON updated and saved to", output_file)
+        if hasattr(output_file, "write"):
+            json.dump(parsed, output_file, indent=2, ensure_ascii=False)
+            output_file.flush()
+        else:
+            with open(output_file, "w", encoding="utf-8") as f:
+                json.dump(parsed, f, indent=2, ensure_ascii=False)
+        print("✅ JSON updated and saved to", getattr(output_file, "name", output_file))
     except Exception as e:
         print("⚠️ Model did not return valid JSON. Raw output below:\n")
         print(updated_json_str)