app.py

import openai
import base64
import io
import time
import logging
import fitz  # PyMuPDF
from PIL import Image
import gradio as gr

# 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
# -------------------------------
class DocumentState:
    def __init__(self):
        self.current_doc_images = []
        self.current_doc_text = ""
        self.doc_type = None

    def clear(self):
        self.current_doc_images = []
        self.current_doc_text = ""
        self.doc_type = None

doc_state = DocumentState()

def process_pdf_file(file_path):
    """Convert PDF to images and extract text using PyMuPDF."""
    try:
        doc = fitz.open(file_path)
        images = []
        text = ""
        for page_num in range(doc.page_count):
            try:
                page = doc[page_num]
                page_text = page.get_text("text")
                if page_text.strip():
                    text += f"Page {page_num + 1}:\n{page_text}\n\n"
                
                # Render page to an image
                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")
                
                # Resize if image is too large
                max_size = 1600
                if max(img.size) > max_size:
                    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)}")
                continue
        doc.close()
        if not images:
            raise ValueError("No valid images could be extracted from the PDF")
        return images, text
    except Exception as e:
        logger.error(f"Error processing PDF file: {str(e)}")
        raise

def process_uploaded_file(file):
    """Process uploaded file 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
        if isinstance(file, dict):
            file_path = file["name"]
        else:
            file_path = file.name
        file_ext = file_path.lower().split('.')[-1]
        image_extensions = {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp'}
        
        if file_ext == 'pdf':
            doc_state.doc_type = 'pdf'
            try:
                doc_state.current_doc_images, doc_state.current_doc_text = process_pdf_file(file_path)
                return f"PDF processed successfully. Total pages: {len(doc_state.current_doc_images)}. You can now ask questions about the content."
            except Exception as e:
                return f"Error processing PDF: {str(e)}. Please try a different PDF file."
        elif file_ext in image_extensions:
            doc_state.doc_type = 'image'
            try:
                img = Image.open(file_path).convert("RGB")
                max_size = 1600
                if max(img.size) > max_size:
                    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)
                doc_state.current_doc_images = [img]
                return "Image loaded successfully. You can now ask questions about the content."
            except Exception as e:
                return f"Error processing image: {str(e)}. Please try a different image file."
        else:
            return f"Unsupported file type: {file_ext}. Please upload a PDF or image file (PNG, JPG, JPEG, GIF, BMP, WEBP)."
    except Exception as e:
        logger.error(f"Error in process_uploaded_file: {str(e)}")
        return "An error occurred while processing the file. Please try again."

# -------------------------------
# Bot Streaming Function Using OpenAI API
# -------------------------------
def bot_streaming(prompt_option, max_new_tokens=4096):
    """
    Generate a response using the OpenAI API.
    
    If an image is available, it is encoded in base64 and appended to the prompt.
    """
    try:
        # Define predetermined prompts
        prompts = {
            "NOC Timesheet": (
                """Extract structured information from the provided timesheet. The extracted details should include:

Name

Position Title

Work Location

Contractor 

NOC ID

Month and Year

Regular Service Days (ONSHORE)

Standby Days (ONSHORE in Doha)

Offshore Days

Standby & Extended Hitch Days (OFFSHORE)

Extended Hitch Days (ONSHORE Rotational)

Service during Weekends & Public Holidays

ONSHORE Overtime Hours (Over 8 hours)

OFFSHORE Overtime Hours (Over 12 hours)

Per Diem Days (ONSHORE/OFFSHORE Rotational Personnel)

Training Days

Travel Days

Noc representative appoval's name as approved_by

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."""
            ),
            "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"
                "   - Work location\n"
                "   - Contractor's name\n"
                "   - NOC ID\n"
                "   - Month and year (in MM/YYYY format)"
            ),
            "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."""
            ),
            "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
        }
    }
}
```"""
            ),
            "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
{
  "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"
  }
}
```"""
            )
        }
        
        # Retrieve the selected 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}"
        full_prompt = selected_prompt + context
        
        # 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 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}]"
        
        # Call the OpenAI API with streaming enabled.
        response = openai.ChatCompletion.create(
            model="qwen/qwen2.5-vl-72b-instruct:free",
            messages=messages,
            max_tokens=max_new_tokens,
            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
                
    except Exception as e:
        logger.error(f"Error in bot_streaming: {str(e)}")
        yield "An error occurred while processing your request. Please try again."

def clear_context():
    """Clear the current document context."""
    doc_state.clear()
    return "Document context cleared. You can upload a new document."

# -------------------------------
# Create the Gradio Interface
# -------------------------------
with gr.Blocks() as demo:
    gr.Markdown("# Document Analyzer with Predetermined Prompts")
    gr.Markdown("Upload a PDF or image (PNG, JPG, JPEG, GIF, BMP, WEBP) and select a prompt to analyze its contents.")
    
    with gr.Row():
        file_upload = gr.File(
            label="Upload Document",
            file_types=[".pdf", ".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp"]
        )
        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"
            ],
            value="NOC Timesheet"
        )
        generate_btn = gr.Button("Generate")
    
    clear_btn = gr.Button("Clear Document Context")
    
    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]
    )

# Launch the interface
demo.launch(debug=True)