"""Hugging Face Spaces and Local Entry Point for DocNeedle.
Provides a standalone, frontend-identical Gradio UI dashboard alongside the complete FastAPI REST API.
Features dynamic field-by-field schema creation matching frontend/src/App.tsx, dual formatted and raw JSON output tabs,
ephemeral OCR deallocation, and model bootstrapping.
Zero dependency on external frontend frameworks.
"""
import html
import json
import logging
import math
import os
import sys
import tempfile
from pathlib import Path
from typing import Any
# ZeroGPU decorator compatibility (applied before CUDA imports)
try:
import spaces
gpu_decorator = spaces.GPU
except Exception:
def gpu_decorator(*args, **kwargs):
if args and callable(args[0]):
return args[0]
return lambda f: f
# Add src and root to sys.path
root_dir = Path(__file__).resolve().parent
src_dir = root_dir / "src"
if src_dir.exists():
sys.path.insert(0, str(src_dir))
sys.path.insert(0, str(root_dir))
# Preload NVIDIA CUDA runtime and cuBLAS shared libraries into global symbol space
# (libcudart.so.12, libcublas.so.12) to ensure llama-cpp-python can link on ZeroGPU
try:
from docneedle.cuda_compat import preload_cuda_runtime_libs
preload_cuda_runtime_libs()
except Exception as e:
logging.getLogger("DocNeedleSpace").debug("CUDA runtime preload info: %s", e)
import gradio as gr # noqa: E402
from docneedle.config import settings # noqa: E402
logger = logging.getLogger("DocNeedleSpace")
logging.basicConfig(level=logging.INFO)
# 1. Automatic model bootstrapping for Hugging Face Spaces and local startup
try:
from scripts.download_models import ensure_all_models_present
logger.info("Checking and bootstrapping quantized models...")
ensure_all_models_present(settings.model_storage_dir)
except Exception as e:
logger.warning(f"Automatic model download check failed or bypassed: {e}")
CUSTOM_CSS = """
/* Dark mode base styles */
:root, body, .gradio-container {
max-width: 1360px !important;
font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif !important;
background-color: #09090b !important;
color: #f4f4f5 !important;
}
/* Headings and labels */
.gradio-container .prose *,
.gradio-container h1,
.gradio-container h2,
.gradio-container h3,
.gradio-container h4,
.gradio-container p {
color: #f4f4f5 !important;
}
/* Container blocks */
.gradio-container .block,
.gradio-container .form {
background-color: #18181b !important;
border: 1px solid #27272a !important;
border-radius: 8px !important;
}
/* Field Card matching frontend */
.dn-field-card {
background-color: #141416 !important;
border: 1px solid #27272a !important;
border-radius: 8px !important;
padding: 12px 14px !important;
margin-bottom: 10px !important;
transition: border-color 0.15s ease-in-out !important;
}
.dn-field-card:hover {
border-color: #3f3f46 !important;
}
/* Add Field dashed button matching frontend */
.dn-add-field-btn {
border: 1px dashed #3f3f46 !important;
background-color: transparent !important;
color: #a1a1aa !important;
font-size: 13px !important;
font-weight: 500 !important;
border-radius: 6px !important;
width: 100% !important;
padding: 10px 14px !important;
cursor: pointer !important;
transition: all 0.2s ease-in-out !important;
}
.dn-add-field-btn:hover {
border-color: #71717a !important;
color: #f4f4f5 !important;
background-color: #27272a !important;
}
/* Delete button */
.dn-del-btn {
background: transparent !important;
border: none !important;
color: #71717a !important;
font-size: 14px !important;
cursor: pointer !important;
padding: 2px 6px !important;
transition: color 0.15s ease-in-out !important;
}
.dn-del-btn:hover {
color: #f43f5e !important;
}
/* Header action container */
.dn-header-actions {
display: flex !important;
gap: 6px !important;
align-items: center !important;
justify-content: flex-end !important;
}
/* Action button for import, export, clear */
.dn-action-btn {
background-color: #27272a !important;
border: 1px solid #3f3f46 !important;
color: #e4e4e7 !important;
font-size: 11px !important;
font-weight: 500 !important;
border-radius: 6px !important;
padding: 4px 10px !important;
min-width: 0 !important;
cursor: pointer !important;
transition: all 0.15s ease-in-out !important;
}
.dn-action-btn:hover {
background-color: #3f3f46 !important;
color: #ffffff !important;
border-color: #52525b !important;
}
/* Clear all button */
.dn-clear-btn {
background: transparent !important;
border: 1px solid transparent !important;
color: #71717a !important;
font-size: 11px !important;
cursor: pointer !important;
}
.dn-clear-btn:hover {
color: #f43f5e !important;
background-color: rgba(244, 63, 94, 0.1) !important;
border-color: rgba(244, 63, 94, 0.2) !important;
}
/* Secondary toolbar button */
.dn-secondary-btn {
background-color: #18181b !important;
border: 1px solid #3f3f46 !important;
color: #cbd5e1 !important;
font-size: 12px !important;
font-weight: 500 !important;
border-radius: 6px !important;
padding: 6px 12px !important;
cursor: pointer !important;
transition: all 0.15s ease-in-out !important;
}
.dn-secondary-btn:hover {
background-color: #27272a !important;
border-color: #71717a !important;
color: #f4f4f5 !important;
}
.dn-raw-hint {
font-size: 11px !important;
color: #a1a1aa !important;
margin-bottom: 6px !important;
}
.dn-input-tabs {
margin-top: 4px !important;
}
/* Primary Action Button */
.dn-primary-btn {
background: linear-gradient(135deg, #f97316 0%, #ea580c 100%) !important;
color: #ffffff !important;
font-weight: 600 !important;
border-radius: 8px !important;
padding: 12px 18px !important;
border: none !important;
font-size: 14px !important;
cursor: pointer !important;
transition: all 0.2s ease-in-out !important;
}
.dn-primary-btn:hover {
background: linear-gradient(135deg, #ea580c 0%, #c2410c 100%) !important;
box-shadow: 0 4px 14px rgba(234, 88, 12, 0.35) !important;
}
/* File upload container & preview */
.gradio-container .file-preview-holder,
.gradio-container .upload-container,
.gradio-container [data-testid="dropzone"] {
background-color: #141416 !important;
border: 1px dashed #3f3f46 !important;
color: #f4f4f5 !important;
}
/* Output Tabs */
.gradio-container .tab-nav {
border-bottom: 1px solid #27272a !important;
}
.gradio-container .tab-nav button {
color: #a1a1aa !important;
font-size: 13px !important;
font-weight: 500 !important;
border-bottom: 2px solid transparent !important;
}
.gradio-container .tab-nav button.selected {
color: #f4f4f5 !important;
border-bottom-color: #f97316 !important;
font-weight: 600 !important;
}
/* JSON viewer */
.gradio-container .json-holder {
background-color: #09090b !important;
border: 1px solid #27272a !important;
border-radius: 6px !important;
}
/* Hide Gradio footer completely */
footer, .gradio-container footer, .gradio-container .footer {
display: none !important;
visibility: hidden !important;
height: 0 !important;
padding: 0 !important;
margin: 0 !important;
}
"""
HEADER_HTML = """
DocNeedle
Document Extraction Microservice
"""
FIELD_TYPE_KEYS = [
"string",
"number",
"money",
"date",
"time",
"percentage",
"person",
"organization",
"id",
"iban",
"email",
"phone",
"address",
"boolean",
"list",
]
TYPE_MAP = {
"string": "string",
"text": "string",
"metin": "string",
"number": "number",
"integer": "number",
"float": "number",
"sayฤฑ": "number",
"money": "money",
"currency": "money",
"amount": "money",
"para": "money",
"date": "date",
"tarih": "date",
"time": "time",
"saat": "time",
"email": "email",
"phone": "phone",
"id": "id",
"iban": "iban",
"percentage": "percentage",
"percent": "percentage",
"oran": "percentage",
"person": "person",
"name": "person",
"organization": "organization",
"company": "organization",
"address": "address",
"boolean": "boolean",
"list": "list",
}
MAX_FIELDS = 10
# ZeroGPU duration estimates must use only metadata available before a GPU worker is allocated.
_DOCUMENT_IMAGE_SUFFIXES = frozenset({".png", ".jpg", ".jpeg", ".tiff", ".bmp", ".webp"})
_DOCUMENT_OFFICE_SUFFIXES = frozenset({".doc", ".docx", ".xls", ".xlsx"})
_DOCUMENT_TEXT_SUFFIXES = frozenset(
{".csv", ".htm", ".html", ".json", ".md", ".txt", ".tsv", ".xml"}
)
_DOCUMENT_PROFILE_KEYS = frozenset(
{
"file_count",
"total_bytes",
"pdf_count",
"pdf_pages",
"image_count",
"office_count",
"text_count",
"other_count",
}
)
INITIAL_STARTER_FIELDS = [
("", "", "string"),
]
DEFAULT_EXAMPLE_SCHEMA: list[dict[str, str]] = [
{
"key": "example_field",
"description": "Describe what you want to extract from the document",
"type": "string",
}
]
def resolve_field_type(raw_type: Any) -> str:
"""Normalize raw field type string against TYPE_MAP and FIELD_TYPE_KEYS."""
t = str(raw_type).strip().lower() if raw_type else "string"
if t in ("none", "n/a", ""):
return "string"
resolved = TYPE_MAP.get(t, t)
return resolved if resolved in FIELD_TYPE_KEYS else "string"
def collect_form_items(
q_vals: list[Any] | tuple[Any, ...],
k_vals: list[Any] | tuple[Any, ...],
t_vals: list[Any] | tuple[Any, ...],
vis_vals: list[bool] | None = None,
) -> list[dict[str, str]]:
"""Collect normalized extraction item dicts from form builder inputs."""
items: list[dict[str, str]] = []
num_fields = min(len(q_vals), len(k_vals), len(t_vals))
for i in range(num_fields):
q = str(q_vals[i]).strip() if q_vals[i] else ""
k = str(k_vals[i]).strip() if k_vals[i] else ""
is_vis = vis_vals[i] if vis_vals is not None and i < len(vis_vals) else True
if not is_vis and not (q or k):
continue
if q or k:
items.append(
{
"key": k or q or f"item_{len(items) + 1}",
"description": q or k,
"type": resolve_field_type(t_vals[i]),
}
)
return items
def unpack_field_args(
field_args: tuple[Any, ...],
num_fields: int = MAX_FIELDS,
) -> tuple[tuple[Any, ...], tuple[Any, ...], tuple[Any, ...]]:
"""Unpack flattened Gradio form arguments into (questions, keys, types) slices."""
q_vals = field_args[0:num_fields]
k_vals = field_args[num_fields : 2 * num_fields]
t_vals = field_args[2 * num_fields : 3 * num_fields]
return q_vals, k_vals, t_vals
def parse_raw_json_schema(raw_json: str) -> list[dict[str, str]]:
"""Parse raw JSON string into a validated list of extraction item dicts."""
if not raw_json or not raw_json.strip():
return []
try:
data = json.loads(raw_json)
except Exception as e:
raise ValueError(f"Invalid JSON syntax: {e}") from e
if isinstance(data, dict):
data = [data]
elif not isinstance(data, list):
raise ValueError("Schema must be a JSON array of objects or a single JSON object.")
parsed_items: list[dict[str, str]] = []
for idx, item in enumerate(data):
if not isinstance(item, dict):
continue
key = str(
item.get("key") or item.get("name") or item.get("id") or item.get("field") or ""
).strip()
desc = str(
item.get("description")
or item.get("question")
or item.get("prompt")
or item.get("desc")
or ""
).strip()
resolved_t = resolve_field_type(item.get("type", "string"))
if key or desc:
final_key = key or desc or f"item_{idx + 1}"
final_desc = desc or key
parsed_items.append(
{
"key": final_key,
"description": final_desc,
"type": resolved_t,
}
)
return parsed_items
def _empty_document_profile() -> dict[str, int]:
"""Return an aggregate-only profile suitable for session state."""
return {key: 0 for key in _DOCUMENT_PROFILE_KEYS}
def _as_file_list(files: Any) -> list[Any]:
"""Normalize Gradio's single/multiple file values without reading file content."""
if not files:
return []
return files if isinstance(files, list) else [files]
def _uploaded_file_path(file_value: Any) -> Path:
"""Resolve a Gradio upload value to its temporary local path."""
return Path(file_value.name if hasattr(file_value, "name") else str(file_value))
def _count_pdf_pages(path: Path) -> int:
"""Read PDF page metadata without extracting document text or running OCR."""
import pypdfium2
try:
pdf = pypdfium2.PdfDocument(path)
except Exception as exc:
raise ValueError(f"Unable to inspect PDF page count for '{path.name}'.") from exc
try:
return len(pdf)
finally:
pdf.close()
def _build_document_profile(files: Any) -> dict[str, int]:
"""Build a privacy-preserving, CPU-only profile for ZeroGPU scheduling."""
profile = _empty_document_profile()
for file_value in _as_file_list(files):
path = _uploaded_file_path(file_value)
if not path.is_file():
raise FileNotFoundError(
f"Uploaded file is no longer available for ZeroGPU profiling: '{path.name}'."
)
suffix = path.suffix.lower()
profile["file_count"] += 1
profile["total_bytes"] += path.stat().st_size
if suffix == ".pdf":
profile["pdf_count"] += 1
profile["pdf_pages"] += _count_pdf_pages(path)
elif suffix in _DOCUMENT_IMAGE_SUFFIXES:
profile["image_count"] += 1
elif suffix in _DOCUMENT_OFFICE_SUFFIXES:
profile["office_count"] += 1
elif suffix in _DOCUMENT_TEXT_SUFFIXES:
profile["text_count"] += 1
else:
profile["other_count"] += 1
return profile
def _is_complete_document_profile(value: Any) -> bool:
"""Check that a Gradio state value is a usable aggregate-only profile."""
return isinstance(value, dict) and all(
isinstance(value.get(key), int) and value[key] >= 0 for key in _DOCUMENT_PROFILE_KEYS
)
def _count_requested_extraction_items(
active_tab: str,
raw_json: str,
vis: list[bool] | None,
field_args: tuple[Any, ...],
) -> int:
"""Mirror UI schema selection using only event inputs available before allocation."""
q_vals, k_vals, t_vals = unpack_field_args(field_args)
form_items = collect_form_items(q_vals, k_vals, t_vals, vis)
raw_items: list[dict[str, str]] = []
if raw_json and raw_json.strip():
try:
raw_items = parse_raw_json_schema(raw_json)
except ValueError:
# The decorated handler returns a validation error immediately for invalid raw JSON.
return 0 if active_tab == "raw" else len(form_items)
default_ex = DEFAULT_EXAMPLE_SCHEMA[0]
raw_items = [
item
for item in raw_items
if not (
item["key"] == default_ex["key"]
and item["description"] == default_ex["description"]
)
]
return len(raw_items) if active_tab == "raw" and raw_items else len(form_items or raw_items)
def serialize_form_to_json(
q_vals: list[Any] | tuple[Any, ...],
k_vals: list[Any] | tuple[Any, ...],
t_vals: list[Any] | tuple[Any, ...],
vis_vals: list[bool] | None = None,
) -> str:
"""Convert form fields to formatted JSON string."""
items = collect_form_to_items_or_default(q_vals, k_vals, t_vals, vis_vals)
return json.dumps(items, indent=2, ensure_ascii=False)
def collect_form_to_items_or_default(
q_vals: list[Any] | tuple[Any, ...],
k_vals: list[Any] | tuple[Any, ...],
t_vals: list[Any] | tuple[Any, ...],
vis_vals: list[bool] | None = None,
) -> list[dict[str, str]]:
"""Collect form items or return the default example schema when empty."""
items = collect_form_items(q_vals, k_vals, t_vals, vis_vals)
return items if items else list(DEFAULT_EXAMPLE_SCHEMA)
def apply_items_to_form_updates(
items: list[dict[str, str]],
) -> tuple[list[bool], list[Any], list[Any], list[Any], list[Any], str]:
"""Given parsed items, generate updates for vis_state, field_groups, q_inputs, k_inputs, t_inputs, and raw_json."""
count = min(len(items), MAX_FIELDS)
new_vis = [(i < count) for i in range(MAX_FIELDS)]
if count == 0:
new_vis[0] = True
group_updates = [gr.update(visible=new_vis[i]) for i in range(MAX_FIELDS)]
q_updates: list[Any] = []
k_updates: list[Any] = []
t_updates: list[Any] = []
for i in range(MAX_FIELDS):
if i < count:
it = items[i]
q_updates.append(gr.update(value=it.get("description", "")))
k_updates.append(gr.update(value=it.get("key", "")))
t_updates.append(gr.update(value=it.get("type", "string")))
else:
q_updates.append(gr.update(value=""))
k_updates.append(gr.update(value=""))
t_updates.append(gr.update(value="string"))
json_str = json.dumps(
items if items else DEFAULT_EXAMPLE_SCHEMA,
indent=2,
ensure_ascii=False,
)
return new_vis, group_updates, q_updates, k_updates, t_updates, json_str
def render_results_html(response_dict: dict[str, Any] | None) -> str:
"""Render extracted results matching frontend/src/App.tsx card design."""
if not response_dict:
return """
No analysis performed yet
Upload documents on the left and define target fields to start extraction.
"""
if "error" in response_dict:
err_msg = html.escape(str(response_dict["error"]))
return f"""
Extraction Error: {err_msg}
"""
total = response_dict.get("total_items", 0)
found = response_dict.get("found_items", 0)
time_ms = response_dict.get("processing_time_ms", 0)
time_s = f"{time_ms / 1000:.2f} s"
results = response_dict.get("results", [])
cards_html = ""
for r in results:
is_found = r.get("found", False)
key = html.escape(str(r.get("key", "")))
raw_val = r.get("value")
val = (
html.escape(str(raw_val))
if raw_val is not None and str(raw_val) != ""
else "Not found in document"
)
conf = int(r.get("confidence", 0) * 100)
evidence = r.get("evidence")
thinking = r.get("thinking")
badge_html = (
'Found'
if is_found
else 'Not Found'
)
conf_html = (
f'%{conf} Confidence'
if is_found
else ""
)
evidence_html = ""
if evidence:
doc_name = html.escape(str(evidence.get("document", "")))
page_num = evidence.get("page")
raw_snippet = evidence.get("snippet", "")
snippet = html.escape(str(raw_snippet)) if raw_snippet else ""
page_str = f" โข Page {int(page_num)}" if isinstance(page_num, int) else ""
snippet_box = (
f'"{snippet}"
'
if snippet
else ""
)
evidence_html = f"""
Source Document: {doc_name}{page_str}
{snippet_box}
"""
thinking_html = ""
if thinking:
safe_thinking = html.escape(str(thinking))
thinking_html = f"""
๐ง Model Reasoning (Thinking Process)
{safe_thinking}
"""
cards_html += f"""
{key}
{badge_html}
{conf_html}
{val}
{evidence_html}
{thinking_html}
"""
summary_html = f"""
Extracted Fields
{found} / {total}
Processing Time
{time_s}
Extraction Completed
"""
return summary_html + cards_html
def make_status_html(message: str, is_success: bool = False, is_loading: bool = False) -> str:
border_color = "#10b981" if is_success else ("#f97316" if is_loading else "#27272a")
bg_color = (
"rgba(16, 185, 129, 0.08)"
if is_success
else ("rgba(249, 115, 22, 0.08)" if is_loading else "#141416")
)
text_color = "#34d399" if is_success else ("#fb923c" if is_loading else "#a1a1aa")
icon = "โ
" if is_success else ("โณ" if is_loading else "๐")
return f"""
{icon}
{message}
"""
def _format_completion_status(found: int, total: int, total_time_ms: float) -> str:
"""Format extraction completion status message with elapsed time."""
time_s = f"{total_time_ms / 1000:.2f} s"
return (
f"Extraction complete: {found} of {total} items found in {time_s} ({total_time_ms:.0f} ms)."
)
async def _async_perform_extraction_ui(
files: Any,
session_chunks: list[Any] | None,
active_tab: str,
raw_json: str,
vis: list[bool],
*field_args: Any,
start_time: float | None = None,
) -> tuple[str, dict[str, Any], list[Any], str]:
"""Execute complete document extraction pipeline from dynamic Gradio UI fields or raw JSON schema.
Uses initial OCR/ingested chunks from session state to avoid redundant OCR runs on queries.
"""
import time
if start_time is None:
start_time = time.perf_counter()
if not files and not session_chunks:
err = {"error": "Please upload at least one document (PDF, DOCX, XLSX, or Image)."}
return (
render_results_html(err),
err,
[],
make_status_html("Please upload at least one document to start.", is_success=False),
)
file_list = files if isinstance(files, list) else ([files] if files else [])
if len(file_list) == 0 and not session_chunks:
err = {"error": "Please upload at least one document."}
return (
render_results_html(err),
err,
[],
make_status_html("Please upload at least one document.", is_success=False),
)
from docneedle.extraction.schemas import ExtractionItemQuery
from docneedle.extraction.service import ExtractionService
from docneedle.ingest.service import IngestionService
from docneedle.retrieval.store import SessionVectorStore
use_raw = active_tab == "raw"
q_vals, k_vals, t_vals = unpack_field_args(field_args)
form_queries = [
ExtractionItemQuery(key=it["key"], description=it["description"], type=it["type"])
for it in collect_form_items(q_vals, k_vals, t_vals, vis)
]
raw_queries: list[ExtractionItemQuery] = []
if raw_json and raw_json.strip():
try:
parsed_items = parse_raw_json_schema(raw_json)
default_ex = DEFAULT_EXAMPLE_SCHEMA[0]
for it in parsed_items:
if (
it["key"] == default_ex["key"]
and it["description"] == default_ex["description"]
):
continue
raw_queries.append(
ExtractionItemQuery(
key=it["key"],
description=it["description"],
type=it["type"],
)
)
except Exception as e:
if use_raw:
err = {"error": f"Invalid JSON schema: {e}"}
return (
render_results_html(err),
err,
list(session_chunks) if session_chunks else [],
make_status_html(f"Schema Error: {e}", is_success=False),
)
queries = raw_queries if (use_raw and raw_queries) else (form_queries or raw_queries)
if not queries:
err = {"error": "Please define at least one field to extract via Form Builder or Raw JSON."}
return (
render_results_html(err),
err,
list(session_chunks) if session_chunks else [],
make_status_html("Please define at least one field to extract.", is_success=False),
)
try:
from docneedle.cuda_compat import preload_cuda_runtime_libs
preload_cuda_runtime_libs()
except Exception:
pass
ingest_svc = IngestionService()
extraction_svc = ExtractionService()
store = SessionVectorStore()
try:
chunks_to_use = list(session_chunks) if session_chunks else []
if not chunks_to_use:
with ingest_svc.parser:
for f in file_list:
file_path = Path(f.name if hasattr(f, "name") else str(f))
content = file_path.read_bytes()
filename = file_path.name
ingest_res = ingest_svc.ingest_bytes(content=content, filename=filename)
chunks_to_use.extend(ingest_res.chunks)
if not chunks_to_use:
err = {"error": "Could not extract readable text from the uploaded documents."}
return (
render_results_html(err),
err,
[],
make_status_html(
"Could not extract readable text from documents.", is_success=False
),
)
logger.info(
f"--> [PIPELINE] Indexing {len(chunks_to_use)} text chunks across {len(file_list)} files into in-memory vector store..."
)
store.add_chunks(chunks_to_use)
logger.info(
f"--> [PIPELINE] Starting extraction pipeline for {len(queries)} target field query(s)..."
)
response = await extraction_svc.extract_items(queries, store)
resp_dict = response.model_dump()
total_time_ms = round((time.perf_counter() - start_time) * 1000.0, 2)
resp_dict["processing_time_ms"] = total_time_ms
status_msg = _format_completion_status(
resp_dict.get("found_items", 0),
resp_dict.get("total_items", 0),
total_time_ms,
)
return (
render_results_html(resp_dict),
resp_dict,
chunks_to_use,
make_status_html(status_msg, is_success=True),
)
except Exception as e:
logger.error(f"UI extraction error: {e}", exc_info=True)
err = {"error": str(e)}
return (
render_results_html(err),
err,
list(session_chunks) if session_chunks else [],
make_status_html(f"Error: {e}", is_success=False),
)
finally:
store.clear()
import gc
gc.collect()
if "torch" in sys.modules:
try:
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
except Exception:
pass
def prepare_documents_ui(
files: Any,
session_chunks: list[Any] | None,
) -> tuple[list[Any], str]:
"""Pre-ingest and parse/OCR documents on CPU before allocating ZeroGPU worker.
Runs on host CPU outside @spaces.GPU so that long OCR rasterization or multi-page
scans do not waste ZeroGPU time quotas or trigger ZeroGPU duration timeouts.
"""
if session_chunks:
return session_chunks, gr.update()
file_list = files if isinstance(files, list) else ([files] if files else [])
if not file_list:
return [], make_status_html(
"Please upload at least one document to start.", is_success=False
)
from docneedle.ingest.service import IngestionService
ingest_svc = IngestionService()
parsed_chunks: list[Any] = []
try:
with ingest_svc.parser:
for f in file_list:
file_path = Path(f.name if hasattr(f, "name") else str(f))
content = file_path.read_bytes()
filename = file_path.name
ingest_res = ingest_svc.ingest_bytes(content=content, filename=filename)
parsed_chunks.extend(ingest_res.chunks)
if not parsed_chunks:
return [], make_status_html(
"Could not extract readable text from documents.", is_success=False
)
status_msg = (
f"โก Documents parsed ({len(parsed_chunks)} text chunk(s)). "
"Dispatching neural extraction to GPU..."
)
return parsed_chunks, make_status_html(status_msg, is_success=True)
except Exception as e:
logger.error(f"CPU document ingestion error: {e}", exc_info=True)
return [], make_status_html(f"Ingestion error: {e}", is_success=False)
def _estimate_extraction_duration(
files: Any = None,
session_chunks: list[Any] | None = None,
document_profile: dict[str, int] | None = None,
active_tab: str = "form",
raw_json: str = "",
vis: list[bool] | None = None,
*field_args: Any,
**kwargs: Any,
) -> int:
"""Estimate a ZeroGPU reservation before the decorated extraction handler runs.
The estimate deliberately uses only upload metadata and event inputs available before GPU
allocation. It never parses document content, runs OCR, or creates embeddings.
"""
requested_fields = _count_requested_extraction_items(
active_tab, raw_json, vis, tuple(field_args)
)
cached_chunk_count = len(session_chunks) if session_chunks else 0
work_seconds = (
settings.zerogpu_duration_base_seconds
+ requested_fields * settings.zerogpu_duration_per_field_seconds
)
if cached_chunk_count:
# Cached chunks skip parsing/OCR but are embedded again into the request-local store.
work_seconds += cached_chunk_count * settings.zerogpu_duration_per_cached_chunk_seconds
else:
profile = (
document_profile
if _is_complete_document_profile(document_profile)
else _build_document_profile(files)
)
input_mib = math.ceil(profile["total_bytes"] / (1024 * 1024))
work_seconds += profile["pdf_count"] * settings.zerogpu_duration_per_pdf_file_seconds
work_seconds += profile["pdf_pages"] * settings.zerogpu_duration_per_pdf_page_seconds
work_seconds += profile["image_count"] * settings.zerogpu_duration_per_image_seconds
work_seconds += (
profile["office_count"] * settings.zerogpu_duration_per_office_document_seconds
)
work_seconds += profile["text_count"] * settings.zerogpu_duration_per_text_document_seconds
work_seconds += (
profile["other_count"] * settings.zerogpu_duration_per_other_document_seconds
)
work_seconds += input_mib * settings.zerogpu_duration_per_mib_seconds
estimated_seconds = (
math.ceil(
work_seconds
* settings.zerogpu_duration_safety_factor
/ settings.zerogpu_duration_granularity_seconds
)
* settings.zerogpu_duration_granularity_seconds
)
return max(
settings.zerogpu_duration_min_seconds,
min(settings.zerogpu_duration_max_seconds, estimated_seconds),
)
@gpu_decorator(duration=_estimate_extraction_duration)
def perform_extraction_ui(
files: Any,
session_chunks: list[Any] | None,
document_profile: dict[str, int] | None,
active_tab: str,
raw_json: str,
vis: list[bool],
*field_args: Any,
) -> tuple[str, dict[str, Any], list[Any], str]:
"""Execute complete document extraction pipeline from dynamic Gradio UI fields or raw JSON schema."""
import asyncio
import concurrent.futures
import time
start_time = time.perf_counter()
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
results_html, resp_dict, chunks_to_use, status_html = pool.submit(
asyncio.run,
_async_perform_extraction_ui(
files,
session_chunks,
active_tab,
raw_json,
vis,
*field_args,
start_time=start_time,
),
).result()
total_time_ms = round((time.perf_counter() - start_time) * 1000.0, 2)
if isinstance(resp_dict, dict) and "results" in resp_dict:
resp_dict["processing_time_ms"] = total_time_ms
results_html = render_results_html(resp_dict)
status_msg = _format_completion_status(
resp_dict.get("found_items", 0),
resp_dict.get("total_items", 0),
total_time_ms,
)
status_html = make_status_html(status_msg, is_success=True)
return results_html, resp_dict, chunks_to_use, status_html
def check_system_health() -> dict[str, Any]:
return {
"status": "healthy",
"app_name": settings.app_name,
"ocr_model": settings.ocr_model_name,
"ocr_use_cuda": settings.ocr_use_cuda,
"llm_model": settings.llm_model_name,
"embedding_model": settings.embedding_model,
"reranker_model": settings.reranker_model,
"gpu_enabled": settings.gpu_enabled,
}
# 2. Build interactive Gradio Dashboard (dark mode & English)
with gr.Blocks(title="DocNeedle - Document Extraction") as demo:
gr.HTML(HEADER_HTML)
with gr.Row():
# Left Column: Documents & Dynamic Fields (5 Columns)
with gr.Column(scale=5):
with gr.Group():
gr.Markdown(
"### ๐ DOCUMENTS\nUpload PDFs, scanned files, spreadsheets, or text files"
)
file_input = gr.File(
label="Upload Documents (Drag & Drop)",
file_count="multiple",
file_types=[
".pdf",
".docx",
".doc",
".xlsx",
".xls",
".png",
".jpg",
".jpeg",
".webp",
".txt",
],
)
session_chunks_state = gr.State([])
document_profile_state = gr.State(_empty_document_profile())
doc_status_view = gr.HTML(
value=make_status_html(
"Upload documents to begin. Processing will run when you click Extract Data."
)
)
with gr.Group():
with gr.Row():
gr.Markdown(
"### ๐ฏ FIELDS TO EXTRACT\nDefine target fields via form or raw JSON"
)
with gr.Row(elem_classes=["dn-header-actions"]):
import_schema_btn = gr.UploadButton(
"๐ฅ Import",
file_types=[".json"],
elem_classes=["dn-action-btn"],
size="sm",
)
export_schema_btn = gr.DownloadButton(
"๐ค Export",
elem_classes=["dn-action-btn"],
size="sm",
)
clear_all_btn = gr.Button(
"๐๏ธ Clear",
elem_classes=["dn-action-btn", "dn-clear-btn"],
size="sm",
)
active_tab_state = gr.State("form")
with gr.Tabs(elem_classes=["dn-input-tabs"]) as input_tabs:
with gr.Tab("๐ Form Builder", id="tab_form_input") as tab_form_input:
vis_state = gr.State(
[True] * len(INITIAL_STARTER_FIELDS)
+ [False] * (MAX_FIELDS - len(INITIAL_STARTER_FIELDS))
)
field_groups: list[Any] = []
q_inputs: list[Any] = []
k_inputs: list[Any] = []
t_inputs: list[Any] = []
del_btns: list[Any] = []
for i in range(MAX_FIELDS):
is_init_visible = i < len(INITIAL_STARTER_FIELDS)
init_q = INITIAL_STARTER_FIELDS[i][0] if is_init_visible else ""
init_k = INITIAL_STARTER_FIELDS[i][1] if is_init_visible else ""
init_t = INITIAL_STARTER_FIELDS[i][2] if is_init_visible else "string"
with gr.Group(
visible=is_init_visible, elem_classes=["dn-field-card"]
) as f_grp:
with gr.Row():
q_box = gr.Textbox(
value=init_q,
label="Field Question / Prompt",
placeholder="e.g. What is the total invoice amount including VAT?",
scale=9,
)
d_btn = gr.Button(
"๐๏ธ",
scale=1,
min_width=36,
elem_classes=["dn-del-btn"],
)
with gr.Row():
k_box = gr.Textbox(
value=init_k,
label="Target Key (Optional)",
placeholder="e.g. total_amount",
scale=1,
)
t_dropdown = gr.Dropdown(
choices=FIELD_TYPE_KEYS,
value=init_t,
label="Expected Type",
scale=1,
)
field_groups.append(f_grp)
q_inputs.append(q_box)
k_inputs.append(k_box)
t_inputs.append(t_dropdown)
del_btns.append(d_btn)
add_field_btn = gr.Button("๏ผ Add Field", elem_classes=["dn-add-field-btn"])
with gr.Tab("๐ Raw JSON View", id="tab_raw_input") as tab_raw_input:
gr.Markdown(
"*Enter or edit schema directly as a JSON list of objects (`key`, `description`, `type`).*",
elem_classes=["dn-raw-hint"],
)
raw_json_input = gr.Code(
value=serialize_form_to_json(
[f[0] for f in INITIAL_STARTER_FIELDS],
[f[1] for f in INITIAL_STARTER_FIELDS],
[f[2] for f in INITIAL_STARTER_FIELDS],
),
language="json",
label="JSON Schema Specification",
lines=12,
)
with gr.Row():
apply_raw_to_form_btn = gr.Button(
"โฌ๏ธ Apply to Form",
elem_classes=["dn-secondary-btn"],
size="sm",
)
sync_form_to_raw_btn = gr.Button(
"โฌ๏ธ Sync from Form",
elem_classes=["dn-secondary-btn"],
size="sm",
)
extract_btn = gr.Button(
"Extract Data",
variant="primary",
elem_classes=["dn-primary-btn"],
)
# Right Column: Result Cards & Provenance (7 Columns) with Dual Views
with gr.Column(scale=7):
with gr.Tabs(elem_classes=["dn-output-tabs"]):
with gr.Tab("๐ Formatted View", id="tab_formatted"):
results_html_view = gr.HTML(value=render_results_html(None))
with gr.Tab("๐ Raw JSON View", id="tab_raw"):
raw_json_output = gr.JSON(label="Complete JSON Extraction Response")
with gr.Accordion("โ๏ธ System Health & API Info", open=False):
health_btn = gr.Button("Check System Health", size="sm")
health_json = gr.JSON(label="System Details")
health_btn.click(
check_system_health, outputs=health_json, api_name=False, queue=False
)
gr.Markdown(
"""
* **Swagger API Documentation:** [/docs](/docs)
* **Extraction Endpoint:** `POST /v1/extract`
* **Health Endpoint:** `GET /healthz`
"""
)
# Event handlers for dynamic schema builder (all internal UI events with api_name=False)
def on_add_field(vis: list[bool], *field_args: Any) -> list[Any]:
new_vis = list(vis)
q_vals = field_args[0:MAX_FIELDS]
k_vals = field_args[MAX_FIELDS : 2 * MAX_FIELDS]
for idx in range(len(new_vis)):
q = str(q_vals[idx]).strip() if idx < len(q_vals) and q_vals[idx] else ""
k = str(k_vals[idx]).strip() if idx < len(k_vals) and k_vals[idx] else ""
if q or k:
new_vis[idx] = True
for idx in range(len(new_vis)):
if not new_vis[idx]:
new_vis[idx] = True
break
return [new_vis] + [gr.update(visible=v) for v in new_vis]
add_field_btn.click(
on_add_field,
inputs=[vis_state] + q_inputs + k_inputs,
outputs=[vis_state] + field_groups,
api_name=False,
queue=False,
)
def make_delete_handler(target_idx: int):
def on_delete(vis: list[bool]) -> tuple[Any, ...]:
new_vis = list(vis)
new_vis[target_idx] = False
return (
new_vis,
gr.update(visible=False),
gr.update(value=""),
gr.update(value=""),
gr.update(value="string"),
)
return on_delete
for idx, d_btn in enumerate(del_btns):
d_btn.click(
make_delete_handler(idx),
inputs=[vis_state],
outputs=[vis_state, field_groups[idx], q_inputs[idx], k_inputs[idx], t_inputs[idx]],
api_name=False,
queue=False,
)
def on_clear_all() -> list[Any]:
new_vis, grp_ups, q_ups, k_ups, t_ups, empty_json = apply_items_to_form_updates([])
return [new_vis] + grp_ups + q_ups + k_ups + t_ups + [empty_json]
clear_all_btn.click(
on_clear_all,
outputs=[vis_state] + field_groups + q_inputs + k_inputs + t_inputs + [raw_json_input],
api_name=False,
queue=False,
)
def on_sync_form_to_raw(vis: list[bool], *field_args: Any) -> str:
q_vals, k_vals, t_vals = unpack_field_args(field_args)
return serialize_form_to_json(q_vals, k_vals, t_vals, vis)
sync_form_to_raw_btn.click(
on_sync_form_to_raw,
inputs=[vis_state] + q_inputs + k_inputs + t_inputs,
outputs=[raw_json_input],
api_name=False,
queue=False,
)
def on_apply_raw_to_form(raw_json: str) -> tuple[Any, ...]:
try:
items = parse_raw_json_schema(raw_json)
if not items:
gr.Warning("No valid fields found in raw JSON schema.")
return (gr.skip(),) * (1 + MAX_FIELDS * 4)
new_vis, grp_ups, q_ups, k_ups, t_ups, _ = apply_items_to_form_updates(items)
gr.Info(f"Loaded {len(items)} fields into Form Builder.")
return tuple([new_vis] + grp_ups + q_ups + k_ups + t_ups)
except Exception as e:
gr.Warning(f"Invalid JSON: {e}")
return (gr.skip(),) * (1 + MAX_FIELDS * 4)
apply_raw_to_form_btn.click(
on_apply_raw_to_form,
inputs=[raw_json_input],
outputs=[vis_state] + field_groups + q_inputs + k_inputs + t_inputs,
api_name=False,
queue=False,
)
def on_import_schema(file: Any) -> tuple[Any, ...]:
if not file:
return (gr.skip(),) * (1 + MAX_FIELDS * 4 + 1)
try:
file_path = Path(file.name if hasattr(file, "name") else str(file))
content = file_path.read_text(encoding="utf-8")
items = parse_raw_json_schema(content)
if not items:
gr.Warning("No valid extraction fields found in imported JSON file.")
return (gr.skip(),) * (1 + MAX_FIELDS * 4 + 1)
new_vis, grp_ups, q_ups, k_ups, t_ups, json_str = apply_items_to_form_updates(items)
gr.Info(f"Successfully imported {len(items)} fields from file.")
return tuple([new_vis] + grp_ups + q_ups + k_ups + t_ups + [json_str])
except Exception as e:
gr.Warning(f"Import failed: {e}")
return (gr.skip(),) * (1 + MAX_FIELDS * 4 + 1)
import_schema_btn.upload(
on_import_schema,
inputs=[import_schema_btn],
outputs=[vis_state] + field_groups + q_inputs + k_inputs + t_inputs + [raw_json_input],
api_name=False,
queue=False,
)
def on_export_schema(
active_tab: str,
raw_json: str,
vis: list[bool],
*field_args: Any,
) -> str:
items: list[dict[str, str]] = []
if active_tab == "raw" and raw_json and raw_json.strip():
try:
items = parse_raw_json_schema(raw_json)
except Exception:
items = []
if not items:
q_vals, k_vals, t_vals = unpack_field_args(field_args)
items = collect_form_to_items_or_default(q_vals, k_vals, t_vals, vis)
tmp = tempfile.NamedTemporaryFile(
mode="w",
suffix=".json",
prefix="docneedle_schema_",
delete=False,
encoding="utf-8",
)
json.dump(items, tmp, indent=2, ensure_ascii=False)
tmp.close()
return tmp.name
export_schema_btn.click(
on_export_schema,
inputs=[active_tab_state, raw_json_input, vis_state] + q_inputs + k_inputs + t_inputs,
outputs=[export_schema_btn],
api_name=False,
queue=False,
)
tab_form_input.select(lambda: "form", outputs=[active_tab_state], api_name=False, queue=False)
def on_select_raw_tab(vis: list[bool], *field_args: Any) -> tuple[str, str]:
q_vals, k_vals, t_vals = unpack_field_args(field_args)
json_str = serialize_form_to_json(q_vals, k_vals, t_vals, vis)
return "raw", json_str
tab_raw_input.select(
on_select_raw_tab,
inputs=[vis_state] + q_inputs + k_inputs + t_inputs,
outputs=[active_tab_state, raw_json_input],
api_name=False,
queue=False,
)
def on_documents_changed(files: Any) -> tuple[str, list[Any], dict[str, int]]:
file_list = _as_file_list(files)
if not file_list:
return (
make_status_html(
"Upload documents to begin. Processing will run when you click Extract Data."
),
[],
_empty_document_profile(),
)
document_profile = _build_document_profile(file_list)
names = [_uploaded_file_path(file_value).name for file_value in file_list]
total_kb = document_profile["total_bytes"] / 1024
names_str = ", ".join(names[:3]) + (f" (+{len(names) - 3} more)" if len(names) > 3 else "")
status_text = (
f"๐ {len(file_list)} document(s) selected ({names_str} ยท {total_kb:.1f} KB). "
"Click Extract Data to process."
)
return make_status_html(status_text, is_success=False), [], document_profile
file_input.change(
on_documents_changed,
inputs=[file_input],
outputs=[doc_status_view, session_chunks_state, document_profile_state],
api_name=False,
queue=False,
)
extract_btn.click(
prepare_documents_ui,
inputs=[file_input, session_chunks_state],
outputs=[session_chunks_state, doc_status_view],
api_name=False,
).then(
perform_extraction_ui,
inputs=[
file_input,
session_chunks_state,
document_profile_state,
active_tab_state,
raw_json_input,
vis_state,
]
+ q_inputs
+ k_inputs
+ t_inputs,
outputs=[results_html_view, raw_json_output, session_chunks_state, doc_status_view],
api_name=False,
)
# ZeroGPU keepalive and compatibility handler
try:
@gpu_decorator(duration=60)
def _zerogpu_compat_probe() -> str:
return ""
_zg_btn = gr.Button("ZeroGPU Init", visible=False)
_zg_btn.click(_zerogpu_compat_probe, outputs=[], api_name=False)
except Exception:
pass
# 3. Mount complete FastAPI REST API onto Gradio App
from gradio.routes import App # noqa: E402
from docneedle.api.routes_extract import router as extract_router # noqa: E402
from docneedle.api.routes_jobs import router as jobs_router # noqa: E402
from docneedle.main import health_check # noqa: E402
gradio_app = App(
title=settings.app_name,
version=settings.app_version,
description="Privacy-first local RAG document extraction microservice",
docs_url="/docs",
redoc_url="/redoc",
)
gradio_app.include_router(extract_router)
gradio_app.include_router(jobs_router)
gradio_app.add_api_route("/healthz", health_check, methods=["GET"])
gradio_app.add_api_route("/health", health_check, methods=["GET"])
if __name__ == "__main__":
app = gradio_app
port = int(os.environ.get("PORT", 7860))
logger.info(f"--> [STARTUP] Starting DocNeedle on 0.0.0.0:{port}...")
demo.launch(
_app=gradio_app,
server_name="0.0.0.0",
server_port=port,
footer_links=[],
css=CUSTOM_CSS,
ssr_mode=False,
)
else:
# Mount Gradio dashboard onto root path for external ASGI runners (e.g. uvicorn app:app)
app = gr.mount_gradio_app(
gradio_app,
demo,
path="/",
footer_links=[],
css=CUSTOM_CSS,
ssr_mode=False,
)