Aniruddha7/QueryLens-Text2SQL_DocVQA-V2

1"""2Vision MCP tool - visual Q&A via Groq Vision API (llama-3.2-11b-vision-preview).3 4Cloud path  : image -> Groq Llama 3.2 Vision -> structured answer (FREE, 14400 req/day)5Local path  : image -> Ollama qwen2-vl -> answer (dev only)6 7Registered in mcp_server.py as 'granite_vision.qa'.8"""9import os10import base6411import requests12from typing import Optional13 14 15def qa(image_path: str, question: str, model_path: Optional[str] = None) -> dict:16    """Answer a question about a document image using Groq Vision or local Ollama."""17 18    print(f"[vision.qa] image={image_path} question={question[:80]}")19 20    # CLOUD MODE: Groq Vision API21    if os.environ.get("USE_HF_CLOUD", "0") in ("1", "true", "True"):22        groq_key = os.environ.get("GROQ_API_KEY")23        groq_model = os.environ.get("GROQ_VISION_MODEL", "meta-llama/llama-4-scout-17b-16e-instruct")24 25        if groq_key:26            print(f"[granite_vision.qa] CLOUD MODE: {groq_model} via Groq")27            try:28                from groq import Groq29 30                client = Groq(api_key=groq_key)31 32                with open(image_path, "rb") as f:33                    img_b64 = base64.b64encode(f.read()).decode("utf-8")34 35                ext = os.path.splitext(image_path)[-1].lower().lstrip(".")36                mime = "image/jpeg" if ext in ("jpg", "jpeg") else f"image/{ext}"37 38                response = client.chat.completions.create(39                    model=groq_model,40                    messages=[41                        {42                            "role": "user",43                            "content": [44                                {45                                    "type": "text",46                                    "text": (47                                        "You are analyzing a document image. "48                                        "Look carefully at the layout, tables, numbers, and text. "49                                        "Answer the following question based only on what you see:\n\n"50                                        + question51                                    )52                                },53                                {54                                    "type": "image_url",55                                    "image_url": {56                                        "url": f"data:{mime};base64,{img_b64}"57                                    }58                                }59                            ]60                        }61                    ],62                    max_tokens=51263                )64 65                answer = response.choices[0].message.content.strip()66                print(f"[granite_vision.qa] Groq answer ({len(answer)} chars): {answer[:200]}")67                return {"answer": answer, "model": f"{groq_model} (Groq Vision)"}68 69            except Exception as e:70                print(f"[granite_vision.qa] Groq failed: {e}")71                return {72                    "answer": f"Groq Vision error: {e}",73                    "model": groq_model,74                    "error": str(e)75                }76        else:77            print("[granite_vision.qa] GROQ_API_KEY not set in HF secrets.")78            return {79                "answer": "GROQ_API_KEY is not configured. Please add it in HuggingFace Space secrets.",80                "model": "none",81                "error": "missing_groq_key"82            }83 84    # LOCAL MODE (dev only): Ollama85    local_model = os.environ.get("GRANITE_MODEL_PATH", "qwen2-vl:2b")86    print(f"[granite_vision.qa] LOCAL MODE: sending image to Ollama ({local_model})")87 88    try:89        with open(image_path, "rb") as f:90            img_b64 = base64.b64encode(f.read()).decode("utf-8")91 92        base_url = os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434").rstrip("/")93        url = f"{base_url}/api/chat"94        payload = {95            "model": local_model,96            "messages": [97                {98                    "role": "user",99                    "content": f"Carefully analyze this document image. {question}",100                    "images": [img_b64]101                }102            ],103            "stream": False,104            "keep_alive": 60,105            "options": {"num_ctx": 1024, "temperature": 0.1}106        }107        r = requests.post(url, json=payload, timeout=300)108        if r.status_code == 200:109            answer = r.json().get("message", {}).get("content", "").strip()110            print(f"[granite_vision.qa] Ollama answer ({len(answer)} chars): {answer[:200]}")111            return {"answer": answer, "model": f"{local_model} (Ollama)"}112        else:113            msg = f"Ollama HTTP {r.status_code}: {r.text}"114            print(f"[granite_vision.qa] {msg}")115            return {"answer": msg, "model": f"{local_model} (Ollama)", "error": "http_error"}116 117    except Exception as e:118        msg = f"[ERROR] Ollama Q&A failed: {e}"119        print(f"[granite_vision.qa] {msg}")120        return {"answer": msg, "model": f"{local_model} (Ollama)", "error": str(e)}