Update app.py

app.py

@@ -1,26 +1,28 @@
-import openai
-import base64
+import os
 import io
 import time
+import base64
 import logging
 import fitz  # PyMuPDF
 from PIL import Image
 import gradio as gr
+from openai import OpenAI  # Use the OpenAI client that supports multimodal messages
+
+# Load API key from environment variable (secrets)
+HF_API_KEY = os.getenv("OPENAI_TOKEN")
+if not HF_API_KEY:
+    raise ValueError("HF_API_KEY environment variable not set")
+
+# Create the client pointing to the Hugging Face Inference endpoint
+client = OpenAI(
+    base_url="https://openrouter.ai/api/v1",
+    api_key=HF_API_KEY
+)
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-import os
-OPENROUTER_API_KEY = os.getenv("OPENAI_TOKEN")
-if not OPENROUTER_API_KEY:
-    raise ValueError("OPENROUTER_API_KEY environment variable not set")
-openai.api_key = OPENROUTER_API_KEY
-
-# Configure the OpenAI API to use OpenRouter
-openai.api_base = "https://openrouter.ai/api/v1"
-openai.api_key = OPENROUTER_API_KEY
-
 # -------------------------------
 # Document State and File Processing
 # -------------------------------
@@ -38,7 +40,7 @@ class DocumentState:
 doc_state = DocumentState()
 
 def process_pdf_file(file_path):
-    """Convert PDF to images and extract text using PyMuPDF."""
+    """Convert PDF pages to images and extract text using PyMuPDF."""
     try:
         doc = fitz.open(file_path)
         images = []
@@ -50,13 +52,12 @@ def process_pdf_file(file_path):
                 if page_text.strip():
                     text += f"Page {page_num + 1}:\n{page_text}\n\n"
                 
-                # Render page to an image
+                # Render page as an image with a zoom factor
                 zoom = 3
                 mat = fitz.Matrix(zoom, zoom)
                 pix = page.get_pixmap(matrix=mat, alpha=False)
                 img_data = pix.tobytes("png")
-                img = Image.open(io.BytesIO(img_data))
-                img = img.convert("RGB")
+                img = Image.open(io.BytesIO(img_data)).convert("RGB")
                 
                 # Resize if image is too large
                 max_size = 1600
@@ -64,7 +65,6 @@ def process_pdf_file(file_path):
                     ratio = max_size / max(img.size)
                     new_size = tuple(int(dim * ratio) for dim in img.size)
                     img = img.resize(new_size, Image.Resampling.LANCZOS)
-                
                 images.append(img)
             except Exception as e:
                 logger.error(f"Error processing page {page_num}: {str(e)}")
@@ -78,13 +78,13 @@ def process_pdf_file(file_path):
         raise
 
 def process_uploaded_file(file):
-    """Process uploaded file and update document state."""
+    """Process an uploaded file (PDF or image) and update document state."""
     try:
         doc_state.clear()
         if file is None:
             return "No file uploaded. Please upload a file."
         
-        # Get the file path and extension
+        # Get the file path from the Gradio upload (may be a dict or file-like object)
         if isinstance(file, dict):
             file_path = file["name"]
         else:
@@ -119,16 +119,17 @@ def process_uploaded_file(file):
         return "An error occurred while processing the file. Please try again."
 
 # -------------------------------
-# Bot Streaming Function Using OpenAI API
+# Bot Streaming Function Using the Multimodal API
 # -------------------------------
-def bot_streaming(prompt_option, max_new_tokens=4096):
+def bot_streaming(prompt_option, max_new_tokens=500):
     """
-    Generate a response using the OpenAI API.
-    
-    If an image is available, it is encoded in base64 and appended to the prompt.
+    Build a multimodal message payload and call the inference API.
+    The payload includes:
+      - A text segment (the selected prompt and any document context).
+      - If available, an image as a data URI (using a base64-encoded PNG).
     """
     try:
-        # Define predetermined prompts
+        # Predetermined prompts (you can adjust these as needed)
         prompts = {
             "NOC Timesheet": (
                 """Extract structured information from the provided timesheet. The extracted details should include:
@@ -173,333 +174,90 @@ Noc representative's date approval_date
 
 Noc representative status as approval_status
 
-The output should be formatted as a JSON instance that conforms to the JSON schema below.
-
-As an example, for the schema {"properties": {"foo": {"title": "Foo", "description": "a list of strings", "type": "array", "items": {"type": "string"}}}, "required": ["foo"]} the object {"foo": ["bar", "baz"]} is a well-formatted instance of the schema. The object {"properties": {"foo": ["bar", "baz"]}} is not well-formatted."""
+Format the output as valid JSON.
+                """
             ),
             "NOC Basic": (
                 "Based on the provided timesheet details, extract the following information:\n"
-                "   - Full name of the person\n"
-                "   - Position title of the person\n"
+                "   - Full name\n"
+                "   - Position title\n"
                 "   - Work location\n"
                 "   - Contractor's name\n"
                 "   - NOC ID\n"
-                "   - Month and year (in MM/YYYY format)"
+                "   - Month and year (MM/YYYY)"
             ),
             "Aramco Full structured": (
-                """You are a document parsing assistant designed to extract structured data from various document types, including invoices, timesheets, purchase orders, and travel bookings. Your goal is to return highly accurate, properly formatted JSON for each document type.
-General Rules:
-1. Always return ONLY valid JSON—no explanations, comments, or additional text.
-2. Use null for any fields that are not present or cannot be extracted.
-3. Ensure all JSON keys are enclosed in double quotes and properly formatted.
-4. Validate financial, time tracking, and contract details carefully before output.
-
-Extraction Instructions:
-
-1. Invoice:
-   - Parse and extract financial and invoice-specific details.
-   - JSON structure:
-     ```json
-     {
-         "invoice": {
-             "date": null,
-             "dueDate": null,
-             "accountNumber": null,
-             "invoiceNumber": null,
-             "customerContact": null,
-             "kintecContact": null,
-             "accountsContact": null,
-             "periodEnd": null,
-             "contractNo": null,
-             "specialistsName": null,
-             "rpoNumber": null,
-             "assignmentProject": null,
-             "workLocation": null,
-             "expenses": null,
-             "regularHours": null,
-             "overtime": null,
-             "mobilisationAllowance": null,
-             "dailyHousing": null,
-             "opPipTechnical": null,
-             "code": null,
-             "vatBasis": null,
-             "vatRate": null,
-             "vatAmount": null,
-             "totalExclVat": null,
-             "totalInclVat": null
-         }
-     }
-     ```
-
-2. Timesheet:
-   - Extract time tracking, work details, and approvals.
-   - JSON structure:
-     ```json
-     {
-         "timesheet": {
-             "Year": null,
-             "RPO_Number": null,
-             "PMC_Name": null,
-             "Project_Location": null,
-             "Project_and_Package": null,
-             "Month": null,
-             "Timesheet_Details": [
-                 {
-                     "Week": null,
-                     "Regular_Hours": null,
-                     "Overtime_Hours": null,
-                     "Total_Hours": null,
-                     "Comments": null
-                 },
-                 {
-                     "Week": null,
-                     "Regular_Hours": null,
-                     "Overtime_Hours": null,
-                     "Total_Hours": null,
-                     "Comments": null
-                 }
-             ],
-             "Monthly_Totals": {
-                 "Regular_Hours": null,
-                 "Overtime_Hours": null,
-                 "Total_Hours": null
-             },
-             "reviewedBy": {
-                 "name": null,
-                 "position": null,
-                 "date": null
-             },
-             "approvedBy": {
-                 "name": null,
-                 "position": null,
-                 "date": null
-             }
-         }
-     }
-     ```
-                
-3. Purchase Order:
-   - Extract contract and pricing details with precision.
-   - JSON structure:
-     ```json
-     {
-         "purchaseOrder": {
-             "contractNo": null,
-             "relPoNo": null,
-             "version": null,
-             "title": null,
-             "startDate": null,
-             "endDate": null,
-             "costCenter": null,
-             "purchasingGroup": null,
-             "contractor": null,
-             "location": null,
-             "workDescription": null,
-             "pricing": {
-                 "regularRate": null,
-                 "overtimeRate": null,
-                 "totalBudget": null
-             }
-         }
-     }
-     ```
-                
-4. Travel Booking:
-   - Parse travel-specific and employee information.
-   - JSON structure:
-     ```json
-     {
-         "travelBooking": {
-             "requestId": null,
-             "approvalStatus": null,
-             "employee": {
-                 "name": null,
-                 "id": null,
-                 "email": null,
-                 "firstName": null,
-                 "lastName": null,
-                 "gradeCodeGroup": null
-             },
-             "defaultManager": {
-                 "name": null,
-                 "email": null
-             },
-             "sender": {
-                 "name": null,
-                 "email": null
-             },
-             "travel": {
-                 "startDate": null,
-                 "endDate": null,
-                 "requestPolicy": null,
-                 "requestType": null,
-                 "employeeType": null,
-                 "travelActivity": null,
-                 "tripType": null
-             },
-             "cost": {
-                 "companyCode": null,
-                 "costObject": null,
-                 "costObjectId": null
-             },
-             "transport": {
-                 "type": null,
-                 "comments": null
-             },
-             "changeRequired": null,
-             "comments": null
-         }
-     }
-     ```
-                
-Use these structures for parsing documents and ensure compliance with the rules and instructions provided for each type."""
+                """You are a document parsing assistant designed to extract structured data from various documents such as invoices, timesheets, purchase orders, and travel bookings. Return only valid JSON with no extra text.
+                """
             ),
             "Aramco Timesheet only": (
                 """Extract time tracking, work details, and approvals.
-- JSON structure:
-```json
-{
-    "timesheet": {
-        "Year": null,
-        "RPO_Number": null,
-        "PMC_Name": null,
-        "Project_Location": null,
-        "Project_and_Package": null,
-        "Month": null,
-        "Timesheet_Details": [
-            {
-                "Week": null,
-                "Regular_Hours": null,
-                "Overtime_Hours": null,
-                "Total_Hours": null,
-                "Comments": null
-            },
-            {
-                "Week": null,
-                "Regular_Hours": null,
-                "Overtime_Hours": null,
-                "Total_Hours": null,
-                "Comments": null
-            }
-        ],
-        "Monthly_Totals": {
-            "Regular_Hours": null,
-            "Overtime_Hours": null,
-            "Total_Hours": null
-        },
-        "reviewedBy": {
-            "name": null,
-            "position": null,
-            "date": null
-        },
-        "approvedBy": {
-            "name": null,
-            "position": null,
-            "date": null
-        }
-    }
-}
-```"""
+Return a JSON object following the specified structure.
+                """
             ),
             "NOC Invoice": (
-                """You are a highly accurate data extraction system. Your task is to analyze the provided image of an invoice and extract all data, paying close attention to the structure and formatting of the document. Organize the extracted data in a clear, structured format, such as JSON. Do not invent any information. If a field cannot be read with high confidence, indicate that with "UNCLEAR" or a similar designation. Be as specific as possible, and do not summarize or combine fields unless explicitly indicated.
-
-Here's the expected output format, in JSON, with all required fields:
-
-```json
+                """You are a highly accurate data extraction system. Analyze the provided invoice image and extract all data into the following JSON format:
 {
-  "invoiceDetails": {
-    "pleaseQuote": "string",
-    "invoiceNumber": "string",
-    "workPeriod": "string",
-    "invoiceDate": "string",
-    "assignmentReference": "string"
-  },
-  "from": {
-    "companyName": "string",
-    "addressLine1": "string",
-    "addressLine2": "string",
-    "city": "string",
-    "postalCode": "string",
-    "country": "string"
-  },
-  "to": {
-    "companyName": "string",
-    "office": "string",
-    "floor": "string",
-    "building": "string",
-    "addressLine1": "string",
-    "poBox": "string",
-    "city": "string"
-  },
-  "services": [
-    {
-      "serviceDetails": "string",
-      "fromDate": "string",
-      "toDate": "string",
-      "currency": "string",
-      "fx": "string",
-      "noOfDays": "number or string (if range)",
-      "rate": "number",
-      "total": "number"
-    }
-  ],
-  "totals": {
-    "subTotal": "number",
-    "tax": "number",
-    "totalDue": "number"
-  },
-  "bankDetails": {
-    "bankName": "string",
-    "descriptionReferenceField": "string",
-    "bankAddress": "string",
-    "swiftBicCode": "string",
-    "ibanNumber": "string",
-    "accountNumber": "string",
-    "beneficiaryName": "string",
-    "accountCurrency": "string",
-    "expectedAmount": "string"
-  }
+  "invoiceDetails": { ... },
+  "from": { ... },
+  "to": { ... },
+  "services": [ ... ],
+  "totals": { ... },
+  "bankDetails": { ... }
 }
-```"""
+                """
             )
         }
         
-        # Retrieve the selected prompt
+        # Select the appropriate prompt
         selected_prompt = prompts.get(prompt_option, "Invalid prompt selected.")
         context = ""
-        if doc_state.current_doc_images:
-            if doc_state.current_doc_text:
-                context = f"\nDocument context:\n{doc_state.current_doc_text}"
+        if doc_state.current_doc_images and doc_state.current_doc_text:
+            context = "\nDocument context:\n" + doc_state.current_doc_text
         full_prompt = selected_prompt + context
+
+        # Build the message payload in the expected format.
+        # The content field is a list of objects—one for text, and (if an image is available) one for the image.
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": full_prompt
+                    }
+                ]
+            }
+        ]
         
-        # Create the messages list for the API call
-        messages = [{"role": "user", "content": full_prompt}]
-        
-        # If an image is available, encode it in base64 and append to the prompt
+        # If an image is available, encode it as a data URI and append it as an image_url message.
         if doc_state.current_doc_images:
             buffered = io.BytesIO()
             doc_state.current_doc_images[0].save(buffered, format="PNG")
-            img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-            messages[0]["content"] += f"\n[Image Data: {img_str}]"
+            img_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+            # Create a data URI (many APIs accept this format in place of a public URL)
+            data_uri = f"data:image/png;base64,{img_b64}"
+            messages[0]["content"].append({
+                "type": "image_url",
+                "image_url": {"url": data_uri}
+            })
         
-        # Call the OpenAI API with streaming enabled.
-        response = openai.ChatCompletion.create(
+        # Call the inference API with streaming enabled.
+        stream = client.chat.completions.create(
             model="qwen/qwen2.5-vl-72b-instruct:free",
             messages=messages,
             max_tokens=max_new_tokens,
-            stream=True,
+            stream=True
         )
         
         buffer = ""
-        for chunk in response:
-            if 'choices' in chunk:
-                delta = chunk['choices'][0].get('delta', {})
-                content = delta.get('content', '')
-                buffer += content
-                time.sleep(0.01)
-                yield buffer
-                
+        for chunk in stream:
+            # The response structure is similar to the reference: each chunk contains a delta.
+            delta = chunk.choices[0].delta.content
+            buffer += delta
+            time.sleep(0.01)
+            yield buffer
+
     except Exception as e:
         logger.error(f"Error in bot_streaming: {str(e)}")
         yield "An error occurred while processing your request. Please try again."
@@ -521,48 +279,21 @@ with gr.Blocks() as demo:
             label="Upload Document",
             file_types=[".pdf", ".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp"]
         )
-        upload_status = gr.Textbox(
-            label="Upload Status",
-            interactive=False
-        )
+        upload_status = gr.Textbox(label="Upload Status", interactive=False)
     
     with gr.Row():
         prompt_dropdown = gr.Dropdown(
             label="Select Prompt",
-            choices=[
-                "NOC Timesheet",
-                "NOC Basic",
-                "Aramco Full structured",
-                "Aramco Timesheet only",
-                "NOC Invoice"
-            ],
+            choices=["NOC Timesheet", "NOC Basic", "Aramco Full structured", "Aramco Timesheet only", "NOC Invoice"],
             value="NOC Timesheet"
         )
         generate_btn = gr.Button("Generate")
     
     clear_btn = gr.Button("Clear Document Context")
+    output_text = gr.Textbox(label="Output", interactive=False)
     
-    output_text = gr.Textbox(
-        label="Output",
-        interactive=False
-    )
-    
-    file_upload.change(
-        fn=process_uploaded_file,
-        inputs=[file_upload],
-        outputs=[upload_status]
-    )
-    
-    generate_btn.click(
-        fn=bot_streaming,
-        inputs=[prompt_dropdown],
-        outputs=[output_text]
-    )
-    
-    clear_btn.click(
-        fn=clear_context,
-        outputs=[upload_status]
-    )
+    file_upload.change(fn=process_uploaded_file, inputs=[file_upload], outputs=[upload_status])
+    generate_btn.click(fn=bot_streaming, inputs=[prompt_dropdown], outputs=[output_text])
+    clear_btn.click(fn=clear_context, outputs=[upload_status])
 
-# Launch the interface
 demo.launch(debug=True)