mirror of
https://github.com/Stirling-Tools/Stirling-PDF.git
synced 2026-09-02 21:03:34 +03:00
364 lines
12 KiB
Python
364 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from typing import Any
|
|
|
|
import models
|
|
from config import SMART_MODEL, STREAMING_ENABLED
|
|
from format_prompts import get_format_prompt
|
|
from html_utils import inject_theme
|
|
from llm_utils import run_ai, stream_ai
|
|
from prompts import (
|
|
field_values_system_prompt,
|
|
html_context_messages,
|
|
html_edit_system_prompt,
|
|
html_system_prompt,
|
|
outline_generator_system_prompt,
|
|
section_draft_system_prompt,
|
|
template_fill_html_system_prompt,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def generate_outline_with_llm(
|
|
prompt: str,
|
|
document_type: str,
|
|
constraints: models.Constraint | None = None,
|
|
) -> models.OutlineResponse:
|
|
"""
|
|
Generate outline for a specific document type.
|
|
|
|
document_type should already be detected - this function only generates the outline.
|
|
"""
|
|
constraint_text = ""
|
|
if constraints:
|
|
tone = constraints.tone
|
|
audience = constraints.audience
|
|
pages = constraints.page_count
|
|
constraint_text = f"Tone: {tone}. Audience: {audience}. Target pages: {pages}."
|
|
|
|
# Try to get a format-specific extraction prompt
|
|
format_prompt, default_sections = get_format_prompt(document_type)
|
|
system_prompt = outline_generator_system_prompt(document_type, constraint_text, format_prompt, default_sections)
|
|
|
|
messages: list[models.ChatMessage] = [
|
|
models.ChatMessage(role="system", content=system_prompt),
|
|
models.ChatMessage(role="user", content=prompt),
|
|
]
|
|
|
|
parsed = run_ai(
|
|
SMART_MODEL,
|
|
messages,
|
|
models.OutlineResponse,
|
|
tag="outline",
|
|
)
|
|
|
|
if parsed:
|
|
logger.info(
|
|
"[OUTLINE] Generated outline doc_type=%s sections=%d filename=%s",
|
|
parsed.doc_type,
|
|
len(parsed.sections),
|
|
parsed.outline_filename,
|
|
)
|
|
return parsed
|
|
|
|
raise RuntimeError("AI outline generation failed.")
|
|
|
|
|
|
def generate_field_values(
|
|
prompt: str,
|
|
document_type: str,
|
|
fields: list[dict[str, Any]],
|
|
constraints: models.Constraint | None = None,
|
|
) -> list[dict[str, str]]:
|
|
constraint_text = ""
|
|
if constraints:
|
|
tone = constraints.tone
|
|
audience = constraints.audience
|
|
pages = constraints.page_count
|
|
constraint_text = f"Tone: {tone}. Audience: {audience}. Target pages: {pages}."
|
|
system_prompt = field_values_system_prompt(constraint_text)
|
|
messages: list[models.ChatMessage] = [
|
|
models.ChatMessage(role="system", content=system_prompt),
|
|
models.ChatMessage(role="user", content=f"Document type: {document_type}"),
|
|
models.ChatMessage(role="user", content=f"Prompt:\n{prompt}"),
|
|
models.ChatMessage(role="user", content=f"Fields:\n{json.dumps(fields, ensure_ascii=True)}"),
|
|
]
|
|
parsed = run_ai(
|
|
SMART_MODEL,
|
|
messages,
|
|
models.LLMFieldValuesResponse,
|
|
tag="field_values",
|
|
)
|
|
if parsed:
|
|
return [{"label": item.label, "value": item.value} for item in parsed.fields]
|
|
|
|
raise RuntimeError("AI field extraction failed.")
|
|
|
|
|
|
def generate_section_draft(
|
|
prompt: str,
|
|
document_type: str,
|
|
outline_text: str,
|
|
constraints: models.Constraint | None = None,
|
|
) -> list[models.DraftSection]:
|
|
constraint_text = ""
|
|
if constraints:
|
|
tone = constraints.tone
|
|
audience = constraints.audience
|
|
pages = constraints.page_count
|
|
constraint_text = f"Tone: {tone}. Audience: {audience}. Target pages: {pages}."
|
|
system_prompt = section_draft_system_prompt(constraint_text)
|
|
messages: list[models.ChatMessage] = [
|
|
models.ChatMessage(role="system", content=system_prompt),
|
|
models.ChatMessage(role="user", content=f"Document type: {document_type}"),
|
|
models.ChatMessage(role="user", content=f"Outline:\n{outline_text}"),
|
|
models.ChatMessage(role="user", content=f"Prompt:\n{prompt}"),
|
|
]
|
|
parsed = run_ai(
|
|
SMART_MODEL,
|
|
messages,
|
|
models.LLMDraftSectionsResponse,
|
|
tag="section_draft",
|
|
)
|
|
return parsed.sections
|
|
|
|
|
|
# ── HTML generation ───────────────────────────────────────────────────────────
|
|
|
|
|
|
def generate_template_fill_html_stream(
|
|
template_html: str,
|
|
document_type: str,
|
|
outline_text: str,
|
|
draft_sections: list[models.DraftSection] | None = None,
|
|
constraints: models.Constraint | None = None,
|
|
theme: dict[str, str] | None = None,
|
|
additional_instructions: str | None = None,
|
|
has_logo: bool = False,
|
|
):
|
|
"""Fill an HTML template by replacing {{PLACEHOLDER}} tokens, streaming chunks."""
|
|
constraints_text = ""
|
|
if constraints:
|
|
tone = constraints.tone
|
|
audience = constraints.audience
|
|
pages = constraints.page_count
|
|
constraints_text = f"Tone: {tone}. Audience: {audience}. Target pages: {pages}."
|
|
if draft_sections:
|
|
constraints_text += "\nUse the section content to inform placeholder values."
|
|
|
|
# Inject theme before sending to AI if overrides are provided
|
|
if theme:
|
|
template_html = inject_theme(template_html, theme)
|
|
|
|
system_prompt = template_fill_html_system_prompt(constraints_text)
|
|
messages: list[models.ChatMessage] = [
|
|
models.ChatMessage(role="system", content=system_prompt),
|
|
models.ChatMessage(role="user", content=f"Document type: {document_type}"),
|
|
models.ChatMessage(
|
|
role="user",
|
|
content=(
|
|
"CRITICAL EXAMPLE — What to preserve vs. what to change:\n"
|
|
"If template has:\n"
|
|
" <div class='vendor-name'>{{VENDOR_NAME}}</div>\n"
|
|
"You MUST output:\n"
|
|
" <div class='vendor-name'>Acme Corporation</div> ← fill the token, keep the tag\n"
|
|
"\n"
|
|
"For multi-row content like {{LINEITEM_ROWS}}, generate full <tr>...</tr> HTML.\n"
|
|
"Match ONLY the columns present in the template's <thead> — do NOT add extra columns.\n"
|
|
"Example (4-column table: #, Description, Qty, Line Total):\n"
|
|
" <tr><td>1</td><td>Office chairs</td><td class='num'>4</td><td class='num'>$480.00</td></tr>\n"
|
|
"\n"
|
|
"Keep ALL HTML tags, CSS, classes, and IDs EXACTLY as-is.\n"
|
|
"Do NOT modify any <style> blocks, structural HTML, or class names."
|
|
),
|
|
),
|
|
models.ChatMessage(
|
|
role="user",
|
|
content=f"Outline/context (data to fill into template):\n{outline_text}",
|
|
),
|
|
]
|
|
if has_logo:
|
|
messages.append(
|
|
models.ChatMessage(
|
|
role="user",
|
|
content=(
|
|
"LOGO:\n"
|
|
"- A company logo is available.\n"
|
|
"- Insert the exact placeholder text {{LOGO_BLOCK}} once where the logo should appear.\n"
|
|
"- Do NOT create your own <img> for the logo; only place {{LOGO_BLOCK}}.\n"
|
|
),
|
|
)
|
|
)
|
|
if additional_instructions and additional_instructions.strip():
|
|
messages.append(
|
|
models.ChatMessage(
|
|
role="user",
|
|
content=(
|
|
"ADDITIONAL INSTRUCTIONS (highest priority):\n"
|
|
f"{additional_instructions.strip()}\n\n"
|
|
"Apply these instructions while preserving the template structure."
|
|
),
|
|
)
|
|
)
|
|
messages.append(
|
|
models.ChatMessage(
|
|
role="user",
|
|
content=(
|
|
f"COMPLETE HTML TEMPLATE (copy ALL of this, replacing ONLY {{{{PLACEHOLDER}}}} tokens):\n\n{template_html}"
|
|
),
|
|
)
|
|
)
|
|
if draft_sections:
|
|
sections_payload = [section.model_dump(by_alias=True, exclude_none=True) for section in draft_sections]
|
|
messages.append(
|
|
models.ChatMessage(
|
|
role="user",
|
|
content=f"Section content (JSON):\n{json.dumps(sections_payload, ensure_ascii=True)}",
|
|
)
|
|
)
|
|
|
|
stream = stream_ai(
|
|
SMART_MODEL,
|
|
messages,
|
|
tag="html_template_fill_stream",
|
|
log_label="html-template-fill-stream",
|
|
)
|
|
yield from stream
|
|
if stream.error or not stream.chunks:
|
|
raise RuntimeError(f"AI HTML template fill failed: {stream.error or 'no chunks streamed'}")
|
|
|
|
|
|
def generate_html_with_llm(
|
|
prompt: str,
|
|
history: list[models.ChatMessage],
|
|
document_type: str,
|
|
template_hint: str | None = None,
|
|
current_html: str | None = None,
|
|
structured_brief: str | None = None,
|
|
) -> str:
|
|
"""Call the LLM for freeform HTML document generation."""
|
|
messages: list[models.ChatMessage] = [
|
|
models.ChatMessage(role="system", content=html_system_prompt(document_type, template_hint))
|
|
]
|
|
messages.extend(history)
|
|
messages.extend(html_context_messages(template_hint, current_html, structured_brief))
|
|
messages.append(models.ChatMessage(role="user", content=prompt))
|
|
|
|
logger.info(
|
|
"[AI] HTML generation model=%s doc_type=%s template_hint=%s",
|
|
SMART_MODEL,
|
|
document_type,
|
|
"yes" if template_hint else "no",
|
|
)
|
|
parsed = run_ai(
|
|
SMART_MODEL,
|
|
messages,
|
|
models.HtmlResponse,
|
|
tag="html_generate",
|
|
)
|
|
if parsed and parsed.html and parsed.html.strip():
|
|
logger.info("[AI] Generated HTML preview (first 500 chars):\n%s", parsed.html[:500])
|
|
return parsed.html
|
|
|
|
raise RuntimeError("AI HTML generation failed or returned empty output.")
|
|
|
|
|
|
def generate_html_with_llm_stream(
|
|
prompt: str,
|
|
history: list[models.ChatMessage],
|
|
document_type: str,
|
|
template_hint: str | None = None,
|
|
current_html: str | None = None,
|
|
structured_brief: str | None = None,
|
|
):
|
|
"""Stream HTML generation from LLM, yielding chunks as they arrive."""
|
|
if not STREAMING_ENABLED:
|
|
full_html = generate_html_with_llm(
|
|
prompt,
|
|
history,
|
|
document_type,
|
|
template_hint,
|
|
current_html,
|
|
structured_brief,
|
|
)
|
|
yield full_html
|
|
return
|
|
|
|
messages: list[models.ChatMessage] = [
|
|
models.ChatMessage(role="system", content=html_system_prompt(document_type, template_hint))
|
|
]
|
|
messages.extend(history)
|
|
messages.extend(html_context_messages(template_hint, current_html, structured_brief))
|
|
messages.append(models.ChatMessage(role="user", content=prompt))
|
|
|
|
logger.info(
|
|
"[AI] Streaming HTML model=%s doc_type=%s template_hint=%s",
|
|
SMART_MODEL,
|
|
document_type,
|
|
"yes" if template_hint else "no",
|
|
)
|
|
stream = stream_ai(
|
|
SMART_MODEL,
|
|
messages,
|
|
tag="html_stream",
|
|
log_label="html-stream",
|
|
)
|
|
yield from stream
|
|
|
|
if stream.error or not stream.chunks:
|
|
raise RuntimeError(f"AI HTML streaming failed: {stream.error or 'no chunks produced'}")
|
|
|
|
|
|
def edit_html_with_llm_stream(
|
|
*,
|
|
document_type: str,
|
|
base_html: str,
|
|
instructions: str,
|
|
):
|
|
"""Stream an updated full HTML document by editing an existing HTML input."""
|
|
if not base_html or not base_html.strip():
|
|
raise RuntimeError("Missing base_html for HTML revision")
|
|
if not instructions or not instructions.strip():
|
|
raise RuntimeError("Missing instructions for HTML revision")
|
|
|
|
system_prompt = html_edit_system_prompt(document_type)
|
|
messages: list[models.ChatMessage] = [
|
|
models.ChatMessage(role="system", content=system_prompt),
|
|
models.ChatMessage(
|
|
role="user",
|
|
content=f"CURRENT_HTML (you MUST preserve everything not explicitly requested):\n\n{base_html}",
|
|
),
|
|
]
|
|
|
|
messages.append(
|
|
models.ChatMessage(
|
|
role="user",
|
|
content=f"INSTRUCTIONS:\n{instructions}\n\nReturn ONLY the full updated HTML document.",
|
|
)
|
|
)
|
|
|
|
stream = stream_ai(
|
|
SMART_MODEL,
|
|
messages,
|
|
tag="html_edit_stream",
|
|
log_label="html-edit-stream",
|
|
)
|
|
yield from stream
|
|
|
|
if stream.error or not stream.chunks:
|
|
raise RuntimeError(f"AI HTML edit streaming failed: {stream.error or 'no chunks produced'}")
|
|
|
|
|
|
__all__ = [
|
|
"generate_field_values",
|
|
"generate_html_with_llm",
|
|
"generate_html_with_llm_stream",
|
|
"generate_outline_with_llm",
|
|
"generate_section_draft",
|
|
"generate_template_fill_html_stream",
|
|
"edit_html_with_llm_stream",
|
|
]
|