Stirling-PDF/engine/src/stirling/agents/pdf_create/agent.py

"""PDF Create Agent — chunked multi-agent pipeline.

Flow:
  1. MetaPlannerAgent (smart_model) analyses the request and produces DocumentMeta:
     title, tone, shared terms, style, and cannot_do_reason. No sections yet.
  2. SectionPlannerAgent (smart_model) reads the meta and produces DocumentSections:
     ordered list of PlannedSection with heading, type, depth, and key_points.
  3. Python assembles DocumentPlan from meta + sections, then groups sections into
     chunks, each staying under the output-token ceiling.
  4. SectionWriterAgents (smart_model) run in parallel via asyncio.gather.
     Each returns a WrittenSections with fully populated DocumentSection objects.
  5. The assembler collects sections in plan order → GeneratedDocument.
  6. Jinja renders the document to HTML. The LLM never writes HTML.

The planner is split into two calls (meta then sections) so each LLM output schema
stays small enough for grammar compilation on all model tiers including Haiku.
"""

from __future__ import annotations

import asyncio
import logging
import re
from dataclasses import dataclass
from pathlib import Path

from jinja2 import Environment, FileSystemLoader
from pydantic_ai import Agent
from pydantic_ai.output import NativeOutput

from stirling.contracts import (
    EditCannotDoResponse,
    EditPlanResponse,
    OrchestratorRequest,
    ToolOperationStep,
    format_conversation_history,
)
from stirling.contracts.pdf_create import (
    DocumentMeta,
    DocumentPlan,
    DocumentSection,
    DocumentSections,
    GeneratedDocument,
    PdfCreateOrchestrateResponse,
    PlannedSection,
    SectionDepth,
    WrittenSections,
)
from stirling.models.agent_tool_models import AgentToolId, CreatePdfFromHtmlAgentParams
from stirling.services import AppRuntime

logger = logging.getLogger(__name__)

_TEMPLATES_DIR = Path(__file__).parent / "templates"

# ── Token budget ──────────────────────────────────────────────────────────────────────────────────

# Conservative per-section token estimates mapped from planner-assigned depth.
_DEPTH_TOKENS: dict[SectionDepth, int] = {
    SectionDepth.BRIEF: 250,
    SectionDepth.STANDARD: 550,
    SectionDepth.DETAILED: 1200,
}

# Maximum output tokens per writer call. Stays well below the quality cliff (~4k).
_CHUNK_CEILING = 3000

# Cap on simultaneous writer calls so a large document doesn't open a burst of LLM
# connections and trip provider rate limits.
_MAX_PARALLEL_WRITERS = 10

# ── Chunk dataclass ───────────────────────────────────────────────────────────────────────────────


@dataclass(frozen=True)
class _Chunk:
    index: int
    sections: list[PlannedSection]
    # Descriptions of neighbouring chunks from the plan — passed to writers as
    # read-only context so they can open/close their sections naturally.
    context_before: str | None
    context_after: str | None


# ── Chunking logic ────────────────────────────────────────────────────────────────────────────────


def _describe_sections(sections: list[PlannedSection]) -> str:
    """One-line summary of a chunk used as neighbour context."""
    return "; ".join(f'"{s.heading}" ({s.type.value})' for s in sections)


def _make_chunks(sections: list[PlannedSection]) -> list[_Chunk]:
    """Group planned sections into chunks, each under _CHUNK_CEILING output tokens.

    Section boundaries are atomic — a section is never split across chunks.
    A single section whose estimated cost exceeds the ceiling gets its own chunk.
    asyncio.gather preserves insertion order so chunk index is only used for logging.
    """
    if not sections:
        return []

    groups: list[list[PlannedSection]] = []
    current: list[PlannedSection] = []
    current_tokens = 0

    for section in sections:
        cost = _DEPTH_TOKENS[section.depth]
        if current and current_tokens + cost > _CHUNK_CEILING:
            groups.append(current)
            current = [section]
            current_tokens = cost
        else:
            current.append(section)
            current_tokens += cost

    if current:
        groups.append(current)

    chunks: list[_Chunk] = []
    for i, group in enumerate(groups):
        context_before = _describe_sections(groups[i - 1]) if i > 0 else None
        context_after = _describe_sections(groups[i + 1]) if i < len(groups) - 1 else None
        chunks.append(
            _Chunk(
                index=i,
                sections=group,
                context_before=context_before,
                context_after=context_after,
            )
        )

    return chunks


# ── Prompts ───────────────────────────────────────────────────────────────────────────────────────

_META_PLANNER_SYSTEM_PROMPT = """\
You are a document planner. Your job is Step 1 of 2: produce the document header — NOT the
section list (that comes in Step 2) and NOT any body text (section writers handle that).

Analyse the user's request and produce a DocumentMeta with:

- title, subtitle (if appropriate), reference_number (only if the user supplies one explicitly)

- tone_brief: one sentence describing register and style
  (e.g. "Formal legal language, third person, present tense." or
  "Professional business tone, active voice.")

- shared_terms: consistent names for key entities AND ground-truth facts used throughout
  the document. Two rules:
  1. Capture EVERY value the user states explicitly that could be referenced in more than
     one section. This includes — but is not limited to:
       · Named parties, organisations, products, or systems
       · Numeric values: amounts, quantities, percentages, durations, limits
       · Identifiers: version numbers, reference codes, model names
       · Dates and time periods
       · Units of measure or currency
  2. For any fact that will appear in two or more sections and that the user did NOT specify
     (e.g. a default time, a standard rate, a typical threshold), assign ONE specific value
     here. Do NOT let multiple writers independently invent the same fact.
  Examples: {"the Agreement": "this Non-Disclosure Agreement", "the Client": "Acme Corp",
             "contract value": "£120,000", "notice period": "30 days"}

- document_context: a single sentence anchoring the temporal or versioning context of the
  document, if the user provides one. Leave empty if the user provides no such context.

- style_primary_color: accent and heading colour. Set ONLY when the user explicitly names a
  colour or colour scheme (e.g. "make it red", "use navy blue"). Use CSS named colours
  (e.g. "magenta", "navy", "crimson") or hex values. Leave null if no colour is stated.
- style_background_color: page background colour. Set only if explicitly requested.
- style_body_text_color: body text colour. Set only if explicitly requested.

- cannot_do_reason: set this ONLY when the request is not asking to create a document at all
  (e.g. a question, a greeting, an edit request to an existing document). Never set it
  because the document is large, complex, or technically detailed. Leave null otherwise.

RULES:
1. Extract ALL information the user provides. Do not invent content.
2. Do not produce any sections — that is Step 2.
"""

_SECTIONS_PLANNER_SYSTEM_PROMPT = """\
You are a document planner. Your job is Step 2 of 2: produce the ordered section list for
a document whose header has already been decided. Do NOT write any body text.

You will be given:
  - The document meta (title, tone, shared terms, etc.) produced in Step 1
  - The original user request

Produce a DocumentSections with an ordered list of PlannedSection objects.

For each section choose:
  type — the most appropriate section type:
    text       — prose paragraphs (narrative, obligations, terms, descriptions)
    key_value  — labelled fields (parties, dates, metadata, identifiers)
    line_items — tables with column headers (expenses, schedules, item lists)
    bullet_list — unordered items (requirements, responsibilities, definitions)
    signature  — sign-off blocks for named parties or roles

  depth — honest estimate of content volume:
    brief    (~250 tokens) — 1-2 items, a short paragraph, or a small table
    standard (~550 tokens) — a few paragraphs, a medium table, or a moderate list
    detailed (~1200 tokens) — long clauses, complex multi-row tables, or dense content

  key_points — specific points this section MUST cover, taken directly from the user's input.
    These are instructions to the writer, not summaries. Be precise and complete.
    Every fact, name, date, amount, and requirement the user provides must appear somewhere.
    For large documents, include enough key_points that the writer can produce substantial
    content.

RULES:
1. Extract ALL information the user provides. Do not invent content.
2. Assign depth honestly — for a long detailed document most sections will be detailed.
3. For large documents, produce as many sections as needed — there is no section count limit.
4. Use the shared_terms from the meta exactly when writing key_points.
"""

_WRITER_SYSTEM_PROMPT = """\
You are a section writer for a structured document.
Write ONLY the sections assigned to you — no extras, no merging, no skipping.

SECTION TYPES — produce sections of exactly the requested type:
  text       — prose paragraphs. Use \\n\\n between paragraphs.
  key_value  — list of (label, value) pairs. Labels ≤ 5 words. Values verbatim from the data.
  line_items — table. Every row must have exactly as many cells as there are columns.
  bullet_list — flat list of items.
  signature  — list of signatory names/roles.

RULES:
1. Write ONLY the sections in your assignment list, in the order given.
2. Cover every key_point listed for each section. Do not omit any.
3. Use the shared_terms exactly — no paraphrasing or substituting alternatives.
   Shared terms are ground truth. If your general knowledge or a common default would
   produce a different value (e.g. a different duration, amount, date, or version number),
   the shared term takes precedence. This applies everywhere in the document, including
   boilerplate, FAQ, and summary sections.
4. Match the depth for each section: brief = concise, standard = moderate, \
detailed = thorough.
5. Maintain the document's tone throughout.
6. Do not reference other sections by number (e.g. "as defined in Section 3").
7. If a document_context is provided, use it to anchor any dates, versions, or time
   references you generate. Do not invent a different temporal or versioning context.
"""


def _build_sections_prompt(meta: DocumentMeta, user_request: str, history: str) -> str:
    lines: list[str] = [
        "Document meta from Step 1:",
        f"  Title: {meta.title}",
        f"  Tone: {meta.tone_brief}",
    ]
    if meta.subtitle:
        lines.append(f"  Subtitle: {meta.subtitle}")
    if meta.document_context:
        lines.append(f"  Document context: {meta.document_context}")
    if meta.shared_terms:
        lines.append("  Shared terms:")
        for term, referent in meta.shared_terms.items():
            lines.append(f"    {term} → {referent}")

    lines.append(f"\nConversation history:\n{history}")
    lines.append(f"\nUser request: {user_request}")
    return "\n".join(lines)


def _build_writer_prompt(plan: DocumentPlan, chunk: _Chunk) -> str:
    lines: list[str] = [
        f"Document: {plan.title}",
        f"Tone: {plan.tone_brief}",
    ]

    if plan.document_context:
        lines.append(f"Document context: {plan.document_context}")

    if plan.shared_terms:
        lines.append("Ground-truth facts and shared terms (use exactly — these override defaults):")
        for term, referent in plan.shared_terms.items():
            lines.append(f"  {term} → {referent}")

    if chunk.context_before:
        lines.append(f"\nThe sections BEFORE yours cover: {chunk.context_before}")
    if chunk.context_after:
        lines.append(f"The sections AFTER yours cover: {chunk.context_after}")

    lines.append(f"\nWrite these {len(chunk.sections)} section(s) in order:")
    for i, s in enumerate(chunk.sections, 1):
        lines.append(f"\n--- Section {i} ---")
        lines.append(f"Heading: {s.heading}")
        lines.append(f"Type: {s.type.value}")
        lines.append(f"Depth: {s.depth.value}")
        lines.append("Key points to cover:")
        for point in s.key_points:
            lines.append(f"  - {point}")

    return "\n".join(lines)


# ── Helpers ───────────────────────────────────────────────────────────────────────────────────────


def _build_jinja_env() -> Environment:
    return Environment(
        loader=FileSystemLoader(str(_TEMPLATES_DIR)),
        autoescape=True,
        trim_blocks=True,
        lstrip_blocks=True,
    )


def _safe_filename(title: str) -> str:
    slug = re.sub(r"[^\w\s-]", "", title.lower())
    slug = re.sub(r"[\s_-]+", "-", slug).strip("-")
    return (slug[:60] or "document") + ".pdf"


# ── Agent ─────────────────────────────────────────────────────────────────────────────────────────


class PdfCreateAgent:
    def __init__(self, runtime: AppRuntime) -> None:
        self.runtime = runtime
        self._jinja_env = _build_jinja_env()

        self._meta_planner: Agent[None, DocumentMeta] = Agent(
            model=runtime.smart_model,
            output_type=NativeOutput(DocumentMeta),
            system_prompt=_META_PLANNER_SYSTEM_PROMPT,
            model_settings={**runtime.smart_model_settings, "temperature": 0.1},
        )

        self._sections_planner: Agent[None, DocumentSections] = Agent(
            model=runtime.smart_model,
            output_type=NativeOutput(DocumentSections),
            system_prompt=_SECTIONS_PLANNER_SYSTEM_PROMPT,
            model_settings={**runtime.smart_model_settings, "temperature": 0.1},
        )

        self._writer: Agent[None, WrittenSections] = Agent(
            model=runtime.smart_model,
            output_type=NativeOutput(WrittenSections),
            system_prompt=_WRITER_SYSTEM_PROMPT,
            model_settings={**runtime.smart_model_settings, "temperature": 0.3},
        )

    async def orchestrate(self, request: OrchestratorRequest) -> PdfCreateOrchestrateResponse:
        history = format_conversation_history(request.conversation_history)

        # ── Phase 1: plan meta ─────────────────────────────────────────────────
        logger.info("[pdf-create] phase 1/6: planning document meta")
        meta_prompt = f"Conversation history:\n{history}\n\nUser request: {request.user_message}"
        meta_result = await self._meta_planner.run(meta_prompt)
        meta = meta_result.output

        if meta.cannot_do_reason:
            logger.info("[pdf-create] cannot_do: %s", meta.cannot_do_reason)
            return EditCannotDoResponse(reason=meta.cannot_do_reason)

        logger.info("[pdf-create] meta: title=%r tone=%r", meta.title, meta.tone_brief)

        # ── Phase 2: plan sections ─────────────────────────────────────────────
        logger.info("[pdf-create] phase 2/6: planning sections")
        sections_prompt = _build_sections_prompt(meta, request.user_message, history)
        sections_result = await self._sections_planner.run(sections_prompt)
        planned_sections = sections_result.output

        if not planned_sections.sections:
            logger.info("[pdf-create] sections planner returned empty sections")
            return EditCannotDoResponse(reason="No document sections could be planned from the request.")

        plan = DocumentPlan.assemble(meta, planned_sections)

        # ── Phase 3: chunk ─────────────────────────────────────────────────────
        chunks = _make_chunks(plan.sections)
        logger.info(
            "[pdf-create] phase 3/6: chunked — sections=%d chunks=%d",
            len(plan.sections),
            len(chunks),
        )

        # ── Phase 4: write in parallel, bounded ────────────────────────────────
        logger.info("[pdf-create] phase 4/6: writing %d chunk(s) in parallel", len(chunks))
        total_chunks = len(chunks)
        semaphore = asyncio.Semaphore(_MAX_PARALLEL_WRITERS)
        written_chunks: list[WrittenSections] = await asyncio.gather(
            *[self._write_chunk(plan, chunk, total_chunks, semaphore) for chunk in chunks]
        )

        # ── Phase 5: assemble in plan order (gather preserves insertion order) ──
        all_sections: list[DocumentSection] = []
        for written in written_chunks:
            all_sections.extend(written.sections)

        logger.info("[pdf-create] phase 5/6: assembled %d sections", len(all_sections))

        doc = GeneratedDocument(
            title=plan.title,
            subtitle=plan.subtitle,
            reference_number=plan.reference_number,
            style=plan.style,
            sections=all_sections,
        )

        # ── Phase 6: render ────────────────────────────────────────────────────
        logger.info("[pdf-create] phase 6/6: rendering HTML")
        html = self._render(doc)
        filename = _safe_filename(plan.title)
        logger.info(
            "[pdf-create] done — filename=%r html_bytes=%d",
            filename,
            len(html),
        )

        return EditPlanResponse(
            summary=f"Created {plan.title}",
            steps=[
                ToolOperationStep(
                    tool=AgentToolId.CREATE_PDF_FROM_HTML_AGENT,
                    parameters=CreatePdfFromHtmlAgentParams(
                        html_content=html,
                        filename=filename,
                    ),
                )
            ],
        )

    async def _write_chunk(
        self, plan: DocumentPlan, chunk: _Chunk, total_chunks: int, semaphore: asyncio.Semaphore
    ) -> WrittenSections:
        async with semaphore:
            prompt = _build_writer_prompt(plan, chunk)
            result = await self._writer.run(prompt)
        logger.info(
            "[pdf-create] chunk %d/%d wrote %d sections",
            chunk.index + 1,
            total_chunks,
            len(result.output.sections),
        )
        return result.output

    def _render(self, doc: GeneratedDocument) -> str:
        template = self._jinja_env.get_template("document.html.jinja2")
        return template.render(doc=doc)