Spaces:
Running
Running
| import os | |
| import base64 | |
| import json | |
| from typing import Any, Dict | |
| import httpx | |
| # Get your OpenRouter API key from env (you'll set this in Hugging Face later) | |
| OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY") | |
| OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1/chat/completions" | |
| MODEL_NAME = "qwen/qwen3-vl-235b-a22b-instruct" | |
| def _file_to_image_block(file_bytes: bytes, content_type: str) -> Dict[str, Any]: | |
| """ | |
| Encode the file as a data URL to feed into the multimodal model. | |
| For demo purposes we treat PDFs and images the same way here. | |
| """ | |
| b64 = base64.b64encode(file_bytes).decode("utf-8") | |
| return { | |
| "type": "input_image", | |
| "image_url": f"data:{content_type};base64,{b64}", | |
| } | |
| async def extract_fields_from_document( | |
| file_bytes: bytes, | |
| content_type: str, | |
| filename: str, | |
| ) -> Dict[str, Any]: | |
| """ | |
| Call OpenRouter with Qwen3-VL and return parsed JSON with fields. | |
| We instruct the model to return JSON only. | |
| """ | |
| if not OPENROUTER_API_KEY: | |
| raise RuntimeError("OPENROUTER_API_KEY environment variable is not set") | |
| image_block = _file_to_image_block(file_bytes, content_type) | |
| system_prompt = ( | |
| "You are a document extraction engine. " | |
| "You analyze invoices, receipts, contracts, reports and similar documents, " | |
| "and output structured JSON only (no explanations or comments)." | |
| ) | |
| user_prompt = ( | |
| "Extract important key-value pairs from the document and respond with JSON only.\n" | |
| "Use this shape:\n" | |
| "{\n" | |
| ' \"doc_type\": \"invoice | receipt | contract | report | other\",\n' | |
| ' \"confidence\": number between 0 and 100,\n' | |
| ' \"fields\": {\n' | |
| ' \"invoice_number\": \"...\",\n' | |
| ' \"date\": \"...\",\n' | |
| ' \"due_date\": \"...\",\n' | |
| ' \"total_amount\": \"...\",\n' | |
| ' \"currency\": \"...\",\n' | |
| ' \"vendor_name\": \"...\",\n' | |
| ' \"line_items\": [\n' | |
| ' {\"description\": \"...\", \"quantity\": \"...\", \"unit_price\": \"...\", \"line_total\": \"...\"}\n' | |
| ' ],\n' | |
| ' \"other_field\": \"...\"\n' | |
| " }\n" | |
| "}\n" | |
| "If fields are missing or not applicable, simply omit them." | |
| ) | |
| payload: Dict[str, Any] = { | |
| "model": MODEL_NAME, | |
| "messages": [ | |
| { | |
| "role": "system", | |
| "content": [{"type": "text", "text": system_prompt}], | |
| }, | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "text", "text": user_prompt}, | |
| image_block, | |
| ], | |
| }, | |
| ], | |
| "max_tokens": 2048, | |
| } | |
| headers = { | |
| "Authorization": f"Bearer {OPENROUTER_API_KEY}", | |
| "Content-Type": "application/json", | |
| # Optional attribution headers | |
| "HTTP-Referer": os.environ.get( | |
| "APP_URL", | |
| "https://huggingface.co/spaces/your-space", | |
| ), | |
| "X-Title": "Document Capture Demo", | |
| } | |
| async with httpx.AsyncClient(timeout=120) as client: | |
| resp = await client.post(OPENROUTER_BASE_URL, headers=headers, json=payload) | |
| resp.raise_for_status() | |
| data = resp.json() | |
| # OpenRouter returns choices[0].message.content | |
| content = data["choices"][0]["message"]["content"] | |
| # content may be a string or a list of content blocks | |
| if isinstance(content, list): | |
| text = "".join(part.get("text", "") for part in content if part.get("type") == "text") | |
| else: | |
| text = content | |
| # Try to parse JSON from the model output | |
| return json.loads(text) | |