diff --git a/.taskfiles/frontend.yml b/.taskfiles/frontend.yml index bcc7c07362..aa4ed9f674 100644 --- a/.taskfiles/frontend.yml +++ b/.taskfiles/frontend.yml @@ -55,6 +55,15 @@ tasks: - editor/src/core/data/ogImageMap.json - editor/public/og-metadata.json + prepare:taxonomy: + internal: true + run: when_changed + desc: "Regenerate the engine taxonomy JSON from the TS source of truth" + cmds: + - npx tsx editor/scripts/generate-taxonomy.mts + sources: + - editor/src/proprietary/data/classificationTaxonomy.ts + prepare: desc: "Set up dev environment" run: when_changed @@ -65,6 +74,7 @@ tasks: vars: { MODE: '{{.MODE}}' } - prepare:icons - prepare:og + - prepare:taxonomy # ============================================================ # Development @@ -384,12 +394,23 @@ tasks: cmds: - node editor/scripts/generate-og-metadata.mjs --check + taxonomy: + desc: "Regenerate the engine classification taxonomy JSON from the TS source" + cmds: + - npx tsx editor/scripts/generate-taxonomy.mts + + taxonomy:check: + desc: "Fail if the committed classification taxonomy JSON is out of date" + cmds: + - npx tsx editor/scripts/generate-taxonomy.mts --check + check:all: desc: "Full CI quality gate" cmds: # Runs first, before prepare regenerates: guards the committed og-metadata.json / # ogImageMap.json that the Cloudflare Pages (plain `vite build`) deploy relies on. - task: og:check + - task: taxonomy:check - task: typecheck:all - task: lint - task: format:check diff --git a/app/common/src/main/java/stirling/software/common/service/PdfMetadataService.java b/app/common/src/main/java/stirling/software/common/service/PdfMetadataService.java index 7b74111870..30b19ccfd7 100644 --- a/app/common/src/main/java/stirling/software/common/service/PdfMetadataService.java +++ b/app/common/src/main/java/stirling/software/common/service/PdfMetadataService.java @@ -7,6 +7,7 @@ import java.time.format.DateTimeFormatter; import java.util.Calendar; import org.apache.pdfbox.pdmodel.PDDocument; +import org.apache.pdfbox.pdmodel.PDDocumentInformation; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.stereotype.Service; @@ -17,6 +18,11 @@ import stirling.software.common.model.PdfMetadata; @Service public class PdfMetadataService { + /** + * ({@code {category, docType, typeConfidence, tags}}). Written by the classify-and-tag tool. + */ + public static final String CLASSIFICATION_KEY = "StirlingPDFClassification"; + private final ApplicationProperties applicationProperties; private final String stirlingPDFLabel; private final UserServiceInterface userService; @@ -177,4 +183,14 @@ public class PdfMetadataService { } pdf.getDocumentInformation().setAuthor(author); } + + /** + * Write the document classifier's JSON result into the custom Info-dictionary field {@link + * #CLASSIFICATION_KEY}, leaving all other metadata untouched. + */ + public void setClassificationMetadata(PDDocument pdf, String classificationJson) { + PDDocumentInformation info = pdf.getDocumentInformation(); + info.setCustomMetadataValue(CLASSIFICATION_KEY, classificationJson); + pdf.setDocumentInformation(info); + } } diff --git a/app/core/src/main/java/stirling/software/SPDF/controller/api/security/GetInfoOnPDF.java b/app/core/src/main/java/stirling/software/SPDF/controller/api/security/GetInfoOnPDF.java index f61a1ce8fe..a66b168163 100644 --- a/app/core/src/main/java/stirling/software/SPDF/controller/api/security/GetInfoOnPDF.java +++ b/app/core/src/main/java/stirling/software/SPDF/controller/api/security/GetInfoOnPDF.java @@ -305,6 +305,23 @@ public class GetInfoOnPDF { } } + /** + * Info-dictionary keys exposed above via typed getters; any other key in the dictionary is + * surfaced as custom metadata (e.g. the classification policy's StirlingPDFClassification + * entry). + */ + private static final java.util.Set STANDARD_INFO_KEYS = + java.util.Set.of( + "Title", + "Author", + "Subject", + "Keywords", + "Producer", + "Creator", + "CreationDate", + "ModDate", + "Trapped"); + private static ObjectNode extractMetadata(PDDocument document) { ObjectNode metadata = objectMapper.createObjectNode(); @@ -335,6 +352,18 @@ public class GetInfoOnPDF { if (modificationDate != null) { metadata.put("ModificationDate", modificationDate); } + + // Surface custom Info-dictionary entries (anything beyond the + // standard fields above) — e.g. StirlingPDFClassification + for (String key : info.getMetadataKeys()) { + if (STANDARD_INFO_KEYS.contains(key)) { + continue; + } + String value = info.getCustomMetadataValue(key); + if (value != null && !value.isBlank()) { + metadata.put(key, value); + } + } } } catch (Exception e) { log.error("Error extracting metadata: {}", e.getMessage()); diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/ClassifyTagController.java b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/ClassifyTagController.java new file mode 100644 index 0000000000..13adb03809 --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/ClassifyTagController.java @@ -0,0 +1,166 @@ +package stirling.software.proprietary.controller.api; + +import com.fasterxml.jackson.annotation.JsonInclude; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; + +import org.apache.pdfbox.pdmodel.PDDocument; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.core.io.Resource; +import org.springframework.http.MediaType; +import org.springframework.http.ResponseEntity; +import org.springframework.web.bind.annotation.PostMapping; +import org.springframework.web.bind.annotation.RequestMapping; +import org.springframework.web.bind.annotation.RequestParam; +import org.springframework.web.bind.annotation.RestController; +import org.springframework.web.multipart.MultipartFile; + +import io.github.pixee.security.Filenames; +import io.swagger.v3.oas.annotations.Hidden; +import io.swagger.v3.oas.annotations.Operation; +import io.swagger.v3.oas.annotations.tags.Tag; + +import lombok.extern.slf4j.Slf4j; + +import stirling.software.common.service.CustomPDFDocumentFactory; +import stirling.software.common.service.PdfMetadataService; +import stirling.software.common.service.UserServiceInterface; +import stirling.software.common.util.TempFileManager; +import stirling.software.common.util.WebResponseUtils; +import stirling.software.proprietary.model.api.ai.AiPageText; +import stirling.software.proprietary.service.AiEngineClient; +import stirling.software.proprietary.service.PdfContentExtractor; + +import tools.jackson.databind.JsonNode; +import tools.jackson.databind.ObjectMapper; +import tools.jackson.databind.node.ObjectNode; + +/** + * Dispatchable tool that classifies a PDF and writes the result into its metadata. + * + *

Runs as a Classification-policy pipeline step: it reads a bounded page window, asks the AI + * engine to classify the document, and stores the engine's JSON answer — minus the transport-only + * {@code outcome} field — in the custom Info-dictionary key {@link + * PdfMetadataService#CLASSIFICATION_KEY}. Returns the tagged PDF. Not intended for direct client + * use. + */ +@Slf4j +@Hidden +@RestController +@RequestMapping("/api/v1/ai/tools") +@Tag(name = "AI Tools", description = "Dispatchable AI-backed tools.") +public class ClassifyTagController { + + /** Pages read from each end of the document — mirrors the engine's window. */ + private static final int WINDOW_PAGES = 2; + + private static final String CLASSIFY_ENDPOINT = "/api/v1/documents/classify"; + + private final CustomPDFDocumentFactory pdfDocumentFactory; + private final TempFileManager tempFileManager; + private final PdfContentExtractor pdfContentExtractor; + private final PdfMetadataService pdfMetadataService; + private final AiEngineClient aiEngineClient; + private final ObjectMapper objectMapper; + private final UserServiceInterface userService; + + public ClassifyTagController( + CustomPDFDocumentFactory pdfDocumentFactory, + TempFileManager tempFileManager, + PdfContentExtractor pdfContentExtractor, + PdfMetadataService pdfMetadataService, + AiEngineClient aiEngineClient, + ObjectMapper objectMapper, + @Autowired(required = false) UserServiceInterface userService) { + this.pdfDocumentFactory = pdfDocumentFactory; + this.tempFileManager = tempFileManager; + this.pdfContentExtractor = pdfContentExtractor; + this.pdfMetadataService = pdfMetadataService; + this.aiEngineClient = aiEngineClient; + this.objectMapper = objectMapper; + this.userService = userService; + } + + @PostMapping(value = "/classify-and-tag", consumes = MediaType.MULTIPART_FORM_DATA_VALUE) + @Operation( + summary = "Classify a PDF and tag its metadata", + description = + "Reads the first two and last two pages, classifies the document via the AI" + + " engine, and stores the result in the StirlingPDFClassification" + + " metadata field. Dispatched by the Classification policy; not" + + " intended for direct client use.") + public ResponseEntity classifyAndTag( + @RequestParam("fileInput") MultipartFile fileInput) throws IOException { + try (PDDocument document = pdfDocumentFactory.load(fileInput, true)) { + String fileName = safeFileName(fileInput.getOriginalFilename()); + + List pages = extractWindow(document); + String requestBody = + objectMapper.writeValueAsString( + new ClassifyEngineRequest(fileName, pages, resolveTaxonomyOverride())); + + String userId = userService != null ? userService.getCurrentUsername() : null; + String responseJson = aiEngineClient.post(CLASSIFY_ENDPOINT, requestBody, userId); + + pdfMetadataService.setClassificationMetadata(document, toMetadataValue(responseJson)); + log.debug("[classify-and-tag] tagged {} ({} window pages)", fileName, pages.size()); + + return WebResponseUtils.pdfDocToWebResponse(document, fileName, tempFileManager); + } + } + + private List extractWindow(PDDocument document) throws IOException { + List pages = new ArrayList<>(); + for (int pageNumber : windowPageNumbers(document.getNumberOfPages(), WINDOW_PAGES)) { + String text = pdfContentExtractor.extractPageTextRaw(document, pageNumber); + if (text != null && !text.isBlank()) { + pages.add(new AiPageText(pageNumber, text)); + } + } + return pages; + } + + /** First and last {@code window} page numbers (1-based), de-duplicated and in order. */ + static List windowPageNumbers(int pageCount, int window) { + Set numbers = new LinkedHashSet<>(); + for (int page = 1; page <= Math.min(window, pageCount); page++) { + numbers.add(page); + } + for (int page = Math.max(1, pageCount - window + 1); page <= pageCount; page++) { + numbers.add(page); + } + return new ArrayList<>(numbers); + } + + /** Drop the transport-only {@code outcome} discriminator; keep the rest verbatim. */ + private String toMetadataValue(String engineResponseJson) { + JsonNode node = objectMapper.readTree(engineResponseJson); + if (node instanceof ObjectNode object) { + object.remove("outcome"); + } + return objectMapper.writeValueAsString(node); + } + + private static String safeFileName(String originalFilename) { + String name = Filenames.toSimpleFileName(originalFilename); + return (name == null || name.isBlank()) ? "classified.pdf" : name; + } + + /** + * Override point for a future per-org / DB-configured taxonomy: resolve the caller's taxonomy + * here and return it (engine shape) to classify against; {@code null} falls back to the engine's + * generated default. Always null today. + */ + private JsonNode resolveTaxonomyOverride() { + return null; + } + + /** Request body for the engine's {@code /api/v1/documents/classify} endpoint. */ + @JsonInclude(JsonInclude.Include.NON_NULL) + private record ClassifyEngineRequest( + String fileName, List pages, JsonNode taxonomy) {} +} diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/controller/api/ClassifyTagControllerTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/controller/api/ClassifyTagControllerTest.java new file mode 100644 index 0000000000..0ad76283de --- /dev/null +++ b/app/proprietary/src/test/java/stirling/software/proprietary/controller/api/ClassifyTagControllerTest.java @@ -0,0 +1,103 @@ +package stirling.software.proprietary.controller.api; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.List; + +import org.apache.pdfbox.pdmodel.PDDocument; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.mockito.junit.jupiter.MockitoSettings; +import org.mockito.quality.Strictness; +import org.springframework.web.multipart.MultipartFile; + +import stirling.software.common.service.CustomPDFDocumentFactory; +import stirling.software.common.service.PdfMetadataService; +import stirling.software.common.util.TempFileManager; +import stirling.software.proprietary.service.AiEngineClient; +import stirling.software.proprietary.service.PdfContentExtractor; + +import tools.jackson.databind.JsonNode; +import tools.jackson.databind.ObjectMapper; +import tools.jackson.databind.json.JsonMapper; + +@ExtendWith(MockitoExtension.class) +@MockitoSettings(strictness = Strictness.LENIENT) +class ClassifyTagControllerTest { + + @Mock private CustomPDFDocumentFactory pdfDocumentFactory; + @Mock private TempFileManager tempFileManager; + @Mock private PdfContentExtractor pdfContentExtractor; + @Mock private PdfMetadataService pdfMetadataService; + @Mock private AiEngineClient aiEngineClient; + + private final ObjectMapper objectMapper = JsonMapper.builder().build(); + private ClassifyTagController controller; + + @BeforeEach + void setUp() { + controller = + new ClassifyTagController( + pdfDocumentFactory, + tempFileManager, + pdfContentExtractor, + pdfMetadataService, + aiEngineClient, + objectMapper, + null); + } + + @Test + void classifyAndTag_writesClassificationWithoutOutcome() throws Exception { + PDDocument document = mock(PDDocument.class); + MultipartFile file = mock(MultipartFile.class); + when(file.getOriginalFilename()).thenReturn("invoice.pdf"); + when(pdfDocumentFactory.load(any(MultipartFile.class), eq(true))).thenReturn(document); + when(document.getNumberOfPages()).thenReturn(1); + when(pdfContentExtractor.extractPageTextRaw(document, 1)) + .thenReturn("Invoice total due 100.00"); + when(aiEngineClient.post(eq("/api/v1/documents/classify"), anyString(), isNull())) + .thenReturn( + "{\"outcome\":\"classification\",\"category\":\"invoice\"," + + "\"docType\":\"invoice\",\"typeConfidence\":0.98," + + "\"tags\":[\"finance\"]}"); + + try { + controller.classifyAndTag(file); + } catch (Exception ignored) { + // WebResponseUtils.pdfDocToWebResponse needs a real temp file; the metadata write we + // assert on has already happened by the time it runs. + } + + ArgumentCaptor value = ArgumentCaptor.forClass(String.class); + verify(pdfMetadataService).setClassificationMetadata(eq(document), value.capture()); + + JsonNode written = objectMapper.readTree(value.getValue()); + assertThat(written.has("outcome")).isFalse(); + assertThat(written.get("category").asText()).isEqualTo("invoice"); + assertThat(written.get("docType").asText()).isEqualTo("invoice"); + assertThat(written.get("tags").get(0).asText()).isEqualTo("finance"); + } + + @Test + void windowPageNumbers_takesFirstAndLastWithoutOverlap() { + assertEquals(List.of(1, 2, 4, 5), ClassifyTagController.windowPageNumbers(5, 2)); + assertEquals(List.of(1, 2, 3), ClassifyTagController.windowPageNumbers(3, 2)); + // Short docs clamp + dedupe rather than throwing or going out of range. + assertEquals(List.of(1, 2), ClassifyTagController.windowPageNumbers(2, 2)); + assertEquals(List.of(1), ClassifyTagController.windowPageNumbers(1, 2)); + assertEquals(List.of(), ClassifyTagController.windowPageNumbers(0, 2)); + } +} diff --git a/engine/src/stirling/agents/__init__.py b/engine/src/stirling/agents/__init__.py index c22bd6c977..7bc2cb63a2 100644 --- a/engine/src/stirling/agents/__init__.py +++ b/engine/src/stirling/agents/__init__.py @@ -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", diff --git a/engine/src/stirling/agents/default_taxonomy.generated.json b/engine/src/stirling/agents/default_taxonomy.generated.json new file mode 100644 index 0000000000..d44c2c40e5 --- /dev/null +++ b/engine/src/stirling/agents/default_taxonomy.generated.json @@ -0,0 +1,568 @@ +{ + "categories": [ + { + "id": "invoice", + "label": "Invoice", + "docTypes": [ + { + "id": "invoice", + "label": "Invoice" + }, + { + "id": "receipt", + "label": "Receipt" + }, + { + "id": "credit_note", + "label": "Credit note" + }, + { + "id": "purchase_order", + "label": "Purchase order" + }, + { + "id": "quote", + "label": "Quote" + } + ] + }, + { + "id": "contract", + "label": "Contract", + "docTypes": [ + { + "id": "nda", + "label": "Non-disclosure agreement" + }, + { + "id": "employment_agreement", + "label": "Employment agreement" + }, + { + "id": "service_agreement", + "label": "Service agreement" + }, + { + "id": "lease_agreement", + "label": "Lease agreement" + }, + { + "id": "master_service_agreement", + "label": "Master service agreement" + }, + { + "id": "statement_of_work", + "label": "Statement of work" + }, + { + "id": "terms_of_service", + "label": "Terms of service" + } + ] + }, + { + "id": "financial_statement", + "label": "Financial statement", + "docTypes": [ + { + "id": "balance_sheet", + "label": "Balance sheet" + }, + { + "id": "income_statement", + "label": "Income statement" + }, + { + "id": "cash_flow_statement", + "label": "Cash flow statement" + }, + { + "id": "bank_statement", + "label": "Bank statement" + }, + { + "id": "annual_report", + "label": "Annual report" + } + ] + }, + { + "id": "report", + "label": "Report", + "docTypes": [ + { + "id": "business_report", + "label": "Business report" + }, + { + "id": "project_report", + "label": "Project report" + }, + { + "id": "research_report", + "label": "Research report" + }, + { + "id": "status_report", + "label": "Status report" + }, + { + "id": "incident_report", + "label": "Incident report" + } + ] + }, + { + "id": "letter", + "label": "Letter", + "docTypes": [ + { + "id": "business_letter", + "label": "Business letter" + }, + { + "id": "cover_letter", + "label": "Cover letter" + }, + { + "id": "recommendation_letter", + "label": "Recommendation letter" + }, + { + "id": "complaint_letter", + "label": "Complaint letter" + }, + { + "id": "demand_letter", + "label": "Demand letter" + } + ] + }, + { + "id": "form", + "label": "Form", + "docTypes": [ + { + "id": "application_form", + "label": "Application form" + }, + { + "id": "registration_form", + "label": "Registration form" + }, + { + "id": "consent_form", + "label": "Consent form" + }, + { + "id": "survey", + "label": "Survey" + }, + { + "id": "questionnaire", + "label": "Questionnaire" + } + ] + }, + { + "id": "resume", + "label": "Resume", + "docTypes": [ + { + "id": "resume", + "label": "Resume" + }, + { + "id": "curriculum_vitae", + "label": "Curriculum vitae" + }, + { + "id": "portfolio", + "label": "Portfolio" + }, + { + "id": "reference_sheet", + "label": "Reference sheet" + } + ] + }, + { + "id": "tax_form", + "label": "Tax form", + "docTypes": [ + { + "id": "tax_return", + "label": "Tax return" + }, + { + "id": "w2", + "label": "W-2" + }, + { + "id": "w9", + "label": "W-9" + }, + { + "id": "form_1099", + "label": "Form 1099" + }, + { + "id": "vat_return", + "label": "VAT return" + } + ] + }, + { + "id": "expense_report", + "label": "Expense report", + "docTypes": [ + { + "id": "expense_report", + "label": "Expense report" + }, + { + "id": "reimbursement_request", + "label": "Reimbursement request" + }, + { + "id": "mileage_log", + "label": "Mileage log" + }, + { + "id": "per_diem_claim", + "label": "Per diem claim" + } + ] + }, + { + "id": "presentation", + "label": "Presentation", + "docTypes": [ + { + "id": "slide_deck", + "label": "Slide deck" + }, + { + "id": "pitch_deck", + "label": "Pitch deck" + }, + { + "id": "training_deck", + "label": "Training deck" + }, + { + "id": "webinar_deck", + "label": "Webinar deck" + } + ] + }, + { + "id": "medical_record", + "label": "Medical record", + "docTypes": [ + { + "id": "lab_result", + "label": "Lab result" + }, + { + "id": "prescription", + "label": "Prescription" + }, + { + "id": "discharge_summary", + "label": "Discharge summary" + }, + { + "id": "medical_history", + "label": "Medical history" + }, + { + "id": "imaging_report", + "label": "Imaging report" + }, + { + "id": "vaccination_record", + "label": "Vaccination record" + } + ] + }, + { + "id": "legal_filing", + "label": "Legal filing", + "docTypes": [ + { + "id": "court_filing", + "label": "Court filing" + }, + { + "id": "complaint", + "label": "Complaint" + }, + { + "id": "motion", + "label": "Motion" + }, + { + "id": "subpoena", + "label": "Subpoena" + }, + { + "id": "affidavit", + "label": "Affidavit" + }, + { + "id": "deposition", + "label": "Deposition" + } + ] + }, + { + "id": "identity_document", + "label": "Identity document", + "docTypes": [ + { + "id": "passport", + "label": "Passport" + }, + { + "id": "drivers_license", + "label": "Driver's license" + }, + { + "id": "national_id", + "label": "National ID" + }, + { + "id": "birth_certificate", + "label": "Birth certificate" + }, + { + "id": "visa", + "label": "Visa" + } + ] + }, + { + "id": "insurance", + "label": "Insurance", + "docTypes": [ + { + "id": "insurance_policy", + "label": "Insurance policy" + }, + { + "id": "insurance_claim", + "label": "Insurance claim" + }, + { + "id": "certificate_of_insurance", + "label": "Certificate of insurance" + }, + { + "id": "explanation_of_benefits", + "label": "Explanation of benefits" + } + ] + }, + { + "id": "real_estate", + "label": "Real estate", + "docTypes": [ + { + "id": "deed", + "label": "Deed" + }, + { + "id": "mortgage_agreement", + "label": "Mortgage agreement" + }, + { + "id": "property_appraisal", + "label": "Property appraisal" + }, + { + "id": "closing_disclosure", + "label": "Closing disclosure" + }, + { + "id": "title_report", + "label": "Title report" + } + ] + }, + { + "id": "shipping", + "label": "Shipping", + "docTypes": [ + { + "id": "bill_of_lading", + "label": "Bill of lading" + }, + { + "id": "packing_slip", + "label": "Packing slip" + }, + { + "id": "customs_declaration", + "label": "Customs declaration" + }, + { + "id": "delivery_note", + "label": "Delivery note" + }, + { + "id": "air_waybill", + "label": "Air waybill" + } + ] + }, + { + "id": "hr_document", + "label": "HR document", + "docTypes": [ + { + "id": "offer_letter", + "label": "Offer letter" + }, + { + "id": "performance_review", + "label": "Performance review" + }, + { + "id": "payslip", + "label": "Payslip" + }, + { + "id": "employee_handbook", + "label": "Employee handbook" + }, + { + "id": "termination_letter", + "label": "Termination letter" + }, + { + "id": "timesheet", + "label": "Timesheet" + } + ] + }, + { + "id": "academic_record", + "label": "Academic record", + "docTypes": [ + { + "id": "transcript", + "label": "Transcript" + }, + { + "id": "diploma", + "label": "Diploma" + }, + { + "id": "certificate", + "label": "Certificate" + }, + { + "id": "syllabus", + "label": "Syllabus" + }, + { + "id": "thesis", + "label": "Thesis" + }, + { + "id": "report_card", + "label": "Report card" + } + ] + }, + { + "id": "marketing_material", + "label": "Marketing material", + "docTypes": [ + { + "id": "brochure", + "label": "Brochure" + }, + { + "id": "flyer", + "label": "Flyer" + }, + { + "id": "case_study", + "label": "Case study" + }, + { + "id": "white_paper", + "label": "White paper" + }, + { + "id": "press_release", + "label": "Press release" + } + ] + }, + { + "id": "technical_document", + "label": "Technical document", + "docTypes": [ + { + "id": "user_manual", + "label": "User manual" + }, + { + "id": "specification", + "label": "Specification" + }, + { + "id": "api_documentation", + "label": "API documentation" + }, + { + "id": "installation_guide", + "label": "Installation guide" + }, + { + "id": "datasheet", + "label": "Datasheet" + }, + { + "id": "release_notes", + "label": "Release notes" + } + ] + } + ], + "tags": [ + "finance", + "legal", + "medical", + "hr", + "tax", + "insurance", + "marketing", + "technical", + "operations", + "academic", + "government", + "draft", + "final", + "signed", + "unsigned", + "executed", + "expired", + "amended", + "void", + "confidential", + "internal", + "public", + "pii", + "phi", + "certified", + "notarized", + "scanned", + "redacted", + "template", + "urgent" + ] +} diff --git a/engine/src/stirling/agents/document_classifier.py b/engine/src/stirling/agents/document_classifier.py new file mode 100644 index 0000000000..28d9908573 --- /dev/null +++ b/engine/src/stirling/agents/document_classifier.py @@ -0,0 +1,179 @@ +from __future__ import annotations + +import logging +from pathlib import Path + +from pydantic import Field +from pydantic_ai import Agent +from pydantic_ai.output import NativeOutput + +from stirling.contracts import ( + ClassificationTaxonomy, + ClassifyDocumentRequest, + ClassifyDocumentResponse, + DocumentClassificationResponse, + PageText, +) +from stirling.models import ApiModel +from stirling.services import AppRuntime + +logger = logging.getLogger(__name__) + +# Sentinel for a label that fell outside the supplied vocabulary. +UNKNOWN_LABEL = "unknown" +# An off-list answer can never be reported as more confident than this, so a +# confident-but-wrong model answer can't clear an organisation's accept +# threshold downstream. See the design doc's "Validate" step. +UNKNOWN_MAX_CONFIDENCE = 0.2 +# 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 + +# The built-in vocabulary the classifier falls back to when a request doesn't +# supply its own. GENERATED from the TS source of truth +# (frontend/editor/src/proprietary/data/classificationTaxonomy.ts) via +# `task frontend:taxonomy` — edit that file, not this JSON. Validated into the +# typed contract on import, so a malformed entry fails fast. +_DEFAULT_TAXONOMY_PATH = Path(__file__).with_name("default_taxonomy.generated.json") +DEFAULT_TAXONOMY = ClassificationTaxonomy.model_validate_json(_DEFAULT_TAXONOMY_PATH.read_text(encoding="utf-8")) + + +_SYSTEM_PROMPT = ( + "You identify what a document is, choosing only from a fixed vocabulary you " + "are given. Decide along three axes:\n" + "- category: the document's structural family. Choose EXACTLY ONE category id.\n" + "- doc_type: the specific instrument within that family. Choose EXACTLY ONE " + "doc_type id listed under the category you chose.\n" + "- tags: zero or more descriptor ids from the tag list.\n" + "\n" + "Rules:\n" + "- Use only ids from the supplied vocabulary. If nothing fits, return " + f'"{UNKNOWN_LABEL}" for category and/or doc_type.\n' + "- The doc_type you pick must belong to the category you pick.\n" + "- type_confidence (0.0-1.0) is how sure you are that doc_type is correct.\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 taxonomy.""" + + category: str = Field(description="A category id from the vocabulary, or 'unknown'.") + doc_type: str = Field(description="A doc_type id belonging to the chosen category, or 'unknown'.") + type_confidence: float = Field(ge=0.0, le=1.0, description="Confidence that doc_type is correct.") + tags: list[str] = Field(default_factory=list, description="Descriptor ids drawn from the tag list.") + + +def render_taxonomy(taxonomy: ClassificationTaxonomy) -> str: + """Render the vocabulary for the prompt, ids first so the model echoes them.""" + lines = ["Categories (id (label): doc_types as id (label)):"] + for category in taxonomy.categories: + types = ", ".join(f"{doc_type.id} ({doc_type.label})" for doc_type in category.doc_types) or "(none)" + lines.append(f"- {category.id} ({category.label}): {types}") + lines.append(f"Tags: {', '.join(taxonomy.tags) or '(none)'}") + return "\n".join(lines) + + +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_against_taxonomy( + output: _ClassifierOutput, + taxonomy: ClassificationTaxonomy, +) -> DocumentClassificationResponse: + """Coerce a raw model answer onto the supplied vocabulary. + + An off-list category collapses both axes to ``unknown``; a doc_type that + isn't a child of its (valid) category collapses the type alone. Either + collapse caps confidence. Tags are filtered to the known set, de-duplicated, + and returned in the model's order. The model identifies; these rules decide + what is allowed to stand. + """ + categories_by_id = {category.id.lower(): category for category in taxonomy.categories} + allowed_tags = {tag.lower(): tag for tag in taxonomy.tags} + + kept_tags: list[str] = [] + for tag in output.tags: + canonical = allowed_tags.get(tag.strip().lower()) + if canonical is not None and canonical not in kept_tags: + kept_tags.append(canonical) + + category = categories_by_id.get(output.category.strip().lower()) + if category is None: + return DocumentClassificationResponse( + category=UNKNOWN_LABEL, + doc_type=UNKNOWN_LABEL, + type_confidence=min(output.type_confidence, UNKNOWN_MAX_CONFIDENCE), + tags=kept_tags, + ) + + types_by_id = {doc_type.id.lower(): doc_type for doc_type in category.doc_types} + doc_type = types_by_id.get(output.doc_type.strip().lower()) + if doc_type is None: + return DocumentClassificationResponse( + category=category.id, + doc_type=UNKNOWN_LABEL, + type_confidence=min(output.type_confidence, UNKNOWN_MAX_CONFIDENCE), + tags=kept_tags, + ) + + return DocumentClassificationResponse( + category=category.id, + doc_type=doc_type.id, + type_confidence=output.type_confidence, + tags=kept_tags, + ) + + +class DocumentClassifierAgent: + """Identifies a document's category, type, and tags against a taxonomy. + + 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: + # Override point: a request-supplied taxonomy (e.g. a future per-org / DB + # vocabulary the backend resolves) wins; the generated default is the fallback. + taxonomy = request.taxonomy or DEFAULT_TAXONOMY + window = select_window(request.pages) + prompt = self._build_prompt(request.file_name, taxonomy, window) + logger.debug("[classify] prompt:\n%s", prompt) + result = await self._agent.run(prompt) + return validate_against_taxonomy(result.output, taxonomy) + + @staticmethod + def _build_prompt(file_name: str, taxonomy: ClassificationTaxonomy, window: list[PageText]) -> str: + return ( + f"{render_taxonomy(taxonomy)}\n\n" + f"Document file name: {file_name}\n" + f"Document content (first and last pages):\n{format_window(window)}" + ) diff --git a/engine/src/stirling/api/app.py b/engine/src/stirling/api/app.py index 7e6272c28b..34b0da36fd 100644 --- a/engine/src/stirling/api/app.py +++ b/engine/src/stirling/api/app.py @@ -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) diff --git a/engine/src/stirling/api/dependencies.py b/engine/src/stirling/api/dependencies.py index 780e97c9d7..ed9aa3f041 100644 --- a/engine/src/stirling/api/dependencies.py +++ b/engine/src/stirling/api/dependencies.py @@ -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. diff --git a/engine/src/stirling/api/routes/__init__.py b/engine/src/stirling/api/routes/__init__.py index 1c1e4e1d23..f40e9918f8 100644 --- a/engine/src/stirling/api/routes/__init__.py +++ b/engine/src/stirling/api/routes/__init__.py @@ -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", diff --git a/engine/src/stirling/api/routes/document_classifier.py b/engine/src/stirling/api/routes/document_classifier.py new file mode 100644 index 0000000000..dcd04d28be --- /dev/null +++ b/engine/src/stirling/api/routes/document_classifier.py @@ -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 default taxonomy. + + 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) diff --git a/engine/src/stirling/contracts/__init__.py b/engine/src/stirling/contracts/__init__.py index 6c8d99d120..4de38af991 100644 --- a/engine/src/stirling/contracts/__init__.py +++ b/engine/src/stirling/contracts/__init__.py @@ -36,6 +36,14 @@ from .contradiction import ( ContradictionReport, ContradictionSeverity, ) +from .document_classifier import ( + ClassificationTaxonomy, + ClassifyDocumentRequest, + ClassifyDocumentResponse, + DocumentCategory, + DocumentClassificationResponse, + DocumentType, +) from .documents import ( DeleteDocumentResponse, IngestDocumentRequest, @@ -129,6 +137,9 @@ __all__ = [ "AiToolAgentStep", "ArtifactKind", "CannotContinueExecutionAction", + "ClassificationTaxonomy", + "ClassifyDocumentRequest", + "ClassifyDocumentResponse", "Claim", "CommentSpec", "CompletedExecutionAction", @@ -139,8 +150,11 @@ __all__ = [ "DeleteDocumentResponse", "PurgeOwnerResponse", "Discrepancy", + "DocumentCategory", + "DocumentClassificationResponse", "DocumentMeta", "DocumentSections", + "DocumentType", "DiscrepancyKind", "EditCannotDoResponse", "EditClarificationRequest", diff --git a/engine/src/stirling/contracts/common.py b/engine/src/stirling/contracts/common.py index 8f35c9ceb4..4fde61cab3 100644 --- a/engine/src/stirling/contracts/common.py +++ b/engine/src/stirling/contracts/common.py @@ -63,6 +63,7 @@ class WorkflowOutcome(StrEnum): UNSUPPORTED_CAPABILITY = "unsupported_capability" GENERATE_FILE = "generate_file" CONVERT_MARKDOWN = "convert_markdown" + CLASSIFICATION = "classification" class ArtifactKind(StrEnum): diff --git a/engine/src/stirling/contracts/document_classifier.py b/engine/src/stirling/contracts/document_classifier.py new file mode 100644 index 0000000000..eb55ccd78e --- /dev/null +++ b/engine/src/stirling/contracts/document_classifier.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +from typing import Literal + +from pydantic import Field + +from stirling.models import ApiModel + +from .common import WorkflowOutcome +from .documents import PageText + + +class DocumentType(ApiModel): + """A specific instrument within a category (e.g. ``nda`` inside ``contract``).""" + + id: str = Field(min_length=1) + label: str = Field(min_length=1) + + +class DocumentCategory(ApiModel): + """A structural family of documents, owning the doc_types shaped like it.""" + + id: str = Field(min_length=1) + label: str = Field(min_length=1) + doc_types: list[DocumentType] = Field(default_factory=list) + + +class ClassificationTaxonomy(ApiModel): + """The vocabulary a document is classified against. + + Supplied per request by the backend. When omitted, the engine falls back to + its small built-in default (see ``DEFAULT_TAXONOMY``). Tags are free-standing + descriptors that never own doc_types. + """ + + categories: list[DocumentCategory] = Field(min_length=1) + tags: list[str] = Field(default_factory=list) + + +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. + """ + + file_name: str = Field(min_length=1) + pages: list[PageText] = Field(default_factory=list) + taxonomy: ClassificationTaxonomy | None = None + + +class DocumentClassificationResponse(ApiModel): + """Terminal classification result. + + ``category`` and ``doc_type`` are ids drawn from the taxonomy, or the + sentinel ``"unknown"`` when the model's answer fell outside it. ``tags`` are + the subset of the model's tags that exist in the taxonomy. + """ + + outcome: Literal[WorkflowOutcome.CLASSIFICATION] = WorkflowOutcome.CLASSIFICATION + category: str + doc_type: str + type_confidence: float = Field(ge=0.0, le=1.0) + tags: 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 diff --git a/engine/tests/test_document_classifier.py b/engine/tests/test_document_classifier.py new file mode 100644 index 0000000000..59a5d93563 --- /dev/null +++ b/engine/tests/test_document_classifier.py @@ -0,0 +1,177 @@ +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from stirling.agents.document_classifier import ( + DEFAULT_TAXONOMY, + UNKNOWN_LABEL, + UNKNOWN_MAX_CONFIDENCE, + DocumentClassifierAgent, + _ClassifierOutput, + render_taxonomy, + select_window, + validate_against_taxonomy, +) +from stirling.contracts import ( + ClassificationTaxonomy, + ClassifyDocumentRequest, + DocumentCategory, + DocumentClassificationResponse, + PageText, +) +from stirling.services.runtime import AppRuntime + + +def _page(number: int, text: str = "x") -> PageText: + return PageText(page_number=number, text=text) + + +# ── select_window ─────────────────────────────────────────────────────────── + + +def test_select_window_returns_short_documents_whole() -> None: + pages = [_page(1), _page(2), _page(3), _page(4)] + assert select_window(pages, window=2) == pages + + +def test_select_window_takes_both_ends_without_overlap() -> None: + pages = [_page(n) for n in range(1, 6)] # 5 pages + selected = select_window(pages, window=2) + assert [p.page_number for p in selected] == [1, 2, 4, 5] + + +def test_select_window_handles_empty() -> None: + assert select_window([], window=2) == [] + + +def test_select_window_zero_returns_all() -> None: + pages = [_page(1), _page(2), _page(3)] + assert select_window(pages, window=0) == pages + + +# ── validate_against_taxonomy ──────────────────────────────────────────────── + + +def test_valid_classification_is_preserved() -> None: + output = _ClassifierOutput(category="contract", doc_type="nda", type_confidence=0.95, tags=["legal", "signed"]) + result = validate_against_taxonomy(output, DEFAULT_TAXONOMY) + assert isinstance(result, DocumentClassificationResponse) + assert result.category == "contract" + assert result.doc_type == "nda" + assert result.type_confidence == 0.95 + assert result.tags == ["legal", "signed"] + + +def test_off_list_category_collapses_to_unknown_with_capped_confidence() -> None: + output = _ClassifierOutput(category="spaceship", doc_type="warp_core", type_confidence=0.99) + result = validate_against_taxonomy(output, DEFAULT_TAXONOMY) + assert result.category == UNKNOWN_LABEL + assert result.doc_type == UNKNOWN_LABEL + assert result.type_confidence == UNKNOWN_MAX_CONFIDENCE + + +def test_off_list_type_keeps_category_but_unknown_type() -> None: + output = _ClassifierOutput(category="contract", doc_type="invoice", type_confidence=0.9) + result = validate_against_taxonomy(output, DEFAULT_TAXONOMY) + assert result.category == "contract" + assert result.doc_type == UNKNOWN_LABEL + assert result.type_confidence == UNKNOWN_MAX_CONFIDENCE + + +def test_type_from_a_different_category_is_not_a_child() -> None: + # "lab_result" is a valid type, but only under medical_record, not contract. + output = _ClassifierOutput(category="contract", doc_type="lab_result", type_confidence=0.8) + result = validate_against_taxonomy(output, DEFAULT_TAXONOMY) + assert result.category == "contract" + assert result.doc_type == UNKNOWN_LABEL + + +def test_matching_is_case_insensitive_and_returns_canonical_ids() -> None: + output = _ClassifierOutput(category="Contract", doc_type="NDA", type_confidence=0.7, tags=["LEGAL"]) + result = validate_against_taxonomy(output, DEFAULT_TAXONOMY) + assert result.category == "contract" + assert result.doc_type == "nda" + assert result.tags == ["legal"] + + +def test_unknown_tags_dropped_and_deduplicated_in_order() -> None: + output = _ClassifierOutput( + category="invoice", + doc_type="invoice", + type_confidence=0.9, + tags=["finance", "made-up", "finance", "legal"], + ) + result = validate_against_taxonomy(output, DEFAULT_TAXONOMY) + assert result.tags == ["finance", "legal"] + + +def test_low_confidence_is_not_raised_when_collapsing() -> None: + output = _ClassifierOutput(category="nope", doc_type="nope", type_confidence=0.05) + result = validate_against_taxonomy(output, DEFAULT_TAXONOMY) + assert result.type_confidence == 0.05 # min(0.05, 0.2) + + +# ── render_taxonomy ────────────────────────────────────────────────────────── + + +def test_render_taxonomy_lists_ids_and_tags() -> None: + rendered = render_taxonomy(DEFAULT_TAXONOMY) + assert "contract" in rendered + assert "nda" in rendered + assert "finance" in rendered + + +def test_render_taxonomy_handles_category_without_types() -> None: + taxonomy = ClassificationTaxonomy( + categories=[DocumentCategory(id="memo", label="Memo", doc_types=[])], + tags=[], + ) + rendered = render_taxonomy(taxonomy) + assert "(none)" in rendered + + +# ── DocumentClassifierAgent (inline page text) ─────────────────────────────── + + +@pytest.mark.anyio +async def test_classify_validates_model_output_against_default_taxonomy(runtime: AppRuntime) -> None: + agent = DocumentClassifierAgent(runtime) + agent._agent.run = AsyncMock( + return_value=SimpleNamespace( + output=_ClassifierOutput(category="invoice", doc_type="invoice", type_confidence=0.97, tags=["finance"]) + ) + ) + + result = await agent.classify( + ClassifyDocumentRequest( + file_name="invoice.pdf", + pages=[PageText(page_number=1, text="Invoice INV-1 total due 100.00")], + ) + ) + + assert isinstance(result, DocumentClassificationResponse) + assert result.category == "invoice" + assert result.doc_type == "invoice" + assert result.tags == ["finance"] + + +@pytest.mark.anyio +async def test_classify_collapses_off_list_model_answer(runtime: AppRuntime) -> None: + agent = DocumentClassifierAgent(runtime) + agent._agent.run = AsyncMock( + return_value=SimpleNamespace( + output=_ClassifierOutput(category="boarding_pass", doc_type="seat", type_confidence=0.9) + ) + ) + + result = await agent.classify( + ClassifyDocumentRequest(file_name="weird.pdf", pages=[PageText(page_number=1, text="Some text")]) + ) + + assert isinstance(result, DocumentClassificationResponse) + assert result.category == UNKNOWN_LABEL + assert result.doc_type == UNKNOWN_LABEL + assert result.type_confidence == UNKNOWN_MAX_CONFIDENCE diff --git a/engine/tests/test_document_classifier_routes.py b/engine/tests/test_document_classifier_routes.py new file mode 100644 index 0000000000..4f1eebd884 --- /dev/null +++ b/engine/tests/test_document_classifier_routes.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from collections.abc import Iterator + +import pytest +from fastapi.testclient import TestClient + +from stirling.api import app +from stirling.api.dependencies import get_document_classifier_agent +from stirling.contracts import ( + ClassifyDocumentRequest, + ClassifyDocumentResponse, + DocumentClassificationResponse, +) + + +class StubClassifierAgent: + """Stands in for DocumentClassifierAgent so route tests don't call a model.""" + + def __init__(self, response: ClassifyDocumentResponse) -> None: + self._response = response + + async def classify(self, _request: ClassifyDocumentRequest) -> ClassifyDocumentResponse: + return self._response + + +@pytest.fixture +def classification_client() -> Iterator[TestClient]: + app.dependency_overrides[get_document_classifier_agent] = lambda: StubClassifierAgent( + DocumentClassificationResponse( + category="contract", doc_type="nda", type_confidence=0.96, tags=["legal", "signed"] + ) + ) + try: + yield TestClient(app) + finally: + app.dependency_overrides.pop(get_document_classifier_agent, None) + + +def test_classify_returns_camel_cased_result(classification_client: TestClient) -> None: + response = classification_client.post( + "/api/v1/documents/classify", + json={"fileName": "nda.pdf", "pages": [{"pageNumber": 1, "text": "Mutual NDA between A and B."}]}, + ) + assert response.status_code == 200 + body = response.json() + assert body["outcome"] == "classification" + assert body["category"] == "contract" + assert body["docType"] == "nda" + assert body["typeConfidence"] == 0.96 + assert body["tags"] == ["legal", "signed"] + + +def test_classify_accepts_empty_pages(classification_client: TestClient) -> None: + response = classification_client.post( + "/api/v1/documents/classify", + json={"fileName": "blank.pdf", "pages": []}, + ) + assert response.status_code == 200 + + +def test_classify_rejects_empty_file_name(classification_client: TestClient) -> None: + response = classification_client.post( + "/api/v1/documents/classify", + json={"fileName": "", "pages": []}, + ) + assert response.status_code == 422 diff --git a/frontend/editor/public/locales/en-US/translation.toml b/frontend/editor/public/locales/en-US/translation.toml index 5440077934..71586b517d 100644 --- a/frontend/editor/public/locales/en-US/translation.toml +++ b/frontend/editor/public/locales/en-US/translation.toml @@ -3702,6 +3702,7 @@ backToFolder = "Back to {{folder}}" backToMyFiles = "Back to My Files" breadcrumbs = "Folder path" cancel = "Cancel" +classification = "Classification" clearSearch = "Clear search" clearSelection = "Clear selection" closeDetails = "Close details" @@ -3859,12 +3860,14 @@ uploadFilesFailedDetail = "Could not upload files: {{message}}" [filesPage.field] added = "Added" +category = "Category" +confidence = "Confidence" count = "Files" folder = "Folder" modified = "Modified" name = "Name" size = "Size" -toolHistory = "Tool history" +tags = "Tags" toolHistoryAtVersion = "Cumulative tool chain" totalSize = "Total size" type = "Type" @@ -4295,6 +4298,10 @@ desc = "Compares and shows the differences between 2 PDF Documents" tags = "difference,compare,diff,compare PDFs,compare documents,find differences,show differences,changes,what changed,track changes,revisions,version compare,side by side,contrast,delta" title = "Compare" +[home.classify] +desc = "Identify a document's type and tag its metadata." +title = "Classify" + [home.compress] desc = "Compress PDFs to reduce their file size." tags = "shrink,reduce,optimize,compress,smaller,downsize,file size,reduce size,minimize,make smaller,decrease size,optimize size" diff --git a/frontend/editor/scripts/generate-taxonomy.mts b/frontend/editor/scripts/generate-taxonomy.mts new file mode 100644 index 0000000000..abd0dc1c7e --- /dev/null +++ b/frontend/editor/scripts/generate-taxonomy.mts @@ -0,0 +1,56 @@ +/** + * Generates the engine's default classification taxonomy JSON from the type-safe + * TS source of truth (src/proprietary/data/classificationTaxonomy.ts). + * + * The Python engine can't import TypeScript, so it reads the generated JSON at + * startup. Editing the .ts and regenerating keeps the two in lockstep — the .ts + * is type-checked, so a malformed entry fails the build rather than shipping. + * + * Run: `npx tsx editor/scripts/generate-taxonomy.mts` (writes the JSON) + * `npx tsx editor/scripts/generate-taxonomy.mts --check` (CI drift guard) + * + * .mts (not .ts) so `import.meta.url` resolves paths relative to this script — + * Task invokes it from the workspace root (frontend/), same as setup-env.mts. + */ + +import { existsSync, readFileSync, writeFileSync } from "node:fs"; +import { createRequire } from "node:module"; +import { dirname, resolve } from "node:path"; +import { fileURLToPath } from "node:url"; + +// frontend/package.json has no "type": "module", so tsx treats the source .ts as +// CommonJS. require() it (tsx hooks require for .ts) to read its named export +// reliably — a named ESM import can't see the export of a CJS-interpreted file. +const require = createRequire(import.meta.url); +const { + DEFAULT_CLASSIFICATION_TAXONOMY, +} = require("../src/proprietary/data/classificationTaxonomy"); + +const here = dirname(fileURLToPath(import.meta.url)); +// editor/scripts -> repo root is three levels up (scripts -> editor -> frontend). +const repoRoot = resolve(here, "../../.."); +const outPath = resolve( + repoRoot, + "engine/src/stirling/agents/default_taxonomy.generated.json", +); + +const json = JSON.stringify(DEFAULT_CLASSIFICATION_TAXONOMY, null, 2) + "\n"; + +if (process.argv.includes("--check")) { + const current = existsSync(outPath) ? readFileSync(outPath, "utf8") : ""; + if (current !== json) { + console.error( + "default_taxonomy.generated.json is stale. Run `task frontend:taxonomy` " + + "(npx tsx editor/scripts/generate-taxonomy.mts).", + ); + process.exit(1); + } + console.log("default_taxonomy.generated.json is up to date."); +} else { + writeFileSync(outPath, json); + const categories = DEFAULT_CLASSIFICATION_TAXONOMY.categories.length; + const tags = DEFAULT_CLASSIFICATION_TAXONOMY.tags.length; + console.log( + `Wrote ${outPath}\n ${categories} categories, ${tags} loose tags`, + ); +} diff --git a/frontend/editor/src/core/components/filesPage/FileDetailsPanel.tsx b/frontend/editor/src/core/components/filesPage/FileDetailsPanel.tsx index aaab5d3b08..eff1e44257 100644 --- a/frontend/editor/src/core/components/filesPage/FileDetailsPanel.tsx +++ b/frontend/editor/src/core/components/filesPage/FileDetailsPanel.tsx @@ -20,15 +20,61 @@ import { downloadFileFromStorage, downloadMultipleFiles, } from "@app/utils/downloadUtils"; -import ToolChain from "@app/components/shared/ToolChain"; import ShareManagementModal from "@app/components/shared/ShareManagementModal"; import { useSharingEnabled } from "@app/hooks/useSharingEnabled"; import { fileStorage } from "@app/services/fileStorage"; +import { extractPDFMetadata } from "@app/services/pdfMetadataService"; import { VersionTimeline, DetailField, } from "@app/components/filesPage/VersionTimeline"; +/** Custom PDF Info-dictionary key the classify-and-tag tool writes (must match + * the backend's PdfMetadataService.CLASSIFICATION_KEY). */ +const CLASSIFICATION_KEY = "StirlingPDFClassification"; + +/** Reading classification means loading the file's bytes through PDF.js, so cap + * the auto-read by size — the app handles very large PDFs and we won't pull a + * multi-GB file into memory just to surface a metadata tag. */ +const MAX_CLASSIFICATION_READ_BYTES = 25 * 1024 * 1024; + +interface DocumentClassification { + category: string; + docType: string; + typeConfidence?: number; + tags: string[]; +} + +/** Parse the classification JSON stored in PDF metadata; null if absent/invalid. */ +function parseClassification(value: string): DocumentClassification | null { + try { + const raw = JSON.parse(value) as Record; + const category = typeof raw.category === "string" ? raw.category : ""; + const docType = typeof raw.docType === "string" ? raw.docType : ""; + if (!category && !docType) return null; + return { + category, + docType, + typeConfidence: + typeof raw.typeConfidence === "number" ? raw.typeConfidence : undefined, + tags: Array.isArray(raw.tags) + ? raw.tags.filter((tag): tag is string => typeof tag === "string") + : [], + }; + } catch { + return null; + } +} + +/** "lab_result" → "Lab result" for display. */ +function prettyLabel(id: string): string { + return id + .split(/[_\s]+/) + .filter(Boolean) + .map((word) => word[0].toUpperCase() + word.slice(1)) + .join(" "); +} + interface FileDetailsPanelProps { selectedFileIds: FileId[]; fileMap: Map; @@ -76,6 +122,12 @@ export function FileDetailsPanel({ // Metadata (size/type/dates) is collapsed by default so the panel stays // short and the action buttons keep their pinned footer in view. const [fieldsOpen, setFieldsOpen] = useState(false); + // Version journey is collapsed by default so the panel stays short. + const [versionsOpen, setVersionsOpen] = useState(false); + // Document classification read from PDF metadata, plus its (collapsed) section. + const [classification, setClassification] = + useState(null); + const [classificationOpen, setClassificationOpen] = useState(false); // Version chain for the selected file; empty for v1 or multi-select. const [versionChain, setVersionChain] = useState([]); const singleFileForChain = files.length === 1 ? files[0] : null; @@ -101,6 +153,35 @@ export function FileDetailsPanel({ }; }, [singleFileForChain]); + // Read the classification the policy wrote into PDF metadata + useEffect(() => { + setClassification(null); + const stub = singleFileForChain; + if (!stub) return; + if (stub.type && !stub.type.toLowerCase().includes("pdf")) return; + if (stub.size > MAX_CLASSIFICATION_READ_BYTES) return; + let cancelled = false; + (async () => { + try { + const file = await fileStorage.getStirlingFile(stub.id); + if (cancelled || !file) return; + const result = await extractPDFMetadata(file); + if (cancelled || !result.success) return; + const entry = result.metadata.customMetadata.find( + (item) => item.key === CLASSIFICATION_KEY, + ); + if (!entry) return; + const parsed = parseClassification(entry.value); + if (parsed && !cancelled) setClassification(parsed); + } catch (err) { + console.error("Failed to read classification metadata", err); + } + })(); + return () => { + cancelled = true; + }; + }, [singleFileForChain]); + if (files.length === 0) { return null; } @@ -232,17 +313,74 @@ export function FileDetailsPanel({ /> )} - {single.toolHistory && single.toolHistory.length > 0 && ( -

-
- {t("filesPage.field.toolHistory", "Tool history")} -
- -
+ {classification && ( + <> + + {classificationOpen && ( +
+ {classification.category && ( + + )} + {classification.docType && ( + + )} + {classification.typeConfidence != null && ( + + )} + {classification.tags.length > 0 && ( +
+ + {t("filesPage.field.tags", "Tags")} + + + {classification.tags.map((tag) => ( + + {tag} + + ))} + +
+ )} +
+ )} + )} {/* Version journey. Each tool run writes a new StirlingFile with the same `originalFileId` and an incremented @@ -266,12 +404,37 @@ export function FileDetailsPanel({ )} ) : ( - + <> + + {versionsOpen && ( + + )} + ))} ) : ( diff --git a/frontend/editor/src/core/components/filesPage/VersionTimeline.tsx b/frontend/editor/src/core/components/filesPage/VersionTimeline.tsx index 937db8c384..0a76c5bfa0 100644 --- a/frontend/editor/src/core/components/filesPage/VersionTimeline.tsx +++ b/frontend/editor/src/core/components/filesPage/VersionTimeline.tsx @@ -55,6 +55,7 @@ export interface VersionTimelineProps { currentId: FileId; onAddToWorkspace: (fileIds: FileId[]) => void; onRemove: (fileIds: FileId[]) => void; + hideHeader?: boolean; } /** Version timeline with per-row tool deltas and collapse-when-long. */ @@ -63,6 +64,7 @@ export function VersionTimeline({ currentId, onAddToWorkspace, onRemove, + hideHeader = false, }: VersionTimelineProps) { const { t } = useTranslation(); const [expandedIds, setExpandedIds] = useState>(new Set()); @@ -120,15 +122,17 @@ export function VersionTimeline({ return (
-
- - {t("filesPage.field.versionHistory", "Version journey")} - - {t("filesPage.versionsCount", "{{count}} versions", { - count: ordered.length, - })} - -
+ {!hideHeader && ( +
+ + {t("filesPage.field.versionHistory", "Version journey")} + + {t("filesPage.versionsCount", "{{count}} versions", { + count: ordered.length, + })} + +
+ )}
    {rows.map((row, idx) => { const isLast = idx === rows.length - 1; diff --git a/frontend/editor/src/core/components/shared/FileSidebarFileItem.tsx b/frontend/editor/src/core/components/shared/FileSidebarFileItem.tsx index 58a318e0d8..4a7997491e 100644 --- a/frontend/editor/src/core/components/shared/FileSidebarFileItem.tsx +++ b/frontend/editor/src/core/components/shared/FileSidebarFileItem.tsx @@ -1,4 +1,4 @@ -import { useState, useCallback, useRef } from "react"; +import { useState, useCallback, useRef, type ReactNode } from "react"; import { createPortal } from "react-dom"; import { Menu, Tooltip } from "@mantine/core"; import { useTranslation } from "react-i18next"; @@ -135,6 +135,8 @@ export interface FileItemFolderRef { export interface FileItemPolicyRef { id: string; name: string; + /** Badge glyph — the policy's own icon; falls back to a shield. */ + icon?: ReactNode; /** CSS colour for the badge (matches the policy's accent). */ accentColor: string; /** True only just after the policy was applied — drives the one-off glow, so @@ -313,7 +315,9 @@ export function FileItem({ className="file-sidebar-policy-badge" style={{ color: policy.accentColor }} > - + {policy.icon ?? ( + + )} ))} diff --git a/frontend/editor/src/proprietary/components/policies/Policies.css b/frontend/editor/src/proprietary/components/policies/Policies.css index 07fcb1865f..86d93f0095 100644 --- a/frontend/editor/src/proprietary/components/policies/Policies.css +++ b/frontend/editor/src/proprietary/components/policies/Policies.css @@ -273,6 +273,36 @@ margin: 0 0 var(--space-2); } +/* A section label rendered as a collapse toggle (Recent Activity): strip the + button chrome but keep the .pol-section-label typography, chevron pushed right. */ +.pol-section-toggle { + display: flex; + align-items: center; + gap: var(--space-2); + width: 100%; + background: none; + border: none; + padding: 0; + cursor: pointer; + text-align: left; + font-family: inherit; +} +.pol-section-chevron { + margin-left: auto; + color: var(--color-text-4); + transition: transform 0.15s ease; +} +.pol-section-chevron.is-open { + transform: rotate(180deg); +} + +/* Recent-activity feed: cap to ~4.5 rows (and never more than ~45% of the + viewport) then scroll, so a long history doesn't push the stats footer away. */ +.pol-activity-list { + max-height: min(22rem, 45vh); + overflow-y: auto; +} + /* Sub-section header inside a settings card (e.g. "Output filename"). The field directly below it carries data-first so the borders don't double up. */ .pol-subhead { diff --git a/frontend/editor/src/proprietary/components/policies/PoliciesSidebar.tsx b/frontend/editor/src/proprietary/components/policies/PoliciesSidebar.tsx index 0e87d9f982..b81879e279 100644 --- a/frontend/editor/src/proprietary/components/policies/PoliciesSidebar.tsx +++ b/frontend/editor/src/proprietary/components/policies/PoliciesSidebar.tsx @@ -439,6 +439,7 @@ export function PolicyDetailTakeover() { state.backendId, item.fileId as FileId, item.doc, + item.runId, ); } }} diff --git a/frontend/editor/src/proprietary/components/policies/PolicyDetailPanel.tsx b/frontend/editor/src/proprietary/components/policies/PolicyDetailPanel.tsx index aed4062869..1874858a32 100644 --- a/frontend/editor/src/proprietary/components/policies/PolicyDetailPanel.tsx +++ b/frontend/editor/src/proprietary/components/policies/PolicyDetailPanel.tsx @@ -7,6 +7,7 @@ import DescriptionIcon from "@mui/icons-material/Description"; import CheckCircleIcon from "@mui/icons-material/CheckCircle"; import WarningAmberIcon from "@mui/icons-material/WarningAmber"; import AutorenewIcon from "@mui/icons-material/Autorenew"; +import KeyboardArrowDownIcon from "@mui/icons-material/KeyboardArrowDown"; import LockIcon from "@mui/icons-material/Lock"; import DeleteOutlineIcon from "@mui/icons-material/DeleteOutlined"; import { PanelHeader } from "@shared/components/PanelHeader"; @@ -115,6 +116,7 @@ export function PolicyDetailPanel({ }: PolicyDetailPanelProps) { const { t } = useTranslation(); const isPaused = status === "paused"; + const [activityOpen, setActivityOpen] = useState(true); // Real configured steps drive the flow; fall back to the preset's rule labels. const enforceItems = steps && steps.length > 0 @@ -181,70 +183,87 @@ export function PolicyDetailPanel({ {/* Recent Activity */}
    -

    - {t("policies.detail.recentActivity", "Recent Activity")} -

    - {activityItems.length > 0 ? ( - - {activityItems.map((item, i) => ( - 0} - leadingTone={ - item.status === "flagged" - ? "warning" - : item.status === "processing" - ? "info" - : "success" - } - leading={ - item.status === "flagged" ? ( - - ) : item.status === "processing" ? ( - - ) : ( - - ) - } - title={item.doc} - description={ - item.status === "flagged" ? ( - - ) : ( - item.action - ) - } - meta={item.time} - trailing={ - item.status === "flagged" && onRetry ? ( - - ) : undefined - } + + {activityOpen && + (activityItems.length > 0 ? ( + +
    + {activityItems.map((item, i) => ( + 0} + leadingTone={ + item.status === "flagged" + ? "warning" + : item.status === "processing" + ? "info" + : "success" + } + leading={ + item.status === "flagged" ? ( + + ) : item.status === "processing" ? ( + + ) : ( + + ) + } + title={item.doc} + description={ + item.status === "flagged" ? ( + + ) : ( + item.action + ) + } + meta={item.time} + trailing={ + item.status === "flagged" && onRetry ? ( + + ) : undefined + } + /> + ))} +
    +     + ) : ( + + } + title={t( + "policies.detail.noActivityTitle", + "No activity yet", + )} + description={t( + "policies.detail.noActivityDescription", + "Documents will appear here once this policy runs.", + )} /> - ))} - - ) : ( - - } - title={t("policies.detail.noActivityTitle", "No activity yet")} - description={t( - "policies.detail.noActivityDescription", - "Documents will appear here once this policy runs.", - )} - /> - - )} +     + ))}
    {/* Stats — one grouped card with divided columns, intentionally diff --git a/frontend/editor/src/proprietary/components/policies/policyStatus.ts b/frontend/editor/src/proprietary/components/policies/policyStatus.ts index fab615ab25..d4364f1efe 100644 --- a/frontend/editor/src/proprietary/components/policies/policyStatus.ts +++ b/frontend/editor/src/proprietary/components/policies/policyStatus.ts @@ -22,6 +22,7 @@ export const STATUS_LABEL: Record = { */ export const ROW_ACCENT: Record = { ingestion: "blue", + classification: "orange", security: "purple", compliance: "green", routing: "amber", diff --git a/frontend/editor/src/proprietary/components/policies/policyToolChains.ts b/frontend/editor/src/proprietary/components/policies/policyToolChains.ts index e3cd231ea1..a7c224d1ea 100644 --- a/frontend/editor/src/proprietary/components/policies/policyToolChains.ts +++ b/frontend/editor/src/proprietary/components/policies/policyToolChains.ts @@ -8,6 +8,9 @@ export const POLICY_TOOL_CHAINS: Record = { // Security: redact PII + watermark + sanitize (strips JS). Which are enabled // by default comes from the preset's defaultOperations, not this list. security: ["redact", "watermark", "sanitize"], + // Classification: a single backend step that classifies the document and + // writes the result into its metadata. + classification: ["classify"], }; /** The configurable tool chain for a category, or null if it has none yet. */ diff --git a/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.retry.test.tsx b/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.retry.test.tsx index a0891a28f5..4adc812a13 100644 --- a/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.retry.test.tsx +++ b/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.retry.test.tsx @@ -33,7 +33,10 @@ vi.mock("@app/contexts/IndexedDBContext", () => ({ useIndexedDB: () => ({ bumpRevision: vi.fn() }), })); -import { usePolicyAutoRun } from "@app/components/policies/usePolicyAutoRun"; +import { + usePolicyAutoRun, + runPolicyOnFile, +} from "@app/components/policies/usePolicyAutoRun"; import { recordRunStart, getRun, @@ -107,3 +110,61 @@ describe("auto-run queue-rejection retry", () => { expect(getRun("run-2")?.status).toBe("PENDING"); }); }); + +describe("manual retry in place", () => { + it("replaces the failed run's row instead of stacking a second", async () => { + getFile.mockResolvedValue({ size: 999 } as never); + runStored.mockResolvedValue("run-new"); + + // A previously-failed run sits in the activity feed. + recordRunStart({ + runId: "run-old", + categoryId: "security", + fileId: "file-1", + fileName: "doc.pdf", + fileSize: 999, + status: "FAILED", + outputs: [], + error: "AI engine unreachable", + startedAt: 0, + }); + + // Retrying that specific run (passing its id) drops it as the new run records. + await runPolicyOnFile( + "security", + "backend-1", + "file-1" as never, + "doc.pdf", + "run-old", + ); + + expect(getRun("run-old")).toBeUndefined(); + expect(getRun("run-new")?.status).toBe("PENDING"); + }); + + it("without a replace id, leaves any existing run untouched", async () => { + getFile.mockResolvedValue({ size: 999 } as never); + runStored.mockResolvedValue("run-new"); + recordRunStart({ + runId: "run-old", + categoryId: "security", + fileId: "file-2", + fileName: "doc2.pdf", + fileSize: 999, + status: "FAILED", + outputs: [], + error: "boom", + startedAt: 0, + }); + + await runPolicyOnFile( + "security", + "backend-1", + "file-2" as never, + "doc2.pdf", + ); + + expect(getRun("run-old")).toBeDefined(); + expect(getRun("run-new")?.status).toBe("PENDING"); + }); +}); diff --git a/frontend/editor/src/proprietary/data/classificationTaxonomy.ts b/frontend/editor/src/proprietary/data/classificationTaxonomy.ts new file mode 100644 index 0000000000..fbbb51dc4b --- /dev/null +++ b/frontend/editor/src/proprietary/data/classificationTaxonomy.ts @@ -0,0 +1,300 @@ +/** + * Document-classification vocabulary — the single, type-safe source of truth. + * + * The Python engine can't import TypeScript, so this is GENERATED into + * `engine/src/stirling/agents/default_taxonomy.generated.json` by + * `editor/scripts/generate-taxonomy.mts` (`task frontend:taxonomy`, drift-guarded + * by `task frontend:taxonomy:check`). Edit THIS file, never the generated JSON. + * + * Shape mirrors the engine's `ClassificationTaxonomy` contract; the camelCase + * keys here map onto that model's aliases. + * + * Override points for later (kept simple now, designed to drop in): + * - This is the built-in DEFAULT. A per-org / DB-configured taxonomy is meant to + * layer on top, not replace this file. + * - The engine already accepts a `taxonomy` per classify request; when one is + * supplied it wins, and this default is the fallback. + * - The backend `ClassifyTagController.resolveTaxonomyOverride()` is the seam to + * load the caller's org/DB taxonomy and pass it through; today it returns none. + */ + +/** A specific instrument within a category — category-scoped (e.g. `nda` only + * under `contract`); the engine enforces it can't apply to another category. */ +export interface DocumentType { + id: string; + label: string; +} + +export interface DocumentCategory { + id: string; + label: string; + /** Category-scoped tags. */ + docTypes: DocumentType[]; +} + +export interface ClassificationTaxonomy { + /** Ordered most-common → least-common. */ + categories: DocumentCategory[]; + /** Loose, cross-cutting tags that aren't tied to any single category. */ + tags: string[]; +} + +export const DEFAULT_CLASSIFICATION_TAXONOMY: ClassificationTaxonomy = { + categories: [ + { + id: "invoice", + label: "Invoice", + docTypes: [ + { id: "invoice", label: "Invoice" }, + { id: "receipt", label: "Receipt" }, + { id: "credit_note", label: "Credit note" }, + { id: "purchase_order", label: "Purchase order" }, + { id: "quote", label: "Quote" }, + ], + }, + { + id: "contract", + label: "Contract", + docTypes: [ + { id: "nda", label: "Non-disclosure agreement" }, + { id: "employment_agreement", label: "Employment agreement" }, + { id: "service_agreement", label: "Service agreement" }, + { id: "lease_agreement", label: "Lease agreement" }, + { id: "master_service_agreement", label: "Master service agreement" }, + { id: "statement_of_work", label: "Statement of work" }, + { id: "terms_of_service", label: "Terms of service" }, + ], + }, + { + id: "financial_statement", + label: "Financial statement", + docTypes: [ + { id: "balance_sheet", label: "Balance sheet" }, + { id: "income_statement", label: "Income statement" }, + { id: "cash_flow_statement", label: "Cash flow statement" }, + { id: "bank_statement", label: "Bank statement" }, + { id: "annual_report", label: "Annual report" }, + ], + }, + { + id: "report", + label: "Report", + docTypes: [ + { id: "business_report", label: "Business report" }, + { id: "project_report", label: "Project report" }, + { id: "research_report", label: "Research report" }, + { id: "status_report", label: "Status report" }, + { id: "incident_report", label: "Incident report" }, + ], + }, + { + id: "letter", + label: "Letter", + docTypes: [ + { id: "business_letter", label: "Business letter" }, + { id: "cover_letter", label: "Cover letter" }, + { id: "recommendation_letter", label: "Recommendation letter" }, + { id: "complaint_letter", label: "Complaint letter" }, + { id: "demand_letter", label: "Demand letter" }, + ], + }, + { + id: "form", + label: "Form", + docTypes: [ + { id: "application_form", label: "Application form" }, + { id: "registration_form", label: "Registration form" }, + { id: "consent_form", label: "Consent form" }, + { id: "survey", label: "Survey" }, + { id: "questionnaire", label: "Questionnaire" }, + ], + }, + { + id: "resume", + label: "Resume", + docTypes: [ + { id: "resume", label: "Resume" }, + { id: "curriculum_vitae", label: "Curriculum vitae" }, + { id: "portfolio", label: "Portfolio" }, + { id: "reference_sheet", label: "Reference sheet" }, + ], + }, + { + id: "tax_form", + label: "Tax form", + docTypes: [ + { id: "tax_return", label: "Tax return" }, + { id: "w2", label: "W-2" }, + { id: "w9", label: "W-9" }, + { id: "form_1099", label: "Form 1099" }, + { id: "vat_return", label: "VAT return" }, + ], + }, + { + id: "expense_report", + label: "Expense report", + docTypes: [ + { id: "expense_report", label: "Expense report" }, + { id: "reimbursement_request", label: "Reimbursement request" }, + { id: "mileage_log", label: "Mileage log" }, + { id: "per_diem_claim", label: "Per diem claim" }, + ], + }, + { + id: "presentation", + label: "Presentation", + docTypes: [ + { id: "slide_deck", label: "Slide deck" }, + { id: "pitch_deck", label: "Pitch deck" }, + { id: "training_deck", label: "Training deck" }, + { id: "webinar_deck", label: "Webinar deck" }, + ], + }, + { + id: "medical_record", + label: "Medical record", + docTypes: [ + { id: "lab_result", label: "Lab result" }, + { id: "prescription", label: "Prescription" }, + { id: "discharge_summary", label: "Discharge summary" }, + { id: "medical_history", label: "Medical history" }, + { id: "imaging_report", label: "Imaging report" }, + { id: "vaccination_record", label: "Vaccination record" }, + ], + }, + { + id: "legal_filing", + label: "Legal filing", + docTypes: [ + { id: "court_filing", label: "Court filing" }, + { id: "complaint", label: "Complaint" }, + { id: "motion", label: "Motion" }, + { id: "subpoena", label: "Subpoena" }, + { id: "affidavit", label: "Affidavit" }, + { id: "deposition", label: "Deposition" }, + ], + }, + { + id: "identity_document", + label: "Identity document", + docTypes: [ + { id: "passport", label: "Passport" }, + { id: "drivers_license", label: "Driver's license" }, + { id: "national_id", label: "National ID" }, + { id: "birth_certificate", label: "Birth certificate" }, + { id: "visa", label: "Visa" }, + ], + }, + { + id: "insurance", + label: "Insurance", + docTypes: [ + { id: "insurance_policy", label: "Insurance policy" }, + { id: "insurance_claim", label: "Insurance claim" }, + { id: "certificate_of_insurance", label: "Certificate of insurance" }, + { id: "explanation_of_benefits", label: "Explanation of benefits" }, + ], + }, + { + id: "real_estate", + label: "Real estate", + docTypes: [ + { id: "deed", label: "Deed" }, + { id: "mortgage_agreement", label: "Mortgage agreement" }, + { id: "property_appraisal", label: "Property appraisal" }, + { id: "closing_disclosure", label: "Closing disclosure" }, + { id: "title_report", label: "Title report" }, + ], + }, + { + id: "shipping", + label: "Shipping", + docTypes: [ + { id: "bill_of_lading", label: "Bill of lading" }, + { id: "packing_slip", label: "Packing slip" }, + { id: "customs_declaration", label: "Customs declaration" }, + { id: "delivery_note", label: "Delivery note" }, + { id: "air_waybill", label: "Air waybill" }, + ], + }, + { + id: "hr_document", + label: "HR document", + docTypes: [ + { id: "offer_letter", label: "Offer letter" }, + { id: "performance_review", label: "Performance review" }, + { id: "payslip", label: "Payslip" }, + { id: "employee_handbook", label: "Employee handbook" }, + { id: "termination_letter", label: "Termination letter" }, + { id: "timesheet", label: "Timesheet" }, + ], + }, + { + id: "academic_record", + label: "Academic record", + docTypes: [ + { id: "transcript", label: "Transcript" }, + { id: "diploma", label: "Diploma" }, + { id: "certificate", label: "Certificate" }, + { id: "syllabus", label: "Syllabus" }, + { id: "thesis", label: "Thesis" }, + { id: "report_card", label: "Report card" }, + ], + }, + { + id: "marketing_material", + label: "Marketing material", + docTypes: [ + { id: "brochure", label: "Brochure" }, + { id: "flyer", label: "Flyer" }, + { id: "case_study", label: "Case study" }, + { id: "white_paper", label: "White paper" }, + { id: "press_release", label: "Press release" }, + ], + }, + { + id: "technical_document", + label: "Technical document", + docTypes: [ + { id: "user_manual", label: "User manual" }, + { id: "specification", label: "Specification" }, + { id: "api_documentation", label: "API documentation" }, + { id: "installation_guide", label: "Installation guide" }, + { id: "datasheet", label: "Datasheet" }, + { id: "release_notes", label: "Release notes" }, + ], + }, + ], + tags: [ + "finance", + "legal", + "medical", + "hr", + "tax", + "insurance", + "marketing", + "technical", + "operations", + "academic", + "government", + "draft", + "final", + "signed", + "unsigned", + "executed", + "expired", + "amended", + "void", + "confidential", + "internal", + "public", + "pii", + "phi", + "certified", + "notarized", + "scanned", + "redacted", + "template", + "urgent", + ], +}; diff --git a/frontend/editor/src/proprietary/data/policyDefinitions.tsx b/frontend/editor/src/proprietary/data/policyDefinitions.tsx index ce12ee3380..837287cbdb 100644 --- a/frontend/editor/src/proprietary/data/policyDefinitions.tsx +++ b/frontend/editor/src/proprietary/data/policyDefinitions.tsx @@ -16,6 +16,7 @@ import PublicIcon from "@mui/icons-material/Public"; import CloudIcon from "@mui/icons-material/Cloud"; import EmailOutlinedIcon from "@mui/icons-material/EmailOutlined"; import FolderOpenIcon from "@mui/icons-material/FolderOpen"; +import LabelOutlinedIcon from "@mui/icons-material/LabelOutlined"; import type { PolicyCategory, PolicyConfigDef, @@ -42,6 +43,14 @@ export const POLICY_CATEGORIES: PolicyCategory[] = [ icon: , desc: "Detect PII, encrypt, verify authenticity, control access, and certify documents.", }, + { + id: "classification", + label: "Classification", + icon: , + desc: "Identify each document's type on upload and tag its metadata for filing and search.", + // Needs the AI engine to classify; hidden from the policy list when it's off. + requiresAiEngine: true, + }, { id: "compliance", label: "Compliance", @@ -204,6 +213,16 @@ export const POLICY_CONFIG: Record = { // output naming + retries are set in the wizard. fields: [], }, + classification: { + summary: + "Classifies every uploaded document and writes the result to its metadata.", + rules: ["Classify", "Tag metadata"], + // Single backend step: classify the document via the AI engine and store the + // result in the document's StirlingPDFClassification metadata field. + defaultOperations: [{ operation: "classify", parameters: {} }], + scopeLabel: "All PDFs on this device", + fields: [], + }, compliance: { summary: "Validates documents against regulatory frameworks before they leave the system.", diff --git a/frontend/editor/src/proprietary/hooks/tools/classify/useClassifyOperation.ts b/frontend/editor/src/proprietary/hooks/tools/classify/useClassifyOperation.ts new file mode 100644 index 0000000000..5774d30065 --- /dev/null +++ b/frontend/editor/src/proprietary/hooks/tools/classify/useClassifyOperation.ts @@ -0,0 +1,27 @@ +import { ToolType } from "@app/hooks/tools/shared/useToolOperation"; + +/** Classify-and-tag takes no user parameters — it only needs the file. */ +export type ClassifyParameters = Record; + +export const defaultParameters: ClassifyParameters = {}; + +// Static function shared by the registry/automation executor. The backend reads +// only the file: it classifies via the AI engine and writes the result into the +// StirlingPDFClassification metadata field. +export const buildClassifyFormData = ( + _parameters: ClassifyParameters, + file: File, +): FormData => { + const formData = new FormData(); + formData.append("fileInput", file); + return formData; +}; + +export const classifyOperationConfig = { + toolType: ToolType.singleFile, + buildFormData: buildClassifyFormData, + operationType: "classify", + endpoint: "/api/v1/ai/tools/classify-and-tag", + multiFileEndpoint: false, + defaultParameters, +} as const; diff --git a/frontend/editor/src/proprietary/hooks/usePolicyCatalog.ts b/frontend/editor/src/proprietary/hooks/usePolicyCatalog.ts index 9d2bfbfa2c..c233333e39 100644 --- a/frontend/editor/src/proprietary/hooks/usePolicyCatalog.ts +++ b/frontend/editor/src/proprietary/hooks/usePolicyCatalog.ts @@ -1,4 +1,5 @@ import { useMemo } from "react"; +import { useAiEngineEnabled } from "@app/hooks/useAiEngineEnabled"; import { loadPolicyCatalog, type PolicyCatalog, @@ -10,7 +11,20 @@ import { * directly. Memoised; when the catalog becomes a backend fetch, this hook is * where loading/error state would be introduced — its consumers already treat * it as the single source of definitions. + * + * Categories flagged {@link PolicyCategory.requiresAiEngine} are hidden while the + * AI engine is off, so a policy only appears where it can actually run. */ export function usePolicyCatalog(): PolicyCatalog { - return useMemo(() => loadPolicyCatalog(), []); + const aiEngineEnabled = useAiEngineEnabled(); + return useMemo(() => { + const catalog = loadPolicyCatalog(); + if (aiEngineEnabled) return catalog; + return { + ...catalog, + categories: catalog.categories.filter( + (category) => !category.requiresAiEngine, + ), + }; + }, [aiEngineEnabled]); } diff --git a/frontend/editor/src/proprietary/hooks/usePolicyFileBadges.ts b/frontend/editor/src/proprietary/hooks/usePolicyFileBadges.ts index ad4c016895..3e60506d94 100644 --- a/frontend/editor/src/proprietary/hooks/usePolicyFileBadges.ts +++ b/frontend/editor/src/proprietary/hooks/usePolicyFileBadges.ts @@ -1,4 +1,10 @@ -import { useMemo } from "react"; +import { + useMemo, + cloneElement, + isValidElement, + type ReactElement, + type ReactNode, +} from "react"; import { usePolicyRuns } from "@app/components/policies/policyRunStore"; import type { PolicyRunRecord } from "@app/components/policies/policyRunStore"; import { useAllFiles } from "@app/contexts/FileContext"; @@ -18,8 +24,22 @@ const ACCENT_VAR: Record = { green: "var(--color-green)", amber: "var(--color-amber)", red: "var(--color-red)", + orange: "var(--color-orange)", }; +/** Glyph size for the file-sidebar policy badge. */ +const BADGE_ICON_SIZE = "0.7rem"; + +/** Reuse a policy category's own icon at badge size, + * so each badge reflects its policy */ +function toBadgeIcon(icon: ReactNode): ReactNode { + return isValidElement(icon) + ? cloneElement(icon as ReactElement<{ sx?: object }>, { + sx: { fontSize: BADGE_ICON_SIZE }, + }) + : icon; +} + /** Minimal provenance shape needed to resolve a file's inherited badges. */ type LineageStub = { id: string; @@ -58,6 +78,7 @@ export function buildPolicyBadgeMap( stubs: ReadonlyArray, labelById: ReadonlyMap, now: number, + iconById?: ReadonlyMap, ): Map { // Direct badges: a file that IS a policy run's output. const directByFile = new Map(); @@ -71,6 +92,7 @@ export function buildPolicyBadgeMap( list.push({ id: run.categoryId, name, + icon: iconById?.get(run.categoryId), accentColor: ACCENT_VAR[ROW_ACCENT[run.categoryId] ?? "blue"], recent, }); @@ -121,9 +143,17 @@ export function usePolicyFileBadges(): Map { const runs = usePolicyRuns(); const { fileStubs } = useAllFiles(); return useMemo(() => { - const labelById = new Map( - loadPolicyCatalog().categories.map((c) => [c.id, c.label]), + const categories = loadPolicyCatalog().categories; + const labelById = new Map(categories.map((c) => [c.id, c.label])); + const iconById = new Map( + categories.map((c) => [c.id, toBadgeIcon(c.icon)]), + ); + return buildPolicyBadgeMap( + runs, + fileStubs, + labelById, + Date.now(), + iconById, ); - return buildPolicyBadgeMap(runs, fileStubs, labelById, Date.now()); }, [runs, fileStubs]); } diff --git a/frontend/editor/src/proprietary/types/policies.ts b/frontend/editor/src/proprietary/types/policies.ts index c157f4759c..7cab109a2c 100644 --- a/frontend/editor/src/proprietary/types/policies.ts +++ b/frontend/editor/src/proprietary/types/policies.ts @@ -51,6 +51,11 @@ export interface PolicyCategory { * or configured. Only Security is live today. */ comingSoon?: boolean; + /** + * Requires the AI engine to be enabled. Hidden from the catalog when the + * engine is off, so the policy only appears where it can actually run. + */ + requiresAiEngine?: boolean; } /** diff --git a/frontend/shared/components/IconBadge.css b/frontend/shared/components/IconBadge.css index 10e2d28915..b068568a19 100644 --- a/frontend/shared/components/IconBadge.css +++ b/frontend/shared/components/IconBadge.css @@ -34,3 +34,6 @@ .sui-iconbadge--red { --ib-base: var(--color-red); } +.sui-iconbadge--orange { + --ib-base: var(--color-orange); +} diff --git a/frontend/shared/components/IconBadge.tsx b/frontend/shared/components/IconBadge.tsx index 0968b64592..8337e737b4 100644 --- a/frontend/shared/components/IconBadge.tsx +++ b/frontend/shared/components/IconBadge.tsx @@ -1,7 +1,13 @@ import type { ReactNode } from "react"; import "@shared/components/IconBadge.css"; -export type IconBadgeAccent = "blue" | "purple" | "green" | "amber" | "red"; +export type IconBadgeAccent = + | "blue" + | "purple" + | "green" + | "amber" + | "red" + | "orange"; export interface IconBadgeProps { children: ReactNode; diff --git a/frontend/shared/tokens/tokens.css b/frontend/shared/tokens/tokens.css index c8ba2171a0..2699bc2952 100644 --- a/frontend/shared/tokens/tokens.css +++ b/frontend/shared/tokens/tokens.css @@ -52,6 +52,10 @@ --color-amber-light: #fef3c7; --color-amber-border: #fde68a; --color-amber-dark: #92400e; + --color-orange: #f97316; + --color-orange-light: #fff7ed; + --color-orange-border: #fed7aa; + --color-orange-dark: #9a3412; /* Category accents (theme-stable) */ --color-cat-insurance: #0ea5e9; @@ -195,6 +199,10 @@ #fbbf24 — identical to the base — which left Mantine's amber hover a no-op in dark mode. */ --color-amber-dark: #f59e0b; + --color-orange: #fb923c; + --color-orange-light: #2a1408; + --color-orange-border: #7c2d12; + --color-orange-dark: #fdba74; --color-bg: #090c14; --color-bg-alt: #0d1120;