AI assistant panel

api_chat.py

@@ -0,0 +1,241 @@
+"""
+DocuMint AI - Chat API endpoint.
+
+Provides /api/v1/chat for the AI Assistant panel.
+Uses the same Qwen2.5-VL models via HF Inference API for document Q&A.
+"""
+
+import os
+import io
+import base64
+import logging
+from typing import Optional
+
+from fastapi import APIRouter, UploadFile, File, Form, HTTPException
+from fastapi.responses import StreamingResponse
+from PIL import Image
+from huggingface_hub import InferenceClient
+
+logger = logging.getLogger(__name__)
+
+chat_router = APIRouter()
+
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
+
+# Model used for chat — the 72B is best for conversational Q&A
+CHAT_MODEL = "Qwen/Qwen2.5-VL-72B-Instruct"
+
+# Fallback models if primary is unavailable
+CHAT_FALLBACK_MODELS = [
+    "Qwen/Qwen2.5-VL-7B-Instruct",
+    "Qwen/Qwen2.5-VL-3B-Instruct",
+]
+
+SYSTEM_PROMPT = """You are DocuMint AI Assistant, an expert in document analysis and data extraction.
+You help users understand their documents, extract information, and answer questions about document content.
+
+When analyzing a document image:
+- Identify the document type (invoice, receipt, tax form, etc.)
+- Describe key fields and their values accurately
+- Point out any issues or anomalies you notice
+- Be precise with numbers, dates, and amounts
+
+When answering questions about extracted data:
+- Reference specific fields and values
+- Perform calculations if asked (totals, tax rates, etc.)
+- Compare values if multiple documents are discussed
+
+Keep responses concise and helpful. Use markdown formatting for clarity when appropriate."""
+
+
+def _image_to_data_url(image_bytes: bytes) -> str:
+    """Convert image bytes to a data URL for the VLM."""
+    img = Image.open(io.BytesIO(image_bytes))
+    # Convert to RGB if needed
+    if img.mode in ("RGBA", "P", "LA"):
+        img = img.convert("RGB")
+    buf = io.BytesIO()
+    img.save(buf, format="JPEG", quality=85)
+    b64 = base64.b64encode(buf.getvalue()).decode()
+    return f"data:image/jpeg;base64,{b64}"
+
+
+def _build_messages(
+    conversation_json: str,
+    image_data_url: Optional[str] = None,
+    extraction_context: Optional[str] = None,
+) -> list[dict]:
+    """Build the messages array for the VLM from conversation history."""
+    import json
+
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+
+    try:
+        conversation = json.loads(conversation_json)
+    except (json.JSONDecodeError, TypeError):
+        conversation = []
+
+    for msg in conversation:
+        role = msg.get("role", "user")
+        text = msg.get("content", "")
+        msg_image = msg.get("image")
+
+        # Build content blocks
+        if role == "user":
+            content_parts: list[dict] = []
+
+            # Add extraction context if this is the first user message with it
+            if extraction_context and msg == conversation[-1]:
+                content_parts.append({
+                    "type": "text",
+                    "text": f"[Current extraction context]\n{extraction_context}",
+                })
+
+            content_parts.append({"type": "text", "text": text})
+
+            # Attach image if present (either from this message or the provided image)
+            img_url = None
+            if msg == conversation[-1] and image_data_url:
+                img_url = image_data_url
+            elif msg_image:
+                img_url = msg_image
+
+            if img_url:
+                content_parts.append({
+                    "type": "image_url",
+                    "image_url": {"url": img_url},
+                })
+
+            messages.append({"role": "user", "content": content_parts})
+        else:
+            messages.append({"role": "assistant", "content": text})
+
+    return messages
+
+
+@chat_router.post("/api/v1/chat")
+async def api_chat(
+    messages: str = Form(..., description="JSON array of conversation messages"),
+    file: Optional[UploadFile] = File(None, description="Optional image/PDF to discuss"),
+    extraction_context: Optional[str] = Form(None, description="Current extraction result JSON for context"),
+    model: Optional[str] = Form(None, description="Model override"),
+    stream: bool = Form(False, description="Whether to stream the response"),
+):
+    """
+    Chat with the AI assistant about documents.
+
+    Accepts a conversation history (messages) and optionally an image/PDF.
+    Returns the assistant's response.
+    """
+    if not HF_TOKEN:
+        raise HTTPException(status_code=500, detail="HF_TOKEN not configured")
+
+    # Process uploaded file if present
+    image_data_url = None
+    if file:
+        try:
+            file_bytes = await file.read()
+            file_ext = (file.filename or "").lower().split(".")[-1]
+
+            if file_ext == "pdf":
+                # Convert first page of PDF to image
+                try:
+                    import fitz  # PyMuPDF
+                    doc = fitz.open(stream=file_bytes, filetype="pdf")
+                    page = doc[0]
+                    pix = page.get_pixmap(dpi=200)
+                    img_bytes = pix.tobytes("jpeg")
+                    image_data_url = f"data:image/jpeg;base64,{base64.b64encode(img_bytes).decode()}"
+                    doc.close()
+                except ImportError:
+                    # Fallback: use pdf2image
+                    from pdf2image import convert_from_bytes
+                    images = convert_from_bytes(file_bytes, first_page=1, last_page=1, dpi=200)
+                    if images:
+                        buf = io.BytesIO()
+                        images[0].save(buf, format="JPEG", quality=85)
+                        image_data_url = f"data:image/jpeg;base64,{base64.b64encode(buf.getvalue()).decode()}"
+            else:
+                image_data_url = _image_to_data_url(file_bytes)
+        except Exception as e:
+            logger.warning(f"Failed to process uploaded file: {e}")
+            raise HTTPException(status_code=400, detail=f"Failed to process file: {str(e)}")
+
+    # Build VLM messages
+    vlm_messages = _build_messages(messages, image_data_url, extraction_context)
+
+    # Try models in order
+    use_model = model or CHAT_MODEL
+    models_to_try = [use_model] + [m for m in CHAT_FALLBACK_MODELS if m != use_model]
+
+    if stream:
+        return StreamingResponse(
+            _stream_chat(models_to_try, vlm_messages),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "X-Accel-Buffering": "no",
+            },
+        )
+
+    # Non-streaming response
+    errors = []
+    for model_id in models_to_try:
+        try:
+            client = InferenceClient(api_key=HF_TOKEN)
+            response = client.chat_completion(
+                model=model_id,
+                messages=vlm_messages,
+                max_tokens=2048,
+                temperature=0.3,
+            )
+            content = response.choices[0].message.content
+            return {
+                "role": "assistant",
+                "content": content,
+                "model": model_id,
+            }
+        except Exception as e:
+            short = model_id.split("/")[-1]
+            errors.append(f"{short}: {type(e).__name__}")
+            logger.warning(f"Chat failed with {model_id}: {e}")
+
+    raise HTTPException(
+        status_code=502,
+        detail=f"All models failed: {'; '.join(errors)}",
+    )
+
+
+async def _stream_chat(models: list[str], vlm_messages: list[dict]):
+    """Generator for SSE streaming."""
+    import json
+
+    errors = []
+    for model_id in models:
+        try:
+            client = InferenceClient(api_key=HF_TOKEN)
+            stream = client.chat_completion(
+                model=model_id,
+                messages=vlm_messages,
+                max_tokens=2048,
+                temperature=0.3,
+                stream=True,
+            )
+            # Send model info first
+            yield f"data: {json.dumps({'type': 'meta', 'model': model_id})}\n\n"
+
+            for chunk in stream:
+                if chunk.choices and chunk.choices[0].delta.content:
+                    token = chunk.choices[0].delta.content
+                    yield f"data: {json.dumps({'type': 'token', 'content': token})}\n\n"
+
+            yield f"data: {json.dumps({'type': 'done'})}\n\n"
+            return
+
+        except Exception as e:
+            short = model_id.split("/")[-1]
+            errors.append(f"{short}: {type(e).__name__}")
+            logger.warning(f"Stream chat failed with {model_id}: {e}")
+
+    yield f"data: {json.dumps({'type': 'error', 'detail': f'All models failed: {'; '.join(errors)}'})}\n\n"

app.py

@@ -334,9 +334,11 @@ if __name__ == "__main__":
     from api import api_router
     from api_runs import runs_router
     from api_prompts import prompts_router
+    from api_chat import chat_router
     app.include_router(api_router)
     app.include_router(runs_router)
     app.include_router(prompts_router)
+    app.include_router(chat_router)
 
     print(f"✅ REST API routes mounted at {local_url}api/v1/")