app.py

from fastapi import FastAPI, UploadFile, File, Form
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
from pdf2image import convert_from_bytes
from typing import List
import base64
from io import BytesIO
import requests
import os
import json

app = FastAPI()

# Optional: Allow CORS for testing
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

FIREWORKS_API_KEY = os.getenv("FIREWORKS_API_KEY")
FIREWORKS_MODEL = "accounts/fireworks/models/llama4-scout-instruct-basic"
FIREWORKS_ENDPOINT = "https://api.fireworks.ai/inference/v1/chat/completions"

def encode_image(img: Image.Image) -> str:
    buffered = BytesIO()
    img.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")

@app.post("/generate/")
async def generate(image: UploadFile = File(...)):
    image_data = await image.read()

    # Handle image or PDF conversion
    if image.filename.lower().endswith(".pdf"):
        pages = convert_from_bytes(image_data, dpi=150)
        if not pages:
            return {"error": "No pages found in PDF"}
        pil_img = pages[0].convert("RGB")  # Use first page of PDF
    else:
        pil_img = Image.open(BytesIO(image_data)).convert("RGB")

    encoded_img = encode_image(pil_img)
    prompt = f"""
    You are an intelligent invoice data extractor. Given image of invoice (in English or other languages), extract key business fields into the specified JSON format. Return each field along with an estimated accuracy score between 0 and 1.
 
- Accuracy reflects your confidence in the correctness of each field.
- Handle synonyms (e.g., 'total' = 'net', 'tax' = 'GST'/'TDS').
- Detect currency from symbols ($, ₹, €) or keywords (USD, INR, EUR); default to USD if unclear.
- The 'items' list may have multiple entries, each with detailed attributes.
- If a field is missing or not found, return an empty value (`""` or `0`) and set `accuracy` to `0.0`. 
- I just want json response in below format nothing else in response
{{
  "invoice": {{
    "invoice_number": {{"value": "", "accuracy": 0.0}},
    "invoice_date": {{"value": "YYYY-MM-DD", "accuracy": 0.0}},
    "due_date": {{"value": "YYYY-MM-DD", "accuracy": 0.0}},
    "purchase_order_number": {{"value": "", "accuracy": 0.0}},
    "vendor": {{
      "vendor_id": {{"value": "", "accuracy": 0.0}},
      "name": {{"value": "", "accuracy": 0.0}},
      "address": {{
        "line1": {{"value": "", "accuracy": 0.0}},
        "line2": {{"value": "", "accuracy": 0.0}},
        "city": {{"value": "", "accuracy": 0.0}},
        "state": {{"value": "", "accuracy": 0.0}},
        "postal_code": {{"value": "", "accuracy": 0.0}},
        "country": {{"value": "", "accuracy": 0.0}}
      }},
      "contact": {{
        "email": {{"value": "", "accuracy": 0.0}},
        "phone": {{"value": "", "accuracy": 0.0}}
      }},
      "tax_id": {{"value": "", "accuracy": 0.0}}
    }},
    "buyer": {{
      "buyer_id": {{"value": "", "accuracy": 0.0}},
      "name": {{"value": "", "accuracy": 0.0}},
      "address": {{
        "line1": {{"value": "", "accuracy": 0.0}},
        "line2": {{"value": "", "accuracy": 0.0}},
        "city": {{"value": "", "accuracy": 0.0}},
        "state": {{"value": "", "accuracy": 0.0}},
        "postal_code": {{"value": "", "accuracy": 0.0}},
        "country": {{"value": "", "accuracy": 0.0}}
      }},
      "contact": {{
        "email": {{"value": "", "accuracy": 0.0}},
        "phone": {{"value": "", "accuracy": 0.0}}
      }},
      "tax_id": {{"value": "", "accuracy": 0.0}}
    }},
    "items": [
      {{
        "item_id": {{"value": "", "accuracy": 0.0}},
        "description": {{"value": "", "accuracy": 0.0}},
        "quantity": {{"value": 0, "accuracy": 0.0}},
        "unit_of_measure": {{"value": "", "accuracy": 0.0}},
        "unit_price": {{"value": 0, "accuracy": 0.0}},
        "total_price": {{"value": 0, "accuracy": 0.0}},
        "tax_rate": {{"value": 0, "accuracy": 0.0}},
        "tax_amount": {{"value": 0, "accuracy": 0.0}},
        "discount": {{"value": 0, "accuracy": 0.0}},
        "net_amount": {{"value": 0, "accuracy": 0.0}}
      }}
    ],
    "sub_total": {{"value": 0, "accuracy": 0.0}},
    "tax_total": {{"value": 0, "accuracy": 0.0}},
    "discount_total": {{"value": 0, "accuracy": 0.0}},
    "total_amount": {{"value": 0, "accuracy": 0.0}},
    "currency": {{"value": "", "accuracy": 0.0}}
  }}
}}
"""

    payload = {
        "model": FIREWORKS_MODEL,
        "messages": [
            {
                "role": "user",
                "content": [
                    { "type": "text", "text": prompt },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/png;base64,{encoded_img}"
                        }
                    }
                ]
            }
        ],
        "max_tokens": 1024,
        "temperature": 0.2,
        "top_p": 0.9
    }

    headers = {
        "Authorization": f"Bearer {FIREWORKS_API_KEY}",
        "Content-Type": "application/json"
    }

    response = requests.post(FIREWORKS_ENDPOINT, headers=headers, json=payload)
    try:
        llm_output = response.json()["choices"][0]["message"]["content"]
        json_start = llm_output.find("{")
        json_end = llm_output.rfind("}") + 1
        json_str = llm_output[json_start:json_end]
        structured_data = json.loads(json_str)
        return structured_data
    except Exception as e:
        output = f"Error: {response.text}"
        return {"error": f"{output}"}