Classifier Policy (#6898)

## Overview Adds **AI document classification** and a **classification-aware Files sidebar**: uploaded documents are automatically tagged with document-type labels (Invoice, Contract, Lab report, …), and the sidebar groups files under editable parent categories so a large library stays navigable. > [!IMPORTANT] > **This feature only runs in the SaaS build.** Classification depends on the AI engine and team-scoped label storage, so it's gated to SaaS end-to-end: > - The sidebar grouping is a `saas/`-layer override of the `fileSidebarGrouping` seam; every other build (OSS core, self-hosted proprietary, desktop) gets the null stub and renders the **unchanged flat, recency-sorted list** — no categories, no "Other", no picker. > - The classify/labels backend endpoints are gated on `policies.enabled` (on in SaaS) and live in `app/proprietary`, so they're absent from pure OSS and dormant in self-hosted unless explicitly enabled. > - The Python classifier is reached only via that gated path. > > Shared-layer changes that do compile everywhere are inert without the engine (dormant schema/field additions) or intentional (`GetInfoOnPDF` surfacing custom metadata). ## What it does - **Classifier (engine):** reads the first/last two pages of a PDF and assigns document-type labels from an allowed vocabulary. Labels are deliberately document-*type* descriptors — no deep-content/PII detection, since only a page window is read. - **Team label vocabulary:** ~270 built-in defaults across ~15 families, seeded per team. Editable by team leaders/admins in the Classification policy settings (import/export/reset). Team-scoped and shared; **per-user personal labels are intentionally out of scope** — the vocabulary is team-level only. - **Sidebar categories:** files group under parent categories (Financial, Legal, Medical, …), busiest-first, collapsible, with a "Recent" group on top and an "Other" group for anything uncategorised. The category structure (names, icons, membership, custom categories) is **device-local and user-editable** via a "Customize" picker — the only per-user personalization; it never changes the team's label vocabulary. - Classification results are written to PDF metadata (`StirlingPDFClassification`), read back to keep files in their groups without re-parsing. ## Architecture Spans all three layers, mirroring the existing policy/source subsystem conventions: - **`frontend/editor`** — sidebar grouping seam + SaaS override, category manager, labels editor, icon palette, file grouping, tests, `en-US` i18n. - **`app/proprietary` + `app/common` + `app/core`** — `ClassifyLabelController`, team-scoped `ClassificationLabelStore` (Jpa + in-process impls, same shape as `PolicyStore`/`SourceStore`), metadata read/write. - **`engine`** — the document-classifier agent, contracts, routes, tests. ## Screenshots **Files sidebar — grouped by category (SaaS)** ### Loading view <img width="2056" height="1046" alt="Screenshot 2026-07-07 at 5 12 56 PM" src="https://github.com/user-attachments/assets/1d712da5-50ae-4349-b0cd-e62665c3ec0c" /> ### Organized in the sidebar <img width="2056" height="1045" alt="Screenshot 2026-07-07 at 5 14 05 PM" src="https://github.com/user-attachments/assets/3ea4fe21-da51-4cea-bc3a-18ce040d3d05" /> **Customize categories picker** ### Personal settings to change how labels are grouped in an individual users editor <img width="2056" height="1044" alt="Screenshot 2026-07-07 at 5 52 42 PM" src="https://github.com/user-attachments/assets/40be03ce-0f63-4d1e-b58b-cec045d01cb2" /> **Classification labels editor (team settings)** <img width="2056" height="1042" alt="Screenshot 2026-07-07 at 5 53 00 PM" src="https://github.com/user-attachments/assets/337b0739-15c9-4749-9c6b-22e3b20825b8" /> ## Testing - Frontend `task frontend:check` — green (editor + portal tests, typecheck across all flavors, lint, label-drift guard). - Backend `task backend:check` (proprietary) and `:saas:test` — green. - Engine `task engine:check` — green.
2026-07-09 11:47:37 +00:00
parent f29500c138
commit c64369e56c
121 changed files with 10018 additions and 506 deletions
@@ -1,5 +1,6 @@
 """Agent modules for Stirling AI reasoning flows."""

+from .document_classifier import DocumentClassifierAgent
 from .execution import ExecutionPlanningAgent
 from .orchestrator import OrchestratorAgent
 from .pdf_create import PdfCreateAgent
@@ -9,6 +10,7 @@ from .pdf_review import PdfReviewAgent
 from .user_spec import UserSpecAgent

 __all__ = [
+    "DocumentClassifierAgent",
    "ExecutionPlanningAgent",
    "OrchestratorAgent",
    "PdfCreateAgent",
@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import logging
+
+from pydantic import Field
+from pydantic_ai import Agent
+from pydantic_ai.output import NativeOutput
+
+from stirling.contracts import (
+    ClassifyDocumentRequest,
+    ClassifyDocumentResponse,
+    DocumentClassificationResponse,
+    LabelOption,
+    PageText,
+)
+from stirling.models import ApiModel
+from stirling.services import AppRuntime
+
+logger = logging.getLogger(__name__)
+
+# A document carries at most this many labels; anything past the cap is dropped
+# in the model's order.
+MAX_ASSIGNED_LABELS = 5
+# Pages read from each end of the document. A document's type is evident from
+# its opening (and closing) pages, so a fixed window keeps cost and latency flat
+# regardless of length. Promote to AppSettings if it ever needs tuning.
+WINDOW_PAGES = 2
+
+
+_SYSTEM_PROMPT = (
+    "You identify what a document is by assigning labels, choosing only from a "
+    "fixed list of allowed labels you are given.\n"
+    "\n"
+    "Rules:\n"
+    f"- Pick up to {MAX_ASSIGNED_LABELS} labels that describe this document's type.\n"
+    "- Only use labels from the allowed list, spelled exactly as listed.\n"
+    "- Return an empty list if none fit.\n"
+    "- Judge from the document's content and structure, not from keywords alone. "
+    "The document may be in any language.\n"
+    "- You are shown only the first and last pages; that is enough to identify the type."
+)
+
+
+class _ClassifierOutput(ApiModel):
+    """Raw model answer, before it is validated against the allowed vocabulary."""
+
+    labels: list[str] = Field(
+        default_factory=list,
+        description="Labels from the allowed list that describe this document's type.",
+    )
+
+
+def render_labels(labels: list[LabelOption]) -> str:
+    """Render the allowed vocabulary for the prompt — the model reasons over the
+    human names (it never sees ids)."""
+    names = [label.name for label in labels]
+    return f"Allowed labels: {', '.join(names) or '(none)'}"
+
+
+def select_window(pages: list[PageText], window: int = WINDOW_PAGES) -> list[PageText]:
+    """Return the first and last ``window`` pages, never overlapping.
+
+    Documents short enough that the two ends would meet are returned whole. The
+    caller usually sends just the window already; this is a defensive trim in
+    case it sends more.
+    """
+    if window <= 0 or len(pages) <= window * 2:
+        return list(pages)
+    return [*pages[:window], *pages[-window:]]
+
+
+def format_window(pages: list[PageText]) -> str:
+    if not pages:
+        return "(no extractable text)"
+    return "\n\n".join(f"[Page {page.page_number}]\n{page.text}" for page in pages)
+
+
+def validate_labels(output: _ClassifierOutput, allowed: list[LabelOption]) -> DocumentClassificationResponse:
+    """Coerce a raw model answer (names) onto the allowed vocabulary, returning
+    label **ids**.
+
+    The model answers with names; they are matched case-insensitively to the
+    allowed vocabulary and returned as that label's id. Anything off-list is
+    dropped, duplicates collapse to the first occurrence, and the result is
+    capped at ``MAX_ASSIGNED_LABELS`` in the model's order. The model identifies;
+    these rules decide what is allowed to stand.
+    """
+    id_by_lower_name = {label.name.lower(): label.id for label in allowed}
+
+    kept: list[str] = []
+    for name in output.labels:
+        label_id = id_by_lower_name.get(name.strip().lower())
+        if label_id is not None and label_id not in kept:
+            kept.append(label_id)
+        if len(kept) == MAX_ASSIGNED_LABELS:
+            break
+
+    return DocumentClassificationResponse(labels=kept)
+
+
+class DocumentClassifierAgent:
+    """Assigns labels to a document from an allowed vocabulary.
+
+    Reads the bounded page window supplied on the request (first/last
+    ``WINDOW_PAGES``) and runs a single fast-model pass, then validates the
+    answer against the vocabulary so nothing off-list survives.
+    """
+
+    def __init__(self, runtime: AppRuntime) -> None:
+        self.runtime = runtime
+        self._agent = Agent(
+            model=runtime.fast_model,
+            output_type=NativeOutput(_ClassifierOutput),
+            system_prompt=_SYSTEM_PROMPT,
+            model_settings=runtime.fast_model_settings,
+        )
+
+    async def classify(self, request: ClassifyDocumentRequest) -> ClassifyDocumentResponse:
+        # The caller (the backend) always supplies the allowed vocabulary — the
+        # team's stored labels — so the engine holds no vocabulary of its own.
+        allowed = request.labels
+        window = select_window(request.pages)
+        prompt = self._build_prompt(request.file_name, allowed, window)
+        logger.debug("[classify] prompt:\n%s", prompt)
+        result = await self._agent.run(prompt)
+        return validate_labels(result.output, allowed)
+
+    @staticmethod
+    def _build_prompt(file_name: str, allowed: list[LabelOption], window: list[PageText]) -> str:
+        return (
+            f"{render_labels(allowed)}\n\n"
+            f"Document file name: {file_name}\n"
+            f"Document content (first and last pages):\n{format_window(window)}"
+        )
@@ -10,6 +10,7 @@ from pydantic_ai import Agent
 from pydantic_ai.models.instrumented import InstrumentationSettings

 from stirling.agents import (
+    DocumentClassifierAgent,
    ExecutionPlanningAgent,
    OrchestratorAgent,
    PdfEditAgent,
@@ -24,6 +25,7 @@ from stirling.api.middleware import UserIdMiddleware
 from stirling.api.routes import (
    agent_capabilities_router,
    agent_draft_router,
+    document_classifier_router,
    document_router,
    execution_router,
    ledger_router,
@@ -95,6 +97,7 @@ async def lifespan(fast_api: FastAPI):
    fast_api.state.execution_planning_agent = ExecutionPlanningAgent(runtime)
    fast_api.state.math_auditor_agent = MathAuditorAgent(runtime)
    fast_api.state.pdf_comment_agent = PdfCommentAgent(runtime)
+    fast_api.state.document_classifier_agent = DocumentClassifierAgent(runtime)
    tracer_provider = setup_posthog_tracking(settings)
    if tracer_provider:
        Agent.instrument_all(InstrumentationSettings(tracer_provider=tracer_provider))
@@ -131,6 +134,7 @@ app.include_router(document_router, dependencies=_user_gate)
 app.include_router(ledger_router, dependencies=_user_gate)
 app.include_router(pdf_comments_router, dependencies=_user_gate)
 app.include_router(agent_capabilities_router, dependencies=_user_gate)
+app.include_router(document_classifier_router, dependencies=_user_gate)


@app.get("/health", response_model=HealthResponse)
@@ -5,6 +5,7 @@ from typing import Annotated
 from fastapi import Depends, HTTPException, Request, status

 from stirling.agents import (
+    DocumentClassifierAgent,
    ExecutionPlanningAgent,
    OrchestratorAgent,
    PdfEditAgent,
@@ -55,6 +56,10 @@ def get_pdf_comment_agent(request: Request) -> PdfCommentAgent:
    return request.app.state.pdf_comment_agent


+def get_document_classifier_agent(request: Request) -> DocumentClassifierAgent:
+    return request.app.state.document_classifier_agent
+
+
 def require_user_id() -> UserId:
    """FastAPI dependency for routes that touch per-user storage.

@@ -1,5 +1,6 @@
 from .agent_capabilities import router as agent_capabilities_router
 from .agent_drafts import router as agent_draft_router
+from .document_classifier import router as document_classifier_router
 from .documents import router as document_router
 from .execution import router as execution_router
 from .ledger import router as ledger_router
@@ -11,6 +12,7 @@ from .pdf_questions import router as pdf_question_router
 __all__ = [
    "agent_capabilities_router",
    "agent_draft_router",
+    "document_classifier_router",
    "document_router",
    "execution_router",
    "ledger_router",
@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from typing import Annotated
+
+from fastapi import APIRouter, Depends
+
+from stirling.agents import DocumentClassifierAgent
+from stirling.api.dependencies import get_document_classifier_agent
+from stirling.contracts import ClassifyDocumentRequest, ClassifyDocumentResponse
+
+router = APIRouter(prefix="/api/v1/documents/classify", tags=["document-classifier"])
+
+
+@router.post("", response_model=ClassifyDocumentResponse)
+async def classify_document(
+    request: ClassifyDocumentRequest,
+    agent: Annotated[DocumentClassifierAgent, Depends(get_document_classifier_agent)],
+) -> ClassifyDocumentResponse:
+    """Classify a document from its supplied page text against the allowed labels.
+
+    The caller sends the bounded page window inline, so no per-user document
+    storage is touched here — the request is self-contained.
+    """
+    return await agent.classify(request)
@@ -35,6 +35,12 @@ from .contradiction import (
    ContradictionReport,
    ContradictionSeverity,
 )
+from .document_classifier import (
+    ClassifyDocumentRequest,
+    ClassifyDocumentResponse,
+    DocumentClassificationResponse,
+    LabelOption,
+)
 from .documents import (
    DeleteDocumentResponse,
    IngestDocumentRequest,
@@ -128,6 +134,8 @@ __all__ = [
    "AiToolAgentStep",
    "ArtifactKind",
    "CannotContinueExecutionAction",
+    "ClassifyDocumentRequest",
+    "ClassifyDocumentResponse",
    "Claim",
    "CommentSpec",
    "CompletedExecutionAction",
@@ -138,6 +146,8 @@ __all__ = [
    "DeleteDocumentResponse",
    "PurgeOwnerResponse",
    "Discrepancy",
+    "DocumentClassificationResponse",
+    "LabelOption",
    "DocumentMeta",
    "DocumentSections",
    "DiscrepancyKind",
@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from pydantic import Field
+
+from stirling.models import ApiModel
+
+from .documents import PageText
+
+
+class LabelOption(ApiModel):
+    """One allowed label: a stable ``id`` (the identity) plus the human ``name``
+    the model actually reasons over. The model picks names; the engine maps its
+    answer back to ids, so the result is language- and rename-stable."""
+
+    id: str = Field(min_length=1)
+    name: str = Field(min_length=1)
+
+
+class ClassifyDocumentRequest(ApiModel):
+    """Classify one document from its page text.
+
+    The caller sends the page text directly — typically just the bounded window
+    (first/last pages), since the classifier reads no more than that. There is no
+    ingestion or RAG step.
+
+    ``labels`` is the allowed vocabulary for this request — the caller (the
+    backend) always supplies it from the team's stored label set, each an
+    ``{id, name}`` pair. The engine holds no vocabulary of its own, so at least
+    one label is required.
+    """
+
+    file_name: str = Field(min_length=1)
+    pages: list[PageText] = Field(default_factory=list)
+    labels: list[LabelOption] = Field(min_length=1)
+
+
+class DocumentClassificationResponse(ApiModel):
+    """Terminal classification result.
+
+    ``labels`` is the subset of the allowed vocabulary the model assigned to the
+    document, as label **ids**: at most five entries, deduplicated, in the
+    model's order. An empty list is a valid answer — nothing in the vocabulary
+    fit. This is a plain answer from a dedicated endpoint — it carries no
+    ``outcome`` discriminator (it isn't one of the orchestrator's
+    WorkflowOutcome-routed union responses).
+    """
+
+    labels: list[str] = Field(default_factory=list)
+
+
+# Only one response shape today; kept as a named alias so routes and agents have
+# a stable response type to import.
+ClassifyDocumentResponse = DocumentClassificationResponse