api_chat.py

"""
DocuLens - 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, Depends
from fastapi.responses import StreamingResponse
from PIL import Image
from huggingface_hub import InferenceClient

from hf_client import create_inference_client, robust_chat_completion
from middleware import check_rate_limit, validate_file_upload, sanitize_filename

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 DocuLens 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"),
    api_key: str = Depends(check_rate_limit),
):
    """
    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:
        validate_file_upload(file.filename, file.size or 0, file.content_type)
        try:
            file_bytes = await file.read()
            safe_name = sanitize_filename(file.filename)
            file_ext = os.path.splitext(safe_name)[1].lower().lstrip(".")

            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 = create_inference_client(api_key=HF_TOKEN)
            response = robust_chat_completion(
                client,
                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 = create_inference_client(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}")

    error_detail = "All models failed: " + "; ".join(errors)
    yield f"data: {json.dumps({'type': 'error', 'detail': error_detail})}\n\n"