820 lines
31 KiB
Python
820 lines
31 KiB
Python
import json
|
|
import mimetypes
|
|
import os
|
|
import re
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
import uuid
|
|
from typing import Literal
|
|
|
|
from flask import Flask, Response, jsonify, request, send_file, send_from_directory, stream_with_context
|
|
from flask_cors import CORS
|
|
from pydantic import BaseModel
|
|
from werkzeug.exceptions import NotFound
|
|
from werkzeug.security import safe_join
|
|
|
|
import analytics
|
|
import models
|
|
from ai_generation import generate_field_values
|
|
from briefs import _preprocess_intent
|
|
from chat_router import classify_chat_route
|
|
from config import (
|
|
ASSETS_DIR,
|
|
FAST_MODEL,
|
|
FLASK_DEBUG,
|
|
JAVA_BACKEND_API_KEY,
|
|
JAVA_REQUEST_TIMEOUT_SECONDS,
|
|
OUTPUT_DIR,
|
|
SMART_MODEL,
|
|
logger,
|
|
model_max_tokens,
|
|
)
|
|
from document_types import detect_document_type
|
|
from editing import register_edit_routes
|
|
from editing.decisions import answer_conversational_info
|
|
from file_processing_agent import ToolCatalogService
|
|
from html_pdf_utils import compile_html_to_pdf
|
|
from html_utils import inject_theme
|
|
from java_client import java_headers, java_url
|
|
from llm_utils import run_ai
|
|
from pdf_generator import _DEFAULT_TEMPLATES_DIR, PDFGenerator
|
|
from pdf_text_editor import convert_pdf_to_text_editor_document
|
|
from prompts import (
|
|
generate_all_sections_system_prompt,
|
|
pdf_qa_system_prompt,
|
|
section_fill_system_prompt,
|
|
)
|
|
from smart_folder_creator import create_smart_folder_config
|
|
from storage import load_versions
|
|
|
|
app = Flask(__name__)
|
|
CORS(app)
|
|
_tool_catalog_service = ToolCatalogService()
|
|
register_edit_routes(app)
|
|
|
|
|
|
@app.before_request
|
|
def log_job_request_sequence() -> None:
|
|
job_id = request.headers.get("X-Job-Id")
|
|
if not job_id:
|
|
return
|
|
seq = request.headers.get("X-Job-Seq", "?")
|
|
total = request.headers.get("X-Job-Total", "?")
|
|
logger.info("[HTTP] job_id=%s req=%s/%s %s %s", job_id, seq, total, request.method, request.path)
|
|
|
|
|
|
def _json_body[T: BaseModel](model: type[T], request_type: Literal["GET", "POST"] = "POST") -> T:
|
|
if request_type == "GET":
|
|
payload = request.args.to_dict()
|
|
else:
|
|
payload = request.get_json(silent=True)
|
|
return model.model_validate(payload or {})
|
|
|
|
|
|
def _java_request_json[T: BaseModel](
|
|
method: str,
|
|
path: str,
|
|
payload: BaseModel | None,
|
|
response_model: type[T],
|
|
) -> T:
|
|
url = java_url(path)
|
|
data = None
|
|
headers = java_headers()
|
|
headers["Content-Type"] = "application/json"
|
|
if payload is not None:
|
|
payload_data = payload.model_dump(by_alias=True, exclude_none=True)
|
|
data = json.dumps(payload_data).encode("utf-8")
|
|
req = urllib.request.Request(url, data=data, headers=headers, method=method)
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=JAVA_REQUEST_TIMEOUT_SECONDS) as resp:
|
|
body = resp.read().decode("utf-8")
|
|
parsed = json.loads(body) if body else {}
|
|
return response_model.model_validate(parsed)
|
|
except urllib.error.HTTPError as exc:
|
|
detail = exc.read().decode("utf-8") if exc.fp else ""
|
|
logger.error("[JAVA] %s %s failed status=%s detail=%s", method, path, exc.code, detail)
|
|
raise
|
|
|
|
|
|
def _fetch_ai_session(session_id: str):
|
|
auth_header = request.headers.get("Authorization")
|
|
api_key = request.headers.get("X-API-KEY")
|
|
if auth_header or api_key:
|
|
path = f"/api/v1/ai/create/sessions/{session_id}"
|
|
elif JAVA_BACKEND_API_KEY:
|
|
path = f"/api/v1/ai/create/internal/sessions/{session_id}"
|
|
else:
|
|
path = f"/api/v1/ai/create/sessions/{session_id}"
|
|
return _java_request_json("GET", path, None, models.AISession)
|
|
|
|
|
|
def _update_ai_session(session_id: str, payload: models.JavaUpdateSessionRequest):
|
|
return _java_request_json(
|
|
"POST", f"/api/v1/ai/create/internal/sessions/{session_id}/update", payload, models.AISession
|
|
)
|
|
|
|
|
|
@app.route("/api/intent/check", methods=["POST"])
|
|
def intent_check():
|
|
payload = _json_body(models.IntentCheckRequest)
|
|
prompt = payload.prompt
|
|
history = payload.conversation_history
|
|
current_pdf_url = payload.current_pdf_url
|
|
intent = _preprocess_intent(prompt, history, bool(current_pdf_url))
|
|
response = intent.model_copy(update={"doc_type": intent.document_type, "has_pdf": bool(current_pdf_url)})
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
@app.route("/api/chat/route", methods=["POST"])
|
|
def chat_route():
|
|
payload = _json_body(models.ChatRouteRequest)
|
|
decision = classify_chat_route(payload)
|
|
return jsonify(decision.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
@app.route("/api/chat/create-smart-folder", methods=["POST"])
|
|
def create_smart_folder():
|
|
payload = _json_body(models.SmartFolderCreateRequest)
|
|
result = create_smart_folder_config(payload)
|
|
return jsonify(result.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
@app.route("/api/chat/interpret-parameter", methods=["POST"])
|
|
def interpret_parameter():
|
|
"""Use AI to interpret an ambiguous parameter response during Tier 2 workflow collection."""
|
|
payload = _json_body(models.InterpretParameterRequest)
|
|
system_prompt = (
|
|
"You are a parameter extraction assistant for a PDF tool workflow.\n"
|
|
"Given a question asked to the user and the user's response, decide how to interpret it.\n\n"
|
|
"Return JSON with these fields:\n"
|
|
'- "type": one of "value", "confused", "cancel", "default"\n'
|
|
' - "value": the user gave a usable answer, extract it cleanly\n'
|
|
' - "confused": the user is asking for help or does not understand\n'
|
|
' - "cancel": the user wants to stop or cancel\n'
|
|
' - "default": the user is vague; suggest the most sensible default\n'
|
|
'- "extracted_value": the clean extracted value (only for type="value" or type="default")\n'
|
|
'- "help_message": a short helpful message explaining what was interpreted or what options exist\n\n'
|
|
'Be concise. For type="value", extracted_value should be just the raw value (e.g. "English", "3", "DOCX").\n'
|
|
'For type="default", extracted_value should be the suggested default value.'
|
|
)
|
|
user_prompt = (
|
|
f"Tool: {payload.tool_name}\nQuestion asked: {payload.question}\nUser response: {payload.user_response}"
|
|
)
|
|
messages = [
|
|
models.ChatMessage(role="system", content=system_prompt),
|
|
models.ChatMessage(role="user", content=user_prompt),
|
|
]
|
|
result = run_ai(
|
|
FAST_MODEL,
|
|
messages,
|
|
models.InterpretParameterResponse,
|
|
tag="interpret_parameter",
|
|
log_label="interpret-parameter",
|
|
)
|
|
return jsonify(result.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
@app.route("/api/chat/infer-tools", methods=["POST"])
|
|
def infer_tools():
|
|
"""Use AI to infer which PDF tools to apply from a freeform user message."""
|
|
payload = _json_body(models.InferToolsRequest)
|
|
tools_list = "\n".join(f"- {t}" for t in payload.available_tools)
|
|
system_prompt = (
|
|
"You are a tool selection assistant for Stirling PDF.\n"
|
|
"Given a user message, identify which PDF tools to apply to achieve the user's goal.\n"
|
|
"Return tools in the execution order.\n\n"
|
|
"Rules:\n"
|
|
"- Only return tools with HIGH confidence — leave out anything uncertain\n"
|
|
"- Return empty list if the request is a question, ambiguous, or conversational\n"
|
|
"- Return empty list if the request is about creating a new document\n"
|
|
"- Maximum 4 tools\n\n"
|
|
f"Available tools (format: 'id: Display Name'):\n{tools_list}\n\n"
|
|
'Return JSON: { "tools": [{"tool_id": "...", "confidence": "high"|"medium"|"low"}], "reason": "..." }\n'
|
|
'IMPORTANT: tool_id must be the exact identifier before the colon (e.g. "convert", not "convert: Convert PDF").'
|
|
)
|
|
messages_list = [
|
|
models.ChatMessage(role="system", content=system_prompt),
|
|
models.ChatMessage(role="user", content=f"User message: {payload.message}"),
|
|
]
|
|
result = run_ai(
|
|
FAST_MODEL,
|
|
messages_list,
|
|
models.InferToolsResponse,
|
|
tag="infer_tools",
|
|
log_label="infer-tools",
|
|
)
|
|
filtered = models.InferToolsResponse(
|
|
tools=[t for t in result.tools if t.confidence == "high"],
|
|
reason=result.reason,
|
|
)
|
|
return jsonify(filtered.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
@app.route("/api/chat/info", methods=["POST"])
|
|
def chat_info():
|
|
"""Handle conversational queries without requiring a file or session."""
|
|
payload = _json_body(models.ChatInfoRequest)
|
|
tool_catalog = ToolCatalogService()
|
|
assistant_message = answer_conversational_info(
|
|
payload.message,
|
|
payload.history,
|
|
tool_catalog,
|
|
)
|
|
response = models.ChatInfoResponse(assistant_message=assistant_message)
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
@app.route("/api/detect_type", methods=["POST"])
|
|
def detect_type():
|
|
"""
|
|
Detect document type from a prompt using a fast AI model.
|
|
|
|
This endpoint uses a cheap/fast model (FAST_MODEL) for AI classification
|
|
to minimize cost and latency.
|
|
|
|
Request body:
|
|
- prompt: The user's document request
|
|
- explicitType: If provided, skip detection and return this type
|
|
|
|
Response:
|
|
- docType: The detected type
|
|
- confidence: "medium" (AI), "low" (AI)
|
|
- method: "explicit" or "ai"
|
|
"""
|
|
payload = _json_body(models.DetectTypeRequest)
|
|
prompt = payload.prompt
|
|
explicit_type = payload.explicit_type or ""
|
|
|
|
# If user explicitly provided a type, skip detection
|
|
if explicit_type and explicit_type not in ("other", "document", "miscellaneous", ""):
|
|
logger.info("[DETECT] Using explicit doc_type=%s, skipping detection", explicit_type)
|
|
response = models.DetectTypeResponse(doc_type=explicit_type, confidence="high", method="explicit")
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
if not prompt:
|
|
response = models.DetectTypeResponse(error="Missing prompt for detection")
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True)), 400
|
|
|
|
# Use fast AI detection (uses FAST_MODEL for cheap classification)
|
|
ai_type, confidence = detect_document_type(prompt, confidence_threshold=0.7)
|
|
logger.info("[DETECT] Fast AI detected doc_type=%s (confidence=%.2f)", ai_type, confidence)
|
|
response = models.DetectTypeResponse(
|
|
doc_type=ai_type or "other",
|
|
confidence="medium" if confidence >= 0.7 else "low",
|
|
method="ai",
|
|
)
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
@app.route("/api/pdf/answer", methods=["POST"])
|
|
def pdf_answer():
|
|
payload = _json_body(models.PdfAnswerRequest)
|
|
pdf_url = payload.pdf_url
|
|
question = payload.question
|
|
if not pdf_url or not question:
|
|
response = models.PdfAnswerResponse(error="Missing pdfUrl or question")
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True)), 400
|
|
|
|
filename = os.path.basename(pdf_url.split("?")[0])
|
|
if not filename.lower().endswith(".pdf"):
|
|
response = models.PdfAnswerResponse(error="Invalid pdf file")
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True)), 400
|
|
|
|
pdf_path = safe_join(OUTPUT_DIR, filename)
|
|
if pdf_path is None or not os.path.exists(pdf_path):
|
|
response = models.PdfAnswerResponse(error="PDF not found")
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True)), 404
|
|
|
|
doc = convert_pdf_to_text_editor_document(pdf_path)
|
|
|
|
pages = doc.document.pages if doc else []
|
|
snippets: list[str] = []
|
|
for page in pages:
|
|
for elem in page.text_elements:
|
|
text = elem.text
|
|
if text:
|
|
snippets.append(str(text))
|
|
if not snippets:
|
|
response = models.PdfAnswerResponse(error="No readable text in PDF")
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True)), 400
|
|
|
|
# Normalize and limit context
|
|
context = " ".join(snippets)
|
|
context = " ".join(context.split()) # normalize whitespace
|
|
max_context = 10000
|
|
if len(context) > max_context:
|
|
context = context[:max_context]
|
|
|
|
model_name = SMART_MODEL
|
|
system_prompt = pdf_qa_system_prompt() + "\nReturn JSON matching the provided schema."
|
|
user_prompt = f"Question: {question}\n\nPDF text:\n{context}"
|
|
messages = [
|
|
models.ChatMessage(role="system", content=system_prompt),
|
|
models.ChatMessage(role="user", content=user_prompt),
|
|
]
|
|
response = run_ai(
|
|
model_name,
|
|
messages,
|
|
models.PdfAnswer,
|
|
tag="pdf_answer",
|
|
max_tokens=220,
|
|
)
|
|
answer = response.answer.strip()
|
|
title_like = re.match(r"^why pdfs|^minimalist|^author:", answer, re.IGNORECASE)
|
|
normalized_answer = re.sub(r"\s+", " ", answer).strip().lower()
|
|
normalized_context = re.sub(r"\s+", " ", context).strip().lower()
|
|
copied_context = bool(normalized_answer) and normalized_answer in normalized_context
|
|
if title_like or copied_context:
|
|
response = models.PdfAnswerResponse(error="AI answer was invalid or echoed the source text.")
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True)), 500
|
|
|
|
response = models.PdfAnswerResponse(answer=answer, mode="model")
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
@app.route("/api/create/sessions", methods=["POST"])
|
|
def create_session():
|
|
"""Create a new AI session via the Java backend."""
|
|
create_session_request = _json_body(models.CreateSessionRequest)
|
|
|
|
# Doc type should ideally come from chat router; fall back to detection if needed
|
|
detection_method = "provided"
|
|
if create_session_request.doc_type in ("miscellaneous", "other", "document", "unknown", ""):
|
|
# Fall back to detection if doc_type not provided by chat router
|
|
logger.info("[SESSION] No doc_type provided, running fallback detection")
|
|
confidence_threshold = 0.7 # Only accept matches with 70%+ confidence
|
|
detected_type, confidence = detect_document_type(
|
|
create_session_request.prompt, confidence_threshold=confidence_threshold
|
|
)
|
|
if confidence >= confidence_threshold:
|
|
logger.info(
|
|
"[SESSION] Detected doc_type=%s from prompt (confidence=%.2f)",
|
|
detected_type,
|
|
confidence,
|
|
)
|
|
create_session_request.doc_type = detected_type
|
|
detection_method = "fallback_detection"
|
|
else:
|
|
logger.info(
|
|
"[SESSION] Detection confidence too low (%.2f < %.2f), using 'other'",
|
|
confidence,
|
|
confidence_threshold,
|
|
)
|
|
create_session_request.doc_type = "other"
|
|
detection_method = "fallback_low_confidence"
|
|
else:
|
|
logger.info(
|
|
"[SESSION] Using doc_type=%s from chat router",
|
|
create_session_request.doc_type,
|
|
)
|
|
|
|
result = _java_request_json(
|
|
"POST",
|
|
"/api/v1/ai/create/sessions",
|
|
create_session_request,
|
|
models.JavaCreateSessionResponse,
|
|
)
|
|
session_id = result.session_id
|
|
logger.info(
|
|
"[SESSION] Created session %s doc_type=%s method=%s",
|
|
session_id,
|
|
create_session_request.doc_type,
|
|
detection_method,
|
|
)
|
|
# Track session creation in PostHog
|
|
safe_doc_type = re.sub(r"[^a-zA-Z0-9_]+", "", (create_session_request.doc_type or "").lower())
|
|
has_html_template = (_DEFAULT_TEMPLATES_DIR / f"{safe_doc_type}.html").exists()
|
|
analytics.track_session_created(
|
|
user_id=session_id,
|
|
session_id=session_id,
|
|
doc_type=create_session_request.doc_type or "unknown",
|
|
template_id=create_session_request.template_id,
|
|
has_template=has_html_template,
|
|
)
|
|
|
|
response = models.CreateSessionResponse(
|
|
session_id=session_id,
|
|
doc_type=create_session_request.doc_type,
|
|
detection_method=detection_method,
|
|
)
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
@app.route("/api/create/sessions/<session_id>/outline", methods=["POST"])
|
|
def update_outline(session_id: str):
|
|
"""Update session with approved outline."""
|
|
payload = _json_body(models.UpdateOutlineRequest)
|
|
_java_request_json(
|
|
"POST",
|
|
f"/api/v1/ai/create/sessions/{session_id}/outline",
|
|
payload,
|
|
models.AISession,
|
|
)
|
|
response = models.SuccessResponse(success=True)
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
@app.route("/api/create/sessions/<session_id>/draft", methods=["POST"])
|
|
def update_draft(session_id: str):
|
|
"""Update session with approved draft sections."""
|
|
payload = _json_body(models.UpdateDraftRequest)
|
|
_java_request_json(
|
|
"POST",
|
|
f"/api/v1/ai/create/sessions/{session_id}/draft",
|
|
payload,
|
|
models.AISession,
|
|
)
|
|
response = models.SuccessResponse(success=True)
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
@app.route("/api/create/sessions/<session_id>/template", methods=["POST"])
|
|
def update_template(session_id: str):
|
|
"""Update session template selection."""
|
|
payload = _json_body(models.UpdateTemplateRequest)
|
|
_java_request_json(
|
|
"POST",
|
|
f"/api/v1/ai/create/sessions/{session_id}/template",
|
|
payload,
|
|
models.AISession,
|
|
)
|
|
response = models.SuccessResponse(success=True)
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
@app.route("/api/create/sessions/<session_id>/reprompt", methods=["POST"])
|
|
def reprompt_session(session_id: str):
|
|
"""Update session with new prompt."""
|
|
payload = _json_body(models.RepromptRequest)
|
|
_java_request_json(
|
|
"POST",
|
|
f"/api/v1/ai/create/sessions/{session_id}/reprompt",
|
|
payload,
|
|
models.AISession,
|
|
)
|
|
response = models.SuccessResponse(success=True)
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
@app.route("/api/create/sessions/<session_id>/stream", methods=["POST"])
|
|
def create_stream(session_id: str):
|
|
payload = _json_body(models.CreateStreamRequest)
|
|
session = _fetch_ai_session(session_id)
|
|
prompt = session.prompt_latest or session.prompt_initial or ""
|
|
doc_type = session.doc_type or "other"
|
|
template_id = session.template_id
|
|
outline_text = session.outline_text or ""
|
|
outline_filename = session.outline_filename
|
|
constraints = session.outline_constraints
|
|
draft_sections = session.draft_sections
|
|
logo_base64 = payload.theme.logo_base64 if payload.theme else None
|
|
theme = payload.theme.css_overrides() if payload.theme else None
|
|
logger.info(
|
|
"[STREAM] session_id=%s has_theme=%s has_logo=%s",
|
|
session_id,
|
|
bool(theme),
|
|
bool(logo_base64),
|
|
)
|
|
base_html = payload.base_html or session.polished_html or None
|
|
handler = PDFGenerator(
|
|
session_id=session_id,
|
|
phase=payload.phase,
|
|
prompt=prompt,
|
|
doc_type=doc_type,
|
|
template_id=template_id,
|
|
outline_text=outline_text,
|
|
outline_filename=outline_filename,
|
|
constraints=constraints,
|
|
draft_sections=draft_sections,
|
|
update_session=_update_ai_session,
|
|
theme=theme,
|
|
logo_base64=logo_base64,
|
|
base_html=base_html,
|
|
instructions=payload.additional_instructions,
|
|
)
|
|
|
|
return Response(
|
|
stream_with_context(handler.generate()),
|
|
mimetype="text/event-stream",
|
|
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
|
|
)
|
|
|
|
|
|
@app.route("/api/create/sessions/<session_id>/fields", methods=["POST"])
|
|
def fill_fields(session_id: str):
|
|
logger.info("[AI create] fill_fields session_id=%s", session_id)
|
|
session = _fetch_ai_session(session_id)
|
|
|
|
payload = _json_body(models.FillFieldsRequest)
|
|
fields = payload.fields
|
|
extra_prompt = payload.extra_prompt
|
|
if not isinstance(fields, list):
|
|
return jsonify({"error": "Fields must be a list"}), 400
|
|
|
|
prompt = session.prompt_latest or session.prompt_initial or ""
|
|
if extra_prompt:
|
|
prompt = f"{prompt}\n{extra_prompt}"
|
|
doc_type = session.doc_type or "other"
|
|
constraints = session.outline_constraints
|
|
|
|
filled = generate_field_values(prompt, doc_type, fields, constraints)
|
|
response = models.FillFieldsResponse(fields=filled)
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
@app.route("/api/generate_section", methods=["POST"])
|
|
def generate_section():
|
|
"""Generate content for a single section based on a custom prompt."""
|
|
payload = _json_body(models.GenerateSectionRequest)
|
|
_session_id = payload.session_id
|
|
section_label = payload.section_label
|
|
section_index = payload.section_index
|
|
custom_prompt = payload.custom_prompt
|
|
document_prompt = payload.document_prompt
|
|
document_type = payload.doc_type
|
|
existing_sections = payload.existing_sections
|
|
|
|
logger.info(
|
|
"[AI] generate_section label=%s custom_prompt=%s", section_label, custom_prompt[:50] if custom_prompt else ""
|
|
)
|
|
|
|
# Build context from existing sections (don't truncate - AI needs full context for calculations)
|
|
sections_context = "\n".join([f"- {sec.label}: {sec.value}" for sec in existing_sections if sec.value])
|
|
system_prompt = section_fill_system_prompt(
|
|
document_type,
|
|
document_prompt[:500],
|
|
sections_context,
|
|
section_label,
|
|
custom_prompt,
|
|
)
|
|
|
|
messages = [
|
|
models.ChatMessage(role="system", content=system_prompt),
|
|
models.ChatMessage(role="user", content=f"Generate content for: {section_label}"),
|
|
]
|
|
parsed = run_ai(
|
|
SMART_MODEL,
|
|
messages,
|
|
models.SectionContent,
|
|
tag="generate_section",
|
|
max_tokens=500,
|
|
)
|
|
response = models.GenerateSectionResponse(content=parsed.content.strip(), section_index=section_index)
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
def _run_sections_batch(
|
|
batch: list[tuple[int, str]],
|
|
document_prompt: str,
|
|
additional_prompt: str | None,
|
|
system_prompt: str,
|
|
) -> dict[int, str]:
|
|
sections_list = "\n".join([f"{i + 1}. {label}" for i, label in batch])
|
|
user_prompt = (
|
|
f"User's document request:\n{document_prompt}\n\n"
|
|
"Generate ONLY these sections (leave others unchanged):\n"
|
|
f"{sections_list}"
|
|
)
|
|
if additional_prompt:
|
|
user_prompt = f"{user_prompt}\n\nAdditional instructions:\n{additional_prompt}"
|
|
messages = [
|
|
models.ChatMessage(role="system", content=system_prompt),
|
|
models.ChatMessage(role="user", content=user_prompt),
|
|
]
|
|
parsed = run_ai(
|
|
SMART_MODEL,
|
|
messages,
|
|
models.LLMGenerateAllSectionsResponse,
|
|
tag="generate_all_sections",
|
|
max_tokens=model_max_tokens(SMART_MODEL),
|
|
)
|
|
return {section.index - 1: section.value for section in parsed.sections if section.value}
|
|
|
|
|
|
@app.route("/api/generate_all_sections", methods=["POST"])
|
|
def generate_all_sections():
|
|
"""Generate content for selected sections."""
|
|
generation_start = time.time()
|
|
payload = _json_body(models.GenerateAllSectionsRequest)
|
|
_session_id = payload.session_id
|
|
document_prompt = payload.document_prompt
|
|
document_type = payload.doc_type
|
|
current_sections = payload.sections
|
|
only_indices = payload.only_indices # Optional: only generate for these indices
|
|
additional_prompt = payload.additional_prompt # Optional: extra instructions from user
|
|
|
|
# If onlyIndices is provided, only generate for those sections
|
|
if only_indices is not None:
|
|
indices_to_generate = set(only_indices)
|
|
else:
|
|
indices_to_generate = set(range(len(current_sections)))
|
|
|
|
logger.info(
|
|
"[AI] generate_all_sections doc_type=%s total_sections=%d generating=%d",
|
|
document_type,
|
|
len(current_sections),
|
|
len(indices_to_generate),
|
|
)
|
|
|
|
if not current_sections:
|
|
response = models.GenerateAllSectionsResponse(error="No sections provided")
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True)), 400
|
|
|
|
if not indices_to_generate:
|
|
# Nothing to generate, return sections as-is
|
|
response = models.GenerateAllSectionsResponse(sections=current_sections)
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
# Build the section labels list (only for sections we're generating)
|
|
section_labels = [section.label for section in current_sections]
|
|
|
|
# Build prompt only for sections we want to generate
|
|
sections_to_generate = [(i, section_labels[i]) for i in sorted(indices_to_generate) if i < len(current_sections)]
|
|
system_prompt = generate_all_sections_system_prompt(document_type)
|
|
|
|
# Batch to avoid hitting model output token limits — each batch is a separate API call.
|
|
batch_size = 3
|
|
batches = [sections_to_generate[i : i + batch_size] for i in range(0, len(sections_to_generate), batch_size)]
|
|
|
|
generated_content: dict[int, str] = {}
|
|
failed_indices: list[int] = []
|
|
|
|
for batch in batches:
|
|
try:
|
|
generated_content.update(_run_sections_batch(batch, document_prompt, additional_prompt, system_prompt))
|
|
except Exception as batch_err:
|
|
logger.warning(
|
|
"[AI] generate_all_sections batch failed (%s), retrying each section individually",
|
|
batch_err,
|
|
)
|
|
for item in batch:
|
|
try:
|
|
generated_content.update(
|
|
_run_sections_batch([item], document_prompt, additional_prompt, system_prompt)
|
|
)
|
|
except Exception as single_err:
|
|
logger.error(
|
|
"[AI] generate_all_sections failed for section index=%d label=%r: %s",
|
|
item[0],
|
|
item[1],
|
|
single_err,
|
|
)
|
|
failed_indices.append(item[0])
|
|
|
|
# Build final sections list, keeping existing content for non-generated sections
|
|
filled_sections = []
|
|
for i, section in enumerate(current_sections):
|
|
label = section.label
|
|
if i in generated_content:
|
|
# Use newly generated content
|
|
filled_sections.append(models.DraftSection(label=label, value=generated_content[i]))
|
|
else:
|
|
# Keep existing content
|
|
filled_sections.append(models.DraftSection(label=label, value=section.value))
|
|
|
|
generation_duration = (time.time() - generation_start) * 1000
|
|
|
|
# Track section generation
|
|
analytics.track_event(
|
|
user_id=_session_id or "unknown",
|
|
event_name="sections_generated",
|
|
properties={
|
|
"session_id": _session_id,
|
|
"doc_type": document_type,
|
|
"total_sections": len(current_sections),
|
|
"generated_sections": len(indices_to_generate),
|
|
"generation_time_ms": generation_duration,
|
|
"has_additional_prompt": bool(additional_prompt),
|
|
},
|
|
)
|
|
|
|
response = models.GenerateAllSectionsResponse(
|
|
sections=filled_sections,
|
|
incomplete_section_indices=failed_indices if failed_indices else None,
|
|
)
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
@app.route("/output/<path:filename>", methods=["GET"])
|
|
def serve_output_file(filename: str):
|
|
"""Serve generated PDF files and stored assets."""
|
|
mime_type, _ = mimetypes.guess_type(filename)
|
|
try:
|
|
response = send_from_directory(OUTPUT_DIR, filename, mimetype=mime_type or "application/pdf")
|
|
file_size = response.headers.get("Content-Length", "unknown")
|
|
logger.info("[SERVE] Serving file=%s size=%s bytes mime=%s", filename, file_size, mime_type)
|
|
response.headers["Access-Control-Allow-Origin"] = "*"
|
|
response.headers["Access-Control-Allow-Methods"] = "GET"
|
|
return response
|
|
except NotFound:
|
|
logger.warning("[SERVE] File not found: %s", filename)
|
|
return jsonify({"error": "File not found"}), 404
|
|
|
|
|
|
@app.route("/api/versions/<user_id>", methods=["GET"])
|
|
def list_versions(user_id: str):
|
|
response = models.VersionsResponse(versions=load_versions(user_id))
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
@app.route("/api/assets/upload", methods=["POST"])
|
|
def upload_asset():
|
|
file = request.files.get("file")
|
|
if not file:
|
|
return jsonify({"error": "Missing file"}), 400
|
|
|
|
_, ext = os.path.splitext(file.filename or "")
|
|
ext = ext.lower()
|
|
if ext not in {".png", ".jpg", ".jpeg", ".gif"}:
|
|
return jsonify({"error": "Unsupported file type"}), 400
|
|
|
|
asset_id = f"{uuid.uuid4().hex}{ext}"
|
|
output_path = os.path.join(ASSETS_DIR, asset_id)
|
|
os.makedirs(ASSETS_DIR, exist_ok=True)
|
|
file.save(output_path)
|
|
|
|
response = models.UploadAssetResponse(
|
|
asset_id=asset_id,
|
|
asset_url=f"/output/assets/{asset_id}",
|
|
)
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
@app.route("/api/templates/previews", methods=["GET"])
|
|
def list_preview_templates():
|
|
"""Return sorted list of preview template stem names."""
|
|
stems = sorted(p.stem for p in _DEFAULT_TEMPLATES_DIR.glob("*_preview.html"))
|
|
return jsonify({"templates": stems})
|
|
|
|
|
|
@app.route("/api/templates/preview-html", methods=["GET"])
|
|
def get_template_preview_html():
|
|
"""Return preview template HTML with default theme injected.
|
|
|
|
Query param: template — stem name (e.g. 'invoice_preview' or 'invoice_preview.html').
|
|
The frontend will override CSS vars via JS for live color preview.
|
|
"""
|
|
payload = _json_body(models.PreviewTemplateHtmlRequest, "GET")
|
|
template_param = payload.template
|
|
if not template_param:
|
|
return jsonify({"error": "Missing template parameter"}), 400
|
|
|
|
# Strip .html extension if present
|
|
stem = template_param.removesuffix(".html")
|
|
if "_preview" not in stem:
|
|
return jsonify({"error": "Template must be a preview template (name must contain '_preview')"}), 400
|
|
|
|
html_path = _DEFAULT_TEMPLATES_DIR / f"{stem}.html"
|
|
if not html_path.exists():
|
|
return jsonify({"error": f"Template not found: {stem}"}), 404
|
|
|
|
html = html_path.read_text(encoding="utf-8", errors="replace")
|
|
# Inject default theme so vars are defined; frontend will override via JS
|
|
html = inject_theme(html, None)
|
|
return Response(html, mimetype="text/html")
|
|
|
|
|
|
@app.route("/api/templates/render-preview", methods=["POST"])
|
|
def render_preview_to_pdf():
|
|
"""Render a preview template to PDF with the given theme.
|
|
|
|
Body: { template: str, theme: dict | None }
|
|
Returns: PDF file download.
|
|
"""
|
|
payload = _json_body(models.RenderPreviewRequest)
|
|
template_param = payload.template
|
|
theme = payload.theme.css_overrides() if payload.theme else None
|
|
|
|
if not template_param:
|
|
return jsonify({"error": "Missing template parameter"}), 400
|
|
|
|
stem = template_param.removesuffix(".html")
|
|
if "_preview" not in stem:
|
|
return jsonify({"error": "Template must be a preview template"}), 400
|
|
|
|
html_path = _DEFAULT_TEMPLATES_DIR / f"{stem}.html"
|
|
if not html_path.exists():
|
|
return jsonify({"error": f"Template not found: {stem}"}), 404
|
|
|
|
html = html_path.read_text(encoding="utf-8", errors="replace")
|
|
html = inject_theme(html, theme)
|
|
|
|
job_id = f"preview/{uuid.uuid4().hex}"
|
|
result = compile_html_to_pdf(html, job_id, log_errors=True)
|
|
|
|
if result.pdf_path and os.path.exists(result.pdf_path):
|
|
return send_file(result.pdf_path, mimetype="application/pdf", as_attachment=True, download_name=f"{stem}.pdf")
|
|
|
|
return jsonify({"error": result.error or "PDF generation failed"}), 500
|
|
|
|
|
|
@app.route("/health", methods=["GET"])
|
|
def health():
|
|
response = models.HealthResponse(status="ok", engine="puppeteer")
|
|
return jsonify(response.model_dump(by_alias=True, exclude_none=True))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app.run(host="0.0.0.0", port=5001, debug=FLASK_DEBUG)
|