"""API endpoints for document intelligence helper functions. Provides JSON endpoints to call Chart2Summary, Chart2CSV, Chart2Code, Table Extraction, and Image Q&A helpers. """ from __future__ import annotations import asyncio import base64 import io import json from collections.abc import AsyncGenerator import logging import queue import re from collections.abc import AsyncGenerator, Callable, Generator from typing import Any from fastapi import HTTPException, Query from fastapi.responses import JSONResponse, StreamingResponse from PIL import Image as PILImage logger = logging.getLogger(__name__) async def _stream_with_keepalive(sync_fn: Any, session_id: str = "") -> AsyncGenerator[str, None]: """Run *sync_fn* in a thread and yield SSE keepalive comments every 15 s. Yields SSE comment lines (``:\\n\\n``) while the function is running, then a ``data:`` event with the JSON result once it completes. On error, yields an ``event: error`` SSE event. """ logger.debug(f"Calling {sync_fn.__name__} for session {session_id}") loop = asyncio.get_event_loop() task = loop.run_in_executor(None, sync_fn) while True: try: result = await asyncio.wait_for(asyncio.shield(task), timeout=15) logger.info(f"Successfully completed {sync_fn.__name__} for session {session_id}") yield f"data: {json.dumps(result)}\n\n" return except asyncio.TimeoutError: yield ":\n\n" # SSE comment — keeps connection alive except Exception as exc: # noqa: BLE001 logger.error("Helper %s failed for session %s: %s", sync_fn.__name__, session_id, exc, exc_info=True) yield f"event: error\ndata: {json.dumps({'detail': str(exc)})}\n\n" return async def _stream_with_progress( sync_fn: Callable[[Callable[[str], None]], Any], session_id: str = "", ) -> AsyncGenerator[str, None]: """Run *sync_fn(on_progress)* in a thread, yielding SSE progress events in real-time. The sync function receives an ``on_progress(phase)`` callback that it can call at any point. Each call puts a message on a thread-safe queue which the async loop drains as ``event: progress`` SSE events. """ progress_q: queue.Queue[str] = queue.Queue() loop = asyncio.get_event_loop() def _on_progress(phase: str) -> None: progress_q.put(phase) task = loop.run_in_executor(None, sync_fn, _on_progress) while True: while not progress_q.empty(): phase = progress_q.get_nowait() yield f"event: progress\ndata: {json.dumps({'phase': phase})}\n\n" try: result = await asyncio.wait_for(asyncio.shield(task), timeout=1) while not progress_q.empty(): phase = progress_q.get_nowait() yield f"event: progress\ndata: {json.dumps({'phase': phase})}\n\n" logger.info("Upload completed for session %s", session_id) yield f"data: {json.dumps(result)}\n\n" return except asyncio.TimeoutError: yield ":\n\n" except Exception as exc: # noqa: BLE001 logger.error("Upload failed for session %s: %s", session_id, exc, exc_info=True) yield f"event: error\ndata: {json.dumps({'detail': str(exc)})}\n\n" return def create_helper_routes(app: Any, session_states: dict[str, dict[str, Any]]) -> None: """Register helper API routes on the FastAPI app. Args: app: FastAPI application instance. session_states: Global session states dictionary. """ @app.post("/api/helpers/chart2summary") async def api_chart2summary(session_id: str = Query(...)) -> StreamingResponse: # noqa: ANN001 """Extract chart summary for current session figure. Args: session_id: Session ID to get figure from. Returns: JSON with result and session_state. """ logger.info(f"API chart2summary request received for session_id: {session_id}") if session_id not in session_states: logger.error(f"Session not found: {session_id}") raise HTTPException(status_code=404, detail="Session not found") def _work() -> dict: logger.debug(f"Importing extract_summary_helper for session {session_id}") from app import extract_summary_helper logger.debug(f"Calling extract_summary_helper for session {session_id}") result, updated_state = extract_summary_helper(session_states[session_id]) session_states[session_id] = updated_state return {"result": result} return StreamingResponse(_stream_with_keepalive(_work, session_id), media_type="text/event-stream") @app.post("/api/helpers/chart2csv") async def api_chart2csv(session_id: str = Query(...)) -> StreamingResponse: # noqa: ANN001 """Extract chart data as CSV for current session figure. Args: session_id: Session ID to get figure from. Returns: JSON with result and session_state. """ logger.info(f"API chart2csv request received for session_id: {session_id}") if session_id not in session_states: logger.error(f"Session not found: {session_id}") raise HTTPException(status_code=404, detail="Session not found") def _work() -> dict: logger.debug(f"Importing extract_csv_helper for session {session_id}") from app import extract_csv_helper logger.debug(f"Calling extract_csv_helper for session {session_id}") result, updated_state = extract_csv_helper(session_states[session_id]) session_states[session_id] = updated_state return {"result": result} return StreamingResponse(_stream_with_keepalive(_work, session_id), media_type="text/event-stream") @app.post("/api/helpers/chart2code") async def api_chart2code(session_id: str = Query(...)) -> StreamingResponse: # noqa: ANN001 """Extract chart generation code for current session figure. Args: session_id: Session ID to get figure from. Returns: JSON with result and session_state. """ logger.info(f"API chart2code request received for session_id: {session_id}") if session_id not in session_states: logger.error(f"Session not found: {session_id}") raise HTTPException(status_code=404, detail="Session not found") def _work() -> dict: logger.debug(f"Importing extract_code_helper for session {session_id}") from app import extract_code_helper logger.debug(f"Calling extract_code_helper for session {session_id}") result, updated_state = extract_code_helper(session_states[session_id]) session_states[session_id] = updated_state return {"result": result} return StreamingResponse(_stream_with_keepalive(_work, session_id), media_type="text/event-stream") @app.post("/api/helpers/table-extract") async def api_table_extract(session_id: str = Query(...)) -> StreamingResponse: # noqa: ANN001 """Extract table data for current session figure. Args: session_id: Session ID to get figure from. Returns: JSON with result and session_state. """ logger.info(f"API table-extract request received for session_id: {session_id}") if session_id not in session_states: logger.error(f"Session not found: {session_id}") raise HTTPException(status_code=404, detail="Session not found") def _work() -> dict: logger.debug(f"Importing extract_table_helper for session {session_id}") from app import extract_table_helper logger.debug(f"Calling extract_table_helper for session {session_id}") result, updated_state = extract_table_helper(session_states[session_id]) session_states[session_id] = updated_state return {"result": result} return StreamingResponse(_stream_with_keepalive(_work, session_id), media_type="text/event-stream") @app.post("/api/helpers/describe-image") async def api_describe_image(session_id: str = Query(...)) -> StreamingResponse: # noqa: ANN001 """Generate a detailed description of the current session figure. Args: session_id: Session ID to get figure from. Returns: JSON with result and session_state. """ logger.info(f"API describe-image request received for session_id: {session_id}") if session_id not in session_states: logger.error(f"Session not found: {session_id}") raise HTTPException(status_code=404, detail="Session not found") def _work() -> dict: logger.debug(f"Importing describe_image_helper for session {session_id}") from app import describe_image_helper logger.debug(f"Calling describe_image_helper for session {session_id}") result, updated_state = describe_image_helper(session_states[session_id]) session_states[session_id] = updated_state return {"result": result} return StreamingResponse(_stream_with_keepalive(_work, session_id), media_type="text/event-stream") # ------------------------------------------------------------------ # SSE streaming endpoints (token-by-token) # ------------------------------------------------------------------ def _sse_event(data: dict[str, str]) -> str: """Format a dict as an SSE data line.""" return f"data: {json.dumps(data)}\n\n" async def _sse_stream( session_id: str, task_label: str, token_generator: Generator[str, None, None], post_process: Any | None = None, ) -> AsyncGenerator[str, None]: """Shared async generator that wraps a blocking token generator as SSE events. Args: session_id: Session to validate. task_label: Human-readable label for the status event. token_generator: Synchronous generator yielding token strings. post_process: Optional callable(full_text) -> str for cleanup. """ if session_id not in session_states: yield _sse_event({"type": "error", "message": "Session not found"}) return state = session_states[session_id] if state.get("selected_figure") is None: yield _sse_event({"type": "error", "message": "No figure selected"}) return yield _sse_event({"type": "status", "message": task_label}) full_text = "" try: for token in token_generator: full_text += token yield _sse_event({"type": "chunk", "content": token}) await asyncio.sleep(0) if post_process is not None: processed = post_process(full_text) if processed != full_text: yield _sse_event({"type": "replace", "content": processed}) yield _sse_event({"type": "done"}) except Exception as e: # noqa: BLE001 logger.error(f"Streaming error for session {session_id}: {e}", exc_info=True) yield _sse_event({"type": "error", "message": str(e)}) _STREAM_HEADERS = { "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", } def _get_selected_image(session_id: str) -> PILImage.Image: """Return the selected figure image or raise a 404.""" if session_id not in session_states: raise HTTPException(status_code=404, detail="Session not found") fig = session_states[session_id].get("selected_figure") if not fig or fig.get("image") is None: raise HTTPException(status_code=404, detail="No figure selected") return fig["image"] @app.post("/api/helpers/chart2summary/stream") async def api_chart2summary_stream(session_id: str = Query(...)) -> StreamingResponse: from infer_vision_qa import answer_question_stream image = _get_selected_image(session_id) gen = answer_question_stream(image, "", [], None) return StreamingResponse( _sse_stream(session_id, "Generating summary...", gen), media_type="text/event-stream", headers=_STREAM_HEADERS, ) @app.post("/api/helpers/chart2csv/stream") async def api_chart2csv_stream(session_id: str = Query(...)) -> StreamingResponse: from infer_chart2csv import extract_csv_stream image = _get_selected_image(session_id) gen = extract_csv_stream(image) return StreamingResponse( _sse_stream(session_id, "Extracting CSV...", gen), media_type="text/event-stream", headers=_STREAM_HEADERS, ) @app.post("/api/helpers/chart2code/stream") async def api_chart2code_stream(session_id: str = Query(...)) -> StreamingResponse: from app import PROMPT_TEXT_CODE from infer_vision_qa import answer_question_stream image = _get_selected_image(session_id) gen = answer_question_stream(image, PROMPT_TEXT_CODE, [], None) return StreamingResponse( _sse_stream(session_id, "Generating code...", gen), media_type="text/event-stream", headers=_STREAM_HEADERS, ) @app.post("/api/helpers/table-extract/stream") async def api_table_extract_stream(session_id: str = Query(...)) -> StreamingResponse: from infer_vision_qa import answer_question_stream def _clean_table_html(text: str) -> str: text = re.sub(r"^```(?:html)?\s*", "", text.strip()) text = re.sub(r"\s*```$", "", text.strip()) text = re.sub(r"^\[\s*", "", text.strip()) text = re.sub(r"\s*\]$", "", text.strip()) return text image = _get_selected_image(session_id) gen = answer_question_stream(image, "", [], None) return StreamingResponse( _sse_stream(session_id, "Extracting table...", gen, post_process=_clean_table_html), media_type="text/event-stream", headers=_STREAM_HEADERS, ) @app.post("/api/helpers/describe-image/stream") async def api_describe_image_stream(session_id: str = Query(...)) -> StreamingResponse: from infer_vision_qa import answer_question_stream image = _get_selected_image(session_id) gen = answer_question_stream(image, "Describe this image in detail", [], None) return StreamingResponse( _sse_stream(session_id, "Describing image...", gen), media_type="text/event-stream", headers=_STREAM_HEADERS, )