diff --git a/AGENTS.md b/AGENTS.md index a47ed0712..6582747be 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -33,7 +33,7 @@ knowhereapi-main/ │ ├── api/ # FastAPI REST API (port 5005) │ │ ├── app/ │ │ │ ├── api/v1/routes/ # Endpoint handlers -│ │ │ ├── services/ # Business logic (auth, knowledge, billing) +│ │ │ ├── services/ # Business logic (auth, ingestion, billing) │ │ │ └── repositories/ # Data access layer │ │ └── main.py # Entrypoint, runs migrations on start │ ├── worker/ # Celery worker for async document processing @@ -50,7 +50,10 @@ knowhereapi-main/ │ ├── services/retrieval/ # Core retrieval engine │ ├── services/chunks/ # DataFrame → ChunkPayload conversion │ ├── services/ai/ # LLM prompt service & AI client -│ └── utils/ # Text, file, and chunk utilities +│ ├── services/http/ # Public URL validation and outbound HTTP +│ ├── services/redis/ # Redis state, key language, and retry policy +│ ├── services/quota/ # Shared token-pool quota primitives +│ └── utils/ # Generic text, chunk, and API helpers └── deploy/ # Docker Compose & deployment scripts ``` @@ -70,14 +73,14 @@ flowchart TB end subgraph PARSE["② Document Parsing (Worker)"] - Queue --> Router["parse_service.checkerboard_inject_parse"] - Router --> Profiler["doc_profiler.profile_document"] - Profiler --> PDF["pdf_parser → MinerU"] - Profiler --> DOCX["doc_parser.parse_docx"] - Profiler --> PPTX["pptx_parser → iLoveAPI → PDF"] - Profiler --> XLSX["table_parser.parse_xlsx"] - Profiler --> MD["md_parser.parse_md"] - Profiler --> IMG["image_parser.parse_image"] + Queue --> Router["parse_service.checkerboard_parse_output"] + Router --> Profiler["profiling.doc_profiler.profile_document"] + Profiler --> PDF["formats.pdf.parser → MinerU"] + Profiler --> DOCX["formats.docx.parser.parse_docx"] + Profiler --> PPTX["formats.pptx.parser → iLoveAPI → PDF"] + Profiler --> XLSX["formats.excel.table_parser.parse_xlsx"] + Profiler --> MD["formats.markdown.parser.parse_md"] + Profiler --> IMG["formats.image.parser.parse_image"] PDF --> DF["pd.DataFrame (ALL_DF_COLS)"] DOCX --> DF PPTX --> DF @@ -117,36 +120,36 @@ flowchart TB ### Entry Point `apps/worker/app/services/document_parser/parse_service.py` → -`checkerboard_inject_parse()` +`checkerboard_parse_output()` -This is the universal entry for all file types. It: +This is the typed `ParseOutput` entry for all file types. The parser flow: -1. **Profiles** the document via `doc_profiler.profile_document()` to detect +1. **Profiles** the document via `profiling.doc_profiler.profile_document()` to detect file type, page count, and special categories (e.g. `atlas`). 2. **Routes** to the appropriate parser based on file extension. 3. **Post-processes**: cleans up unreferenced images, compresses PNG→JPG. -4. Returns `(output_dir, parsed_df)` — the parsed DataFrame. +4. Returns typed parse output with task-local artifact paths. ### Parser Routing Table | Extension | Parser Module | Strategy | |:---|:---|:---| -| `.pdf` | `pdf_parser.parse_pdfs` | MinerU API → `md_parser` → `layout_parser.pred_titles` | -| `.docx` | `doc_parser.parse_docx` + `convert_doc2dics` | OXML iteration → heading detection → hierarchical tree | -| `.doc` | `legacy_converter.doc_to_docx` → `.docx` pipeline | LibreOffice headless conversion first | -| `.pptx` | `pptx_parser.parse_pptx` | iLoveAPI PPTX→PDF → MinerU pipeline | -| `.xlsx` | `table_parser.parse_xlsx` | Sheet-by-sheet HTML table extraction | -| `.xls` | `legacy_converter.xls_to_xlsx` → `.xlsx` pipeline | LibreOffice conversion first | -| `.md` | `md_parser.parse_md` | Markdown heading parsing + LLM summaries | -| `.txt` | `txt_parser.parse_texts` → `md_parser` | Read lines then route to MD parser | -| `.png/.jpg` | `image_parser.parse_image` | VLM image description + OCR | -| `.fragment` | `fragment_parser.parse_fragment` | Raw text fragment ingestion | - -### Heading Detection: `layout_parser.pred_titles()` +| `.pdf` | `formats.pdf.parser.parse_pdfs` | MinerU API → Markdown parser → `structure.layout_parser.pred_titles` | +| `.docx` | `formats.docx.parser.parse_docx` + `convert_doc2dics` | OXML iteration → heading detection → hierarchical tree | +| `.doc` | `conversion.legacy_converter.doc_to_docx` → `.docx` pipeline | LibreOffice headless conversion first | +| `.pptx` | `formats.pptx.parser.parse_pptx` | iLoveAPI PPTX→PDF → MinerU pipeline | +| `.xlsx` | `formats.excel.table_parser.parse_xlsx` | Sheet-by-sheet HTML table extraction | +| `.xls` | `conversion.legacy_converter.xls_to_xlsx` → `.xlsx` pipeline | LibreOffice conversion first | +| `.md` | `formats.markdown.parser.parse_md` | Markdown heading parsing + LLM summaries | +| `.txt` | `formats.text.parser.parse_texts` → Markdown parser | Read lines then route to MD parser | +| `.png/.jpg` | `formats.image.parser.parse_image` | VLM image description + OCR | +| `.fragment` | `formats.fragment.parser.parse_fragment` | Raw text fragment ingestion | + +### Heading Detection: `structure.layout_parser.pred_titles()` The core hierarchical recognition module. Determines heading levels using: -1. **TOC-first**: If a DOCX TOC exists (`toc_parser.build_docx_toc_hierarchies`), +1. **TOC-first**: If a DOCX TOC exists (`structure.toc_parser.build_docx_toc_hierarchies`), use it as ground truth for heading levels. 2. **Regex patterns**: Match numbered headings like `1.2.3`, `第X章`, `（一）`. 3. **LLM smart parse**: When `smart_title_parse=True`, send candidate headings @@ -155,7 +158,7 @@ The core hierarchical recognition module. Determines heading levels using: 4. **Font clustering (PDF)**: K-means on span heights from MinerU `layout.json` to group headings into 5 discrete tiers. -### DOCX Parsing Deep Dive: `doc_parser.py` +### DOCX Parsing Deep Dive: `formats/docx/parser.py` ```mermaid flowchart LR @@ -188,11 +191,11 @@ Key logic in `parse_docx()`: ```mermaid flowchart LR - PDF[pdf_parser] --> MinerU[MinerU Cloud API] + PDF[formats.pdf.parser] --> MinerU[MinerU Cloud API] MinerU --> MDFile[Markdown + layout.json] - MDFile --> MDParser[md_parser.parse_md] + MDFile --> MDParser[formats.markdown.parser.parse_md] MDParser --> EvalHeadings[eval_md_headings + layout.json] - EvalHeadings --> PredTitles[layout_parser.pred_titles] + EvalHeadings --> PredTitles[structure.layout_parser.pred_titles] PredTitles --> Chunks[Hierarchical Chunks] ``` @@ -204,21 +207,21 @@ flowchart LR | Heading hierarchy recognition | `HIERARCHY_LLM_MODEL` | Falls back to `NORMOL_MODEL` | | Image description (VLM) | `IMAGE_MODEL` | `qwen3.5-flash` | | Image OCR / Q&A | `IMAGE_MODEL_MAX` | `qwen3.5-flash` | -| Atlas classification | VLM via `atlas_classifier` | `IMAGE_MODEL` | +| Atlas classification | VLM via `formats.atlas.classifier` | `IMAGE_MODEL` | --- -## Persisted Knowledge Base Schema (On-Disk Output) +## Persisted Document Corpus Schema (On-Disk Output) -After parsing and chunk conversion, results are persisted to `~/.knowhere/{kb_name}/`. +After parsing and chunk conversion, results are persisted to `~/.knowhere/{corpus_name}/`. This on-disk structure is the **authoritative persisted format** — the intermediate DataFrame is an internal detail. Below is the complete schema. -### KB-Level Directory Layout +### Corpus-Level Directory Layout ```text -~/.knowhere/{kb_name}/ -├── knowledge_graph.json # KB-wide graph: file metadata + cross-doc edges +~/.knowhere/{corpus_name}/ +├── knowledge_graph.json # corpus-wide graph: file metadata + cross-doc edges ├── chunk_stats.json # Per-chunk retrieval hit analytics {chunk_id → stats} ├── {source_file_name}/ # One directory per ingested document │ ├── chunks.json # All parsed chunks for this document @@ -233,12 +236,12 @@ DataFrame is an internal detail. Below is the complete schema. │ └── toc_hierarchies.json # Debug: extracted TOC structure (DOCX only) ``` -### `knowledge_graph.json` — KB-Wide Graph +### `knowledge_graph.json` — Corpus-Wide Graph ```json { "version": "2.0", - "kb_id": "test_kb", + "corpus_id": "test-corpus", "stats": { "total_files": 3, "total_chunks": 364, "total_cross_file_edges": 0 }, "files": { "AI_Security_Report.docx": { @@ -422,7 +425,7 @@ Used by agentic retrieval for 2-level section browsing. Structure: ``` The `HIERARCHY` field is a nested dict representing the full heading tree -discovered by `layout_parser.pred_titles()`. Each key is a heading title; +discovered by `structure.layout_parser.pred_titles()`. Each key is a heading title; its value is a dict of child headings (empty `{}` for leaf nodes). ### Intermediate DataFrame (`ALL_DF_COLS`) @@ -527,6 +530,19 @@ debug CSVs (`preds_*.csv`) are saved alongside for troubleshooting. `shared/services/retrieval/app_service.py` → `run_retrieval_query()` +Core retrieval internals are grouped by ownership: + +- `execution/`: request shaping, route selection, legacy route execution, and public response projection. +- `search/`: lexical channels, scoring, section filters, and candidate ranking. +- `hydration/`: row/path/reference hydration, inline assets, and result assembly. +- `graph/`: document graph publication/query support. +- `stats/`: retrieval hit recording. +- `workflow/`: query planning, step execution, synthesis, and wallet state. +- `agentic/core/`: agentic run types, token budgets, runtime config, and traces. +- `agentic/discovery/`: bottom discovery and document selection. +- `agentic/navigation/`: section-tree navigation, selection hydration, and asset tools. +- `agentic/evidence/`: evidence tree rendering and budget trimming. + ### Two Retrieval Modes The system supports two modes, controlled globally by `RETRIEVAL_AGENTIC_ENABLED` and locally via the per-request `use_agentic` toggle. @@ -543,7 +559,7 @@ flowchart LR T --> RRF RRF --> Graph[Legacy Graph Routing] Graph --> Rank[Dual-priority ranking] - Rank --> Assemble[assemble_retrieval_results] + Rank --> Assemble[hydration.result_assembly] ``` **Channel weights** (default): path=1.0, content=2.0, term=1.5 @@ -578,7 +594,7 @@ Unlike legacy retrieval which relied on static `hydrate_mode` tags, hydration is ### Result Assembly -`assemble_retrieval_results()`: +`hydration.result_assembly.assemble_retrieval_results()`: 1. Filters by `exclude_document_ids` and `exclude_sections` 2. Filters by `allowed_chunk_types` (data_type parameter) @@ -586,10 +602,10 @@ Unlike legacy retrieval which relied on static `hydrate_mode` tags, hydration is 4. Cleans asset path references from content 5. Attaches citation: `{document_id, chunk_id, source_file_name, section_path}` -### Small KB Optimization +### Small Corpus Optimization When `total_chunks <= top_k`, skips the full pipeline and returns all chunks -directly (router: `small_kb_all`). +directly (router: `small_corpus_all`). ### Caching diff --git a/CONTEXT.md b/CONTEXT.md new file mode 100644 index 000000000..c8903a584 --- /dev/null +++ b/CONTEXT.md @@ -0,0 +1,473 @@ +# CONTEXT + +## Purpose + +Knowhere API turns authenticated requests into document ingestion, document +lifecycle, retrieval, billing, and webhook workflows. + +Within this repository, `apps/api` is the coordination layer between HTTP +adapters and the shared implementations in `packages/shared-python/shared`. + +## Core Terms + +### User + +The authenticated owner of jobs, documents, credits, API keys, and webhooks. + +### Namespace + +The isolation scope for retrieval-visible data. The default namespace is +`default`. + +### Job + +The API-side intake and execution handle for a workflow such as file parsing, +URL ingestion, or demo source materialization. + +### Job Result + +The terminal artifact record attached to a Job. It stores delivery metadata, +result bundle references, and the revision that publication uses. + +### Job Transition Outcome + +The typed result of a Job state-machine transition. It preserves whether the +transition succeeded, the target state, previous state when known, attempt +count, and rejection reason while keeping older boolean facades available. + +### Job Post-Commit Effect + +A post-transaction side effect planned during terminal Job finalization and run +only after the database commit succeeds. Current effects include retrieval cache +invalidation and outbound webhook publication. + +### Job Read + +The workflow that lists a User's Jobs and projects one Job into the public Job +Result response shape. + +### Document + +The retrieval-visible knowledge object produced from a Job Result after +publication. + +### Document Section + +The hierarchical navigation node derived from parsed headings and section paths. + +### Document Chunk + +The retrieval-visible text, image, or table row attached to a Document Section. + +### Document Ingestion + +The workflow that creates a Job, accepts a file or URL source, confirms upload +state, and starts parsing work. + +### Document Ingestion Worker Dispatch + +The API-side handoff that advances an uploaded file Job to pending state and +enqueues the worker parse task with the user-aware Celery queue policy. + +### Worker Document Parsing + +The worker-side workflow that turns a source file into parsed DataFrame rows, +parsed assets, and parser debug artifacts before chunk conversion and result +packaging. + +### Worker Document Processing + +The worker-side Document Ingestion coordinator. It prepares task-local source +files, bills the workload estimate, invokes Worker Document Parsing, builds the +result package, uploads artifacts, and finalizes the Job. + +### Temporary Parse Workspace + +The task-local folder set used by Worker Document Processing: input files, +parser output files, and generated result packages. It is temporary worker +storage, not a retrieval Namespace or durable document scope. + +### Parse Output + +The stable parser adapter result with an output directory and optional parsed +DataFrame. + +### Parse Artifact + +The ingestion-side parsed content artifact. It validates parser output before +chunk conversion and result packaging. + +### Generated Result Package + +The generated ZIP bundle metadata used by terminal Job finalization, including +ZIP path, checksum, statistics, and byte size. + +### ZIP Result Packaging + +The storage-side module set that turns parsed chunks and parser artifacts into a +Generated Result Package. `ZipResultService` is the orchestration interface; +resource discovery, schema projection, and physical ZIP writing live in +separate modules. + +### Workload Estimate + +The worker-side estimate used for billing and processing metadata. It records +page count, estimation method, and any fallback reason. + +### Parser Input + +The typed worker-side parse request assembled from Job metadata, parser options, +source-file identity, output naming, and storage transform keys. + +### Document Format Routing + +The Worker Document Parsing module that selects one concrete parser adapter for +the source document format while keeping format-specific conversion details out +of the stable parser entrypoint. + +### Rendered PDF Transform + +The Worker Document Parsing module that reuses or creates rendered PDF artifacts +for PDF-backed parsing paths, including PPTX-to-PDF fallback handling, image-only +PDF rendering, temporary PDF materialization, MinerU handoff, and cleanup. + +### Heading Hierarchy + +The Worker Document Parsing module that predicts section levels from Markdown +lines, DOCX blocks, TOC context, layout metadata, heuristics, and optional LLM +inference. + +### Job Admission + +The policy checks that must pass before a new Job is created: authentication, +guest scope, system limits, billing RPM, concurrent job limits, and daily +quota. + +### Job Admission Route Policy + +The route-aware part of Job Admission that enforces guest API key scope and +system limits from plain route-admission context built by HTTP dependency +adapters. + +### Job Admission Capacity + +The quota-aware part of Job Admission that enforces billing RPM, concurrent +jobs, and daily quota. + +### Publication + +The shared workflow that turns parsed chunks into Documents, Document Sections, +Document Chunks, and document graph state. + +### Publication Content + +The Publication module that replaces a single Document revision's Document +Sections and Document Chunks from parsed chunk rows. + +### Retrieval + +The query workflow that returns cited evidence from published documents. + +### Retrieval Query + +The typed retrieval request that owns cache-shaping fields and route policy: +scope, filters, data type, channels, ranking options, and agentic toggle. + +### Workflow Run Request + +The agentic Retrieval request passed through planning and step execution. It +preserves user scope, filters, channel policy, internal recall, and explicit +ranking policy fields for the workflow path. + +### Workflow Step Request + +The per-step projection of a Workflow Run Request. It applies step-level query, +top-k, and data-type overrides while preserving the request policy. + +### Demo Source + +An API-owned canonical document shipped with the repository for demo and guest +flows. + +### Demo Source Validation + +The repeatable script workflow that validates Demo Source catalog metadata, +canonical chunks, citation projection, original file size, and regenerated +doc_nav.json output. + +### Demo Source Materialization + +The workflow that copies a Demo Source into a user's Namespace as normal Job, +Job Result, Document, and Document Chunk records. + +### Billing Workflow + +The credits purchase, checkout, webhook handling, refund reconciliation, and +tier refresh flows. + +### API Key Authentication + +The auth-time workflow that validates API keys, reads and writes the API-key +cache, and schedules best-effort last-used updates. + +### API Key Management + +The user-facing workflow that creates, lists, reads, revokes, and toggles API +keys. + +### Stripe Purchase + +The Billing Workflow adapter that creates Stripe payment intents and checkout +sessions for credits purchases. + +### Stripe Credits Settlement + +The Billing Workflow adapter that settles successful Stripe checkout and +payment-intent events into credits, payment records, and tier refreshes. + +### Stripe Webhook Reconciliation + +The Billing Workflow adapter that verifies Stripe events and reconciles credits, +payment records, and refunds. + +### Guest API Key + +A guest-tier API key with a restricted route surface. + +### Webhook Management + +The user-facing workflow for storing outbound webhook configuration and reading +delivery logs. + +### QStash Callback + +The verified async callback used to continue background work after external +delivery. + +### Public URL Policy + +The shared URL safety workflow used before Knowhere reaches user-provided or +third-party HTTP targets. It validates public HTTP/HTTPS URLs, pins resolved +addresses for outbound requests, blocks unsafe redirects, and detects URL file +types for Document Ingestion. + +### Redis State + +The shared Redis-backed runtime state used by background work, rate limits, +state-machine progress, distributed locks, and job metadata. It owns Redis key +language and Redis retry policy. + +### Quota Token Pool + +The shared Redis-backed token leasing workflow used by provider-specific quota +managers such as Ali, iLoveAPI, and MinerU. + +## apps/api Module Map + +### HTTP Adapters + +`apps/api/app/api/v1/routes/*` +`apps/api/app/api/dependencies/*` + +These modules translate HTTP requests and dependency context into application +workflow calls. + +### Application Workflows + +`apps/api/app/services/*` + +These modules coordinate Job Admission, Document Ingestion, document lifecycle, +Billing Workflow, Demo Source Materialization, webhook handling, and internal +callbacks. + +### Persistence Adapters + +`apps/api/app/repositories/*` + +These modules own database reads and writes for API-side workflows. + +### Shared Implementations + +`packages/shared-python/shared/*` + +These modules own the lower-level implementations for publication, retrieval, +state machines, storage, Redis-backed metadata, billing primitives, and core +exceptions. Shared Job lifecycle finalization lives under +`packages/shared-python/shared/services/jobs/lifecycle/*`. + +## apps/api Workflow Ownership + +### Document Ingestion + +- `app/api/v1/routes/jobs.py` +- `app/services/document_ingestion/service.py` +- `app/services/document_ingestion/creation_service.py` +- `app/services/document_ingestion/confirmation_service.py` +- `app/services/document_ingestion/handoff_service.py` +- `app/services/document_ingestion/scope_service.py` +- `app/services/document_ingestion/worker_dispatcher.py` +- `app/repositories/job_repository.py` + +### Job Read + +- `app/api/v1/routes/jobs.py` +- `app/services/jobs/read_service.py` +- `app/services/jobs/result_projection.py` +- `app/repositories/job_repository.py` + +### Job Admission + +- `app/api/dependencies/auth.py` +- `app/api/dependencies/current_user.py` +- `app/api/dependencies/route_admission.py` +- `app/api/dependencies/job_admission.py` +- `app/services/auth/*` +- `app/services/rate_limit/*` + +`auth.py`, `current_user.py`, and `route_admission.py` are HTTP dependency +adapters. `job_admission.py` owns only the route-level billing and system-limit +admission dependencies. + +### Document Lifecycle + +- `app/api/v1/routes/documents.py` +- `app/services/documents/lifecycle_service.py` +- `app/repositories/document_repository.py` + +### Retrieval + +- `app/api/v1/routes/retrieval.py` +- `packages/shared-python/shared/services/retrieval/app_service.py` +- `packages/shared-python/shared/services/retrieval/publication_service.py` +- `packages/shared-python/shared/services/retrieval/publication_content.py` +- `packages/shared-python/shared/services/retrieval/publication_models.py` +- `packages/shared-python/shared/services/retrieval/execution/*` +- `packages/shared-python/shared/services/retrieval/search/*` +- `packages/shared-python/shared/services/retrieval/hydration/*` +- `packages/shared-python/shared/services/retrieval/graph/*` +- `packages/shared-python/shared/services/retrieval/stats/*` +- `packages/shared-python/shared/services/retrieval/workflow/*` +- `packages/shared-python/shared/services/retrieval/agentic/core/*` +- `packages/shared-python/shared/services/retrieval/agentic/discovery/*` +- `packages/shared-python/shared/services/retrieval/agentic/navigation/*` +- `packages/shared-python/shared/services/retrieval/agentic/evidence/*` + +### Demo Source Materialization + +- `app/api/v1/routes/demo.py` +- `app/services/demo/*` +- `apps/api/scripts/validate_demo_documents.py` + +### Billing Workflow + +- `app/api/v1/routes/billing.py` +- `app/services/billing/*` +- `app/repositories/payment_record_repository.py` +- shared billing modules in `packages/shared-python/shared/services/billing/*` + +### API Key Management + +- `app/api/v1/routes/api_key.py` +- `app/services/auth/*` +- `app/repositories/api_key_repository.py` + +### Webhook Management + +- `app/api/v1/routes/webhook.py` +- `app/api/v1/routes/webhook_secrets.py` +- `app/services/webhook/*` +- `app/repositories/webhook_repository.py` + +### Internal Storage Events + +- `app/api/v1/routes/s3_events.py` +- `app/services/s3_events/*` + +### Storage Event Intake + +The internal workflow that decodes S3-compatible storage events, sanitizes +headers, acknowledges malformed or unsafe events, and triggers upload handoff. + +### Async Callbacks + +- `app/api/v1/routes/qstash_callbacks.py` +- `app/services/webhook/qstash_callback_service.py` + +The route owns QStash HTTP signature verification and HTTP response projection. +The workflow owns callback parsing, event status resolution, and webhook log +side effects. + +## Shared Workflow Ownership + +### Job Lifecycle Finalization + +- `shared/services/jobs/lifecycle/service.py` +- `shared/services/jobs/lifecycle/success_finalizer.py` +- `shared/services/jobs/lifecycle/failure_finalizer.py` +- `shared/services/jobs/lifecycle/result_writer.py` +- `shared/services/jobs/lifecycle/publication.py` +- `shared/services/jobs/lifecycle/post_commit_effects.py` +- `shared/services/jobs/lifecycle/webhook_outbox.py` + +## apps/worker Workflow Ownership + +### Worker Document Processing + +- `app/services/document_ingestion/service.py` +- `app/services/document_ingestion/processing_run.py` +- `app/services/document_ingestion/source_preparation.py` +- `app/services/document_ingestion/parse_execution.py` +- `app/services/document_ingestion/success_finalization.py` +- `app/services/document_ingestion/workspace.py` +- `app/services/document_ingestion/parse_result_package.py` +- `app/services/document_ingestion/processing_billing.py` + +### Worker Document Parsing + +- `app/services/document_parser/parse_service.py` +- `app/services/document_parser/orchestration/parse_input.py` +- `app/services/document_parser/orchestration/parse_session.py` +- `app/services/document_parser/orchestration/route_parse.py` +- `app/services/document_parser/orchestration/format_router.py` +- `app/services/document_parser/orchestration/format_adapters.py` +- `app/services/document_parser/formats/*` +- `app/services/document_parser/providers/*` +- `app/services/document_parser/structure/*` +- `app/services/document_parser/tables/*` +- `app/services/document_parser/assets/*` +- `app/services/document_parser/support/*` + +### Rendered PDF Transform + +- `app/services/document_parser/formats/pdf/rendered_transform.py` +- `app/services/document_parser/formats/pdf/pptx_rendering.py` +- `app/services/document_parser/formats/pdf/parser.py` +- `app/services/document_parser/formats/pptx/parser.py` + +### Heading Hierarchy + +- `app/services/document_parser/structure/heading_hierarchy.py` +- `app/services/document_parser/structure/layout_parser.py` +- `app/services/document_parser/formats/markdown/parser.py` +- `app/services/document_parser/formats/docx/parser.py` + +## Invariants + +- `apps/api` coordinates workflows. Parsing, publication, retrieval internals, + storage mechanics, and state-machine implementation mostly live outside the + route modules. +- Worker Document Parsing exposes `checkerboard_parse_output` as the stable + parser entrypoint; parser option shaping, format routing, rendered PDF + transforms, typed Parse Output, and heading inference stay behind that + entrypoint. +- A Job and a Document are not the same thing. Jobs track intake and processing; + Documents track retrieval-visible knowledge state. +- Terminal Job finalization should plan post-commit effects with primitive + identifiers and run them after the database transaction commits. +- State-machine callers that need diagnostics should consume Job Transition + Outcome; boolean state-machine methods remain compatibility facades. +- `current_job_result_id` selects the active revision of a Document. +- Namespace is part of the retrieval contract, not a UI-only label. +- Demo Sources should behave like normal Documents after materialization. +- Billing Workflow and Job Admission shape whether work is allowed to start; + they are not worker-only concerns. diff --git a/apps/api/app/api/dependencies/__init__.py b/apps/api/app/api/dependencies/__init__.py new file mode 100644 index 000000000..e0170a914 --- /dev/null +++ b/apps/api/app/api/dependencies/__init__.py @@ -0,0 +1 @@ +"""FastAPI dependency adapters.""" diff --git a/apps/api/app/api/dependencies/auth.py b/apps/api/app/api/dependencies/auth.py new file mode 100644 index 000000000..4ebec70cc --- /dev/null +++ b/apps/api/app/api/dependencies/auth.py @@ -0,0 +1,23 @@ +"""FastAPI authentication dependency adapters.""" + +from app.services.auth.current_user_authentication_service import ( + get_current_user_authentication_service, +) +from fastapi import Depends, Header +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.database import get_db + + +async def get_current_user_id( + authorization: str | None = Header( + default=None, + description="Bearer OR internal signature auth", + ), + db: AsyncSession = Depends(get_db), +) -> str: + """Authenticate the caller and return the current user ID.""" + return await get_current_user_authentication_service().authenticate_authorization_header( + db, + authorization, + ) diff --git a/apps/api/app/api/dependencies/current_user.py b/apps/api/app/api/dependencies/current_user.py new file mode 100644 index 000000000..4c5593c1a --- /dev/null +++ b/apps/api/app/api/dependencies/current_user.py @@ -0,0 +1,25 @@ +"""FastAPI current-user dependency adapters.""" + +from typing import AsyncGenerator + +from app.api.dependencies.auth import get_current_user_id +from app.api.dependencies.route_admission import get_route_admission_context +from app.services.rate_limit.data_structures import ( + CurrentUser, + RouteAdmissionContext, +) +from app.services.rate_limit.job_admission_service import JobAdmissionService +from fastapi import Depends + +_job_admission_service = JobAdmissionService() + + +async def with_current_user( + route_context: RouteAdmissionContext = Depends(get_route_admission_context), + user_id: str = Depends(get_current_user_id), +) -> AsyncGenerator[CurrentUser, None]: + current_user = await _job_admission_service.resolve_current_user( + route_context=route_context, + user_id=user_id, + ) + yield current_user diff --git a/apps/api/app/api/dependencies/job_admission.py b/apps/api/app/api/dependencies/job_admission.py new file mode 100644 index 000000000..ab87c1dc0 --- /dev/null +++ b/apps/api/app/api/dependencies/job_admission.py @@ -0,0 +1,29 @@ +"""FastAPI adapters for the Job Admission workflow.""" + +from typing import AsyncGenerator + +from app.api.dependencies.current_user import with_current_user +from app.api.dependencies.route_admission import get_route_admission_context +from app.services.rate_limit.data_structures import ( + CurrentUser, + RouteAdmissionContext, +) +from app.services.rate_limit.job_admission_service import JobAdmissionService +from fastapi import Depends + +_job_admission_service = JobAdmissionService() + + +async def require_billing_limits( + current_user: CurrentUser = Depends(with_current_user), +) -> AsyncGenerator[CurrentUser, None]: + await _job_admission_service.enforce_billing_limits(current_user=current_user) + yield current_user + + +async def require_route_system_limit( + route_context: RouteAdmissionContext = Depends(get_route_admission_context), +) -> None: + await _job_admission_service.enforce_route_system_limit( + route_context=route_context, + ) diff --git a/apps/api/app/api/dependencies/route_admission.py b/apps/api/app/api/dependencies/route_admission.py new file mode 100644 index 000000000..08457ba09 --- /dev/null +++ b/apps/api/app/api/dependencies/route_admission.py @@ -0,0 +1,36 @@ +"""FastAPI route-fact adapters for the Job Admission workflow.""" + +from app.services.rate_limit.data_structures import RouteAdmissionContext +from fastapi import Request + + +def get_route_admission_context(request: Request) -> RouteAdmissionContext: + """Extract route facts needed by Job Admission from a FastAPI request.""" + return RouteAdmissionContext( + method=request.method, + path=_get_route_path(request), + limit_identifier=_get_route_limit_identifier(request), + ) + + +def _get_route_path(request: Request) -> str: + scope_path = request.scope.get("path", request.url.path) + root_path = request.scope.get("root_path", "") + if isinstance(scope_path, str) and isinstance(root_path, str): + if root_path and scope_path.startswith(root_path): + return scope_path[len(root_path) :] + return scope_path + return request.url.path + + +def _get_route_limit_identifier(request: Request) -> str: + route = request.scope.get("route") + route_path = getattr(route, "path", None) + if isinstance(route_path, str) and route_path: + return route_path + + route_path_format = getattr(route, "path_format", None) + if isinstance(route_path_format, str) and route_path_format: + return route_path_format + + return _get_route_path(request) diff --git a/apps/api/app/api/v1/routes/api_key.py b/apps/api/app/api/v1/routes/api_key.py index 370521f7b..e42fb3dcf 100644 --- a/apps/api/app/api/v1/routes/api_key.py +++ b/apps/api/app/api/v1/routes/api_key.py @@ -2,11 +2,9 @@ API key management endpoints. """ -from app.services.auth.api_key_service import APIKeyService -from app.services.rate_limit.dependencies import ( - CurrentUser, - with_current_user, -) +from app.services.auth.api_key_management_service import APIKeyManagementService +from app.api.dependencies.current_user import with_current_user +from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends from sqlalchemy.ext.asyncio import AsyncSession @@ -25,6 +23,7 @@ ) router = APIRouter(tags=["API Key Management"]) +_api_key_management_service = APIKeyManagementService() @router.post("/create", summary="Create an API key") @@ -34,10 +33,8 @@ async def create_api_key( db: AsyncSession = Depends(get_db), ): """Create an API key.""" - api_key_service = APIKeyService.get_instance() - try: - api_key = await api_key_service.create_api_key( + api_key = await _api_key_management_service.create_api_key( session=db, user_id=current_user.user_id, name=request.name, @@ -68,11 +65,10 @@ async def list_api_keys( db: AsyncSession = Depends(get_db), ): """List API keys for the current user.""" - api_key_service = APIKeyService.get_instance() - try: - api_keys_data = await api_key_service.list_user_api_keys( - db, current_user.user_id + api_keys_data = await _api_key_management_service.list_user_api_keys( + db, + user_id=current_user.user_id, ) api_keys = [ @@ -104,11 +100,11 @@ async def revoke_api_key( db: AsyncSession = Depends(get_db), ): """Revoke an API key.""" - api_key_service = APIKeyService.get_instance() - try: - await api_key_service.revoke_api_key( - session=db, api_key_id=request.api_key_id, user_id=current_user.user_id + await _api_key_management_service.revoke_api_key( + session=db, + api_key_id=request.api_key_id, + user_id=current_user.user_id, ) return {"message": "API key revoked"} @@ -129,11 +125,11 @@ async def get_api_key( db: AsyncSession = Depends(get_db), ): """Get details for a single API key.""" - api_key_service = APIKeyService.get_instance() - try: - api_key = await api_key_service.get_api_key( - db, current_user.user_id, api_key_id + api_key = await _api_key_management_service.get_api_key( + db, + user_id=current_user.user_id, + api_key_id=api_key_id, ) if not api_key: raise NotFoundException( @@ -167,11 +163,11 @@ async def toggle_api_key( db: AsyncSession = Depends(get_db), ): """Enable or disable an API key.""" - api_key_service = APIKeyService.get_instance() - try: - success = await api_key_service.toggle_api_key( - db, current_user.user_id, api_key_id + success = await _api_key_management_service.toggle_api_key( + db, + user_id=current_user.user_id, + api_key_id=api_key_id, ) if success: return {"message": "API key status updated"} diff --git a/apps/api/app/api/v1/routes/billing.py b/apps/api/app/api/v1/routes/billing.py index 2ae4721bd..94379c33c 100644 --- a/apps/api/app/api/v1/routes/billing.py +++ b/apps/api/app/api/v1/routes/billing.py @@ -1,134 +1,72 @@ -""" -Billing API Routes -""" +"""Billing API routes.""" from typing import Optional -from app.services.billing.stripe_service import StripeService -from app.services.rate_limit.dependencies import CurrentUser, with_current_user +from app.services.billing.billing_workflow_service import ( + BillingWorkflowService, + ParseUsageResponse, +) +from app.api.dependencies.current_user import with_current_user +from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends, Query, Request -from pydantic import BaseModel -from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession -from shared.core.billing import MicroDollar -from shared.core.config import settings from shared.core.database import get_db -from shared.core.exceptions.domain_exceptions import StripeServiceException -from shared.models.database.credits_transaction import CreditsTransaction -from shared.models.database.job import Job -from shared.models.database.stripe_price_config import StripePriceConfig from shared.models.schemas.billing import ( BuyCreditsPackageRequest, BuyCreditsRequest, CheckoutSessionResponse, CreditsBalanceResponse, PaymentIntentResponse, - TransactionHistoryResponse, UsageStatsResponse, ) -from shared.services.billing import CreditsService router = APIRouter(tags=["Billing"]) +_billing_workflow_service = BillingWorkflowService() -class ParseUsageResponse(BaseModel): - """Parse usage overview response""" - - request_total: int - mom_growth: float - credits_used: float - estimated_amount: Optional[float] - success_rate: float - avg_processing_time: float - - -@router.post("/buy-credits", summary="Buy Credits") +@router.post("/buy-credits", summary="Buy Credits", response_model=PaymentIntentResponse) async def buy_credits( request: BuyCreditsRequest, current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), -): - """Buy credits via Stripe payment intent""" - stripe_service = StripeService() - - try: - # Calculate amount (100 Credits = ¥2, i.e. 1 Credit = ¥0.02) - amount_cny = request.credits_amount * 0.02 # CNY amount - amount_cents = int(amount_cny * 100) # Convert to cents - - payment_intent = await stripe_service.create_payment_intent( - user_id=current_user.user_id, - amount=amount_cents, - credits_amount=MicroDollar.from_dollars(request.credits_amount).amount, - currency="cny", - ) - - return PaymentIntentResponse( - client_secret=payment_intent["client_secret"], - payment_intent_id=payment_intent["payment_intent_id"], - ) - - except Exception as e: - raise StripeServiceException( - internal_message=f"Failed to buy credits: {str(e)}" - ) +) -> PaymentIntentResponse: + return await _billing_workflow_service.buy_credits( + request=request, + user_id=current_user.user_id, + ) @router.get( - "/credits", summary="Get Credits Balance", response_model=CreditsBalanceResponse + "/credits", + summary="Get Credits Balance", + response_model=CreditsBalanceResponse, ) async def get_credits_balance( current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), -): - """Get the current credits balance for the authenticated user""" - credits_service = CreditsService() - - try: - # Ensure user is initialized - await credits_service.ensure_user_initialized(db, current_user.user_id) - await db.commit() - - balance_micro_dollar = await credits_service.get_balance( - db, current_user.user_id - ) +) -> CreditsBalanceResponse: + return await _billing_workflow_service.get_credits_balance( + db, + user_id=current_user.user_id, + ) - return CreditsBalanceResponse( - credits_balance=MicroDollar(balance_micro_dollar).to_credit() - ) - except Exception as e: - raise StripeServiceException( - internal_message=f"Failed to get credits balance: {str(e)}" - ) - - -@router.get("/usage", summary="Get Usage Statistics") +@router.get( + "/usage", + summary="Get Usage Statistics", + response_model=UsageStatsResponse, +) async def get_usage_stats( period: str = "month", current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), -): - """Get usage statistics for the authenticated user""" - credits_service = CreditsService() - - try: - stats = await credits_service.get_usage_stats(db, current_user.user_id, period) - - return UsageStatsResponse( - period=stats["period"], - total_credits_used=MicroDollar(stats["total_used"]).to_credit(), - api_calls_count=stats["transaction_count"], - success_rate=95.0, # TODO: Calculate actual success rate from usage logs - average_response_time=stats.get("avg_response_time", 0), - top_endpoints=[], # TODO: Get top endpoints from usage logs - ) - - except Exception as e: - raise StripeServiceException( - internal_message=f"Failed to get usage statistics: {str(e)}" - ) +) -> UsageStatsResponse: + return await _billing_workflow_service.get_usage_stats( + db, + user_id=current_user.user_id, + period=period, + ) @router.get( @@ -139,85 +77,11 @@ async def get_usage_stats( async def parse_usage_overview( current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), -): - """ - Returns a usage overview: - - Total request count (deprecated, always 0) - - Month-over-month growth (deprecated, always 0) - - Credits used (from credits_transactions table, including usage and refund types) - - Estimated amount (using the first credits_package unit price: amount_cents / (100 * credits_amount)) - - Success rate (jobs: done out of terminal-state jobs) - - Average processing time (jobs: updated_at - created_at, in seconds) - """ - try: - # request_total and mom_growth: UsageLog is deprecated, hardcoded to 0 - total_requests = 0 - mom_growth = 0.0 - - # Credits used: sum usage and refund types from credits_transactions - # Usage type is negative (deduction), refund type is positive (return) - # Net consumption = abs(sum(usage + refund)), then convert to display credits - credits_row = await db.execute( - select(func.coalesce(func.sum(CreditsTransaction.credits_amount), 0)) - .where(CreditsTransaction.user_id == current_user.user_id) - .where(CreditsTransaction.transaction_type.in_(["usage", "refund"])) - ) - # Cast Decimal to int is safe here because: - # 1. Source column is BigInteger (whole numbers only) - # 2. Postgres returns Decimal to avoid overflow - # 3. Sum of integers has no fractional part, so int() is lossless - total_micro_credits_used = int(abs(credits_row.scalar_one() or 0)) - - # Success rate & average processing time (terminal-state jobs only: done / failed) - job_row = await db.execute( - select( - func.count().filter(Job.status == "done").label("done_cnt"), - func.count() - .filter(Job.status.in_(["done", "failed"])) - .label("terminal_cnt"), - func.avg(func.extract("epoch", Job.updated_at - Job.created_at)) - .filter(Job.status.in_(["done", "failed"])) - .label("avg_secs"), - ).where(Job.user_id == current_user.user_id) - ) - job_stats = job_row.first() or (0, 0, 0.0) - done_cnt = getattr(job_stats, "done_cnt", 0) or 0 - terminal_cnt = getattr(job_stats, "terminal_cnt", 0) or 0 - success_rate = (done_cnt / terminal_cnt * 100) if terminal_cnt > 0 else 0.0 - avg_processing_time = round( - float(getattr(job_stats, "avg_secs", 0.0) or 0.0), 2 - ) - - # Estimated amount: use the first credits_package price config - price_row = await db.execute( - select(StripePriceConfig) - .where(StripePriceConfig.product_type == "credits_package") - .where(StripePriceConfig.is_active.is_(True)) - .order_by(StripePriceConfig.created_at) - .limit(1) - ) - price_cfg = price_row.scalar_one_or_none() - estimated_amount = None - if price_cfg and price_cfg.credits_amount and price_cfg.credits_amount > 0: - estimated_amount = round( - price_cfg.amount_cents - * total_micro_credits_used - / (100 * price_cfg.credits_amount), - 4, - ) - - return ParseUsageResponse( - request_total=total_requests or 0, - mom_growth=round(mom_growth, 2), - credits_used=MicroDollar(total_micro_credits_used).to_credit() or 0, - estimated_amount=estimated_amount, # in dollar - success_rate=round(success_rate, 2), - avg_processing_time=avg_processing_time, - ) - except Exception as e: - raise StripeServiceException( - internal_message=f"Failed to get parse usage overview: {str(e)}" - ) +) -> ParseUsageResponse: + return await _billing_workflow_service.get_parse_usage_overview( + db, + user_id=current_user.user_id, + ) @router.get("/history", summary="Get Transaction History") @@ -226,262 +90,48 @@ async def get_transaction_history( current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), ): - """Get credits transaction history for the authenticated user""" - credits_service = CreditsService() - - try: - transactions = await credits_service.get_transaction_history( - db, current_user.user_id, limit - ) - - transaction_list = [ - TransactionHistoryResponse( - id=tx.id, - credits_amount=MicroDollar(tx.credits_amount).to_credit(), - transaction_type=tx.transaction_type, - description=tx.description, - created_at=tx.created_at, - ) - for tx in transactions - ] - - return transaction_list - - except Exception as e: - raise StripeServiceException( - internal_message=f"Failed to get transaction history: {str(e)}" - ) + return await _billing_workflow_service.get_transaction_history( + db, + user_id=current_user.user_id, + limit=limit, + ) @router.get("/price-configs", summary="Get Price Configurations") async def get_price_configs( product_type: Optional[str] = Query( - None, description="Product type: subscription or credits_package" + None, + description="Product type: subscription or credits_package", ), db: AsyncSession = Depends(get_db), -): - """Get price configuration list (subscriptions or credits packages)""" - try: - from app.services.billing.price_config_service import PriceConfigService +) -> dict[str, list[dict]]: + return await _billing_workflow_service.get_price_configs( + db, + product_type=product_type, + ) - price_config_service = PriceConfigService() - if product_type == "subscription": - # Get all subscription type configs - configs = await price_config_service.repository.get_all_active(db) - subscription_configs = [ - c for c in configs if c.product_type == "subscription" - ] - return { - "subscriptions": [ - { - "id": config.plan_id, - "plan_id": config.plan_id, - "price_id": config.price_id, - "name": ( - config.extra_metadata.get( - "display_name", config.plan_id.upper() - ) - if config.extra_metadata - else config.plan_id.upper() - ), - "description": ( - config.extra_metadata.get("description", "") - if config.extra_metadata - else "" - ), - "features": ( - config.extra_metadata.get("features", []) - if config.extra_metadata - else [] - ), - "popular": ( - config.extra_metadata.get("frontend_config", {}).get( - "popular", False - ) - if config.extra_metadata - else False - ), - "amount_cents": config.amount_cents, - "currency": config.currency, - "metadata": config.extra_metadata or {}, - } - for config in subscription_configs - ], - "credits_packages": [], - } - elif product_type == "credits_package": - # Get all credits package configs - credits_configs = await price_config_service.get_all_credits_packages(db) - return { - "subscriptions": [], - "credits_packages": [ - { - "id": config.plan_id, - "plan_id": config.plan_id, - "price_id": config.price_id, - "name": ( - config.extra_metadata.get( - "display_name", - f"{MicroDollar(config.credits_amount).to_credit()} Credits", - ) - if config.extra_metadata - else f"{MicroDollar(config.credits_amount).to_credit()} Credits" - ), - "description": ( - config.extra_metadata.get("description", "") - if config.extra_metadata - else "" - ), - "credits_amount": MicroDollar( - config.credits_amount - ).to_credit(), - "amount_cents": config.amount_cents, - "currency": config.currency, - "metadata": config.extra_metadata or {}, - } - for config in credits_configs - ], - } - else: - # Get all configs - configs = await price_config_service.repository.get_all_active(db) - subscriptions = [c for c in configs if c.product_type == "subscription"] - credits_packages = [ - c for c in configs if c.product_type == "credits_package" - ] - - return { - "subscriptions": [ - { - "id": config.plan_id, - "plan_id": config.plan_id, - "price_id": config.price_id, - "name": ( - config.extra_metadata.get( - "display_name", config.plan_id.upper() - ) - if config.extra_metadata - else config.plan_id.upper() - ), - "description": ( - config.extra_metadata.get("description", "") - if config.extra_metadata - else "" - ), - "features": ( - config.extra_metadata.get("features", []) - if config.extra_metadata - else [] - ), - "popular": ( - config.extra_metadata.get("frontend_config", {}).get( - "popular", False - ) - if config.extra_metadata - else False - ), - "amount_cents": config.amount_cents, - "currency": config.currency, - "metadata": config.extra_metadata or {}, - } - for config in subscriptions - ], - "credits_packages": [ - { - "id": config.plan_id, - "plan_id": config.plan_id, - "price_id": config.price_id, - "name": ( - config.extra_metadata.get( - "display_name", - f"{MicroDollar(config.credits_amount).to_credit()} Credits", - ) - if config.extra_metadata - else f"{MicroDollar(config.credits_amount).to_credit()} Credits" - ), - "description": ( - config.extra_metadata.get("description", "") - if config.extra_metadata - else "" - ), - "credits_amount": MicroDollar( - config.credits_amount - ).to_credit(), - "amount_cents": config.amount_cents, - "currency": config.currency, - "metadata": config.extra_metadata or {}, - } - for config in credits_packages - ], - } - - except Exception as e: - raise StripeServiceException( - internal_message=f"Failed to get price configurations: {str(e)}" - ) - - -@router.post("/buy-credits-package", summary="Buy Credits Package by Price ID") +@router.post( + "/buy-credits-package", + summary="Buy Credits Package by Price ID", + response_model=CheckoutSessionResponse, +) async def buy_credits_package( request: BuyCreditsPackageRequest, current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), -): - """Buy a credits package by its Stripe price ID""" - from sqlalchemy import select - - from shared.models.database.user import User - - stripe_service = StripeService() - - try: - # Query user email from database - result = await db.execute( - select(User.email).where(User.id == current_user.user_id) - ) - user_email = result.scalar_one_or_none() - - frontend_url = settings.FRONTEND_URL - success_url = f"{frontend_url}/billing?success=true&type=credits_package" - cancel_url = f"{frontend_url}/billing?canceled=true" - - checkout_url = await stripe_service.create_checkout_session_for_credits_package( - db=db, - user_id=current_user.user_id, - price_id=request.price_id, - success_url=success_url, - cancel_url=cancel_url, - quantity=request.quantity, - email=user_email, - ) - - return CheckoutSessionResponse(checkout_url=checkout_url, session_id="") - - except Exception as e: - raise StripeServiceException( - internal_message=f"Failed to create credits package purchase: {str(e)}" - ) +) -> CheckoutSessionResponse: + return await _billing_workflow_service.buy_credits_package( + db, + request=request, + user_id=current_user.user_id, + ) @router.post("/webhook", summary="Stripe Webhook") async def stripe_webhook(request: Request, db: AsyncSession = Depends(get_db)): - """Handle Stripe webhook events""" - stripe_service = StripeService() - - try: - payload = await request.body() - sig_header = request.headers.get("stripe-signature") - if not sig_header: - raise StripeServiceException( - internal_message="Missing stripe-signature header" - ) - - result = await stripe_service.handle_webhook(db, payload, sig_header) - - return result - - except Exception as e: - raise StripeServiceException( - internal_message=f"Failed to handle webhook: {str(e)}" - ) + return await _billing_workflow_service.handle_stripe_webhook( + db, + payload=await request.body(), + stripe_signature=request.headers.get("stripe-signature"), + ) diff --git a/apps/api/app/api/v1/routes/demo.py b/apps/api/app/api/v1/routes/demo.py index 330a5167c..4f8dbf6eb 100644 --- a/apps/api/app/api/v1/routes/demo.py +++ b/apps/api/app/api/v1/routes/demo.py @@ -4,36 +4,47 @@ from typing import Any -from app.services.demo_document_service import DemoDocumentService -from app.services.rate_limit.dependencies import CurrentUser, with_current_user +from app.services.demo.document_service import DemoDocumentService +from app.api.dependencies.current_user import with_current_user +from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends, Query from fastapi.responses import FileResponse -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, field_validator from sqlalchemy.ext.asyncio import AsyncSession from shared.core.database import get_db from shared.core.exceptions.domain_exceptions import NotFoundException +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace router = APIRouter(tags=["Demo Documents"]) -_demo_document_service = DemoDocumentService() +_demo_service = DemoDocumentService() class DemoMaterializeRequest(BaseModel): """Request to copy selected canonical demo sources into a namespace.""" - namespace: str | None = Field(None, description="Target retrieval namespace") + namespace: str | None = Field( + None, + max_length=255, + description="Target retrieval namespace", + ) demo_source_ids: list[str] = Field( default_factory=list, min_length=1, description="Canonical demo source IDs to materialize", ) + @field_validator("namespace") + @classmethod + def normalize_namespace(cls, namespace: str | None) -> str: + return normalize_retrieval_namespace(namespace) + @router.get("/catalog") async def get_demo_catalog() -> dict[str, Any]: """Return API-owned canonical demo source metadata and curated Q/A.""" - return _demo_document_service.get_catalog() + return _demo_service.get_catalog() @router.get("/sources/{demo_source_id}/chunks") @@ -43,7 +54,7 @@ async def list_demo_source_chunks( page_size: int = Query(50, ge=1, le=200, description="Items per page"), ) -> dict[str, Any]: """Return paginated canonical chunks for a demo source.""" - response = _demo_document_service.list_chunks( + response = _demo_service.list_chunks( demo_source_id=demo_source_id, page=page, page_size=page_size, @@ -59,7 +70,7 @@ async def get_demo_source_chunk( demo_chunk_id: str, ) -> dict[str, Any]: """Return one canonical demo chunk for citation focusing.""" - response = _demo_document_service.get_chunk( + response = _demo_service.get_chunk( demo_source_id=demo_source_id, demo_chunk_id=demo_chunk_id, ) @@ -75,7 +86,7 @@ async def get_demo_source_chunk( @router.get("/sources/{demo_source_id}/original") async def get_demo_source_original(demo_source_id: str) -> FileResponse: """Return the canonical original file for preview.""" - file_path = _demo_document_service.get_original_file_path( + file_path = _demo_service.get_original_file_path( demo_source_id=demo_source_id, ) if file_path is None: @@ -95,7 +106,7 @@ async def get_demo_source_asset( asset_path: str, ) -> FileResponse: """Return a canonical parsed media or table asset for preview.""" - file_path = _demo_document_service.get_asset_file_path( + file_path = _demo_service.get_asset_file_path( demo_source_id=demo_source_id, asset_path=asset_path, ) @@ -116,9 +127,9 @@ async def materialize_demo_sources( db: AsyncSession = Depends(get_db), ) -> dict[str, Any]: """Copy canonical demo sources into the authenticated user's namespace.""" - namespace = (payload.namespace or "default").strip() or "default" + namespace = normalize_retrieval_namespace(payload.namespace) try: - materialized_sources = await _demo_document_service.materialize_sources( + materialized_sources = await _demo_service.materialize_sources( db, user_id=current_user.user_id, namespace=namespace, diff --git a/apps/api/app/api/v1/routes/documents.py b/apps/api/app/api/v1/routes/documents.py index 6d21a73cc..a77c5d122 100644 --- a/apps/api/app/api/v1/routes/documents.py +++ b/apps/api/app/api/v1/routes/documents.py @@ -4,13 +4,15 @@ from typing import Literal -from app.services.document_service import DocumentService -from app.services.rate_limit.dependencies import CurrentUser, with_current_user +from app.api.dependencies.current_user import with_current_user +from app.services.documents.lifecycle_service import DocumentService +from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends, Query from sqlalchemy.ext.asyncio import AsyncSession from shared.core.database import get_db from shared.core.exceptions.domain_exceptions import NotFoundException +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace router = APIRouter(tags=["Documents"]) @@ -40,11 +42,11 @@ async def _archive_document_response( @router.get("") async def list_documents( - namespace: str | None = Query(None), + namespace: str | None = Query(None, max_length=255), current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), ): - effective_namespace = namespace or "default" + effective_namespace = normalize_retrieval_namespace(namespace) documents = await _document_service.list_documents( db, user_id=current_user.user_id, diff --git a/apps/api/app/api/v1/routes/guest.py b/apps/api/app/api/v1/routes/guest.py index 42935f7e1..700b5e5a1 100644 --- a/apps/api/app/api/v1/routes/guest.py +++ b/apps/api/app/api/v1/routes/guest.py @@ -1,7 +1,7 @@ """Guest registration routes.""" from app.services.guest.guest_registration_service import GuestRegistrationService -from app.services.rate_limit.dependencies import require_route_system_limit +from app.api.dependencies.job_admission import require_route_system_limit from fastapi import APIRouter, Depends from sqlalchemy.ext.asyncio import AsyncSession diff --git a/apps/api/app/api/v1/routes/jobs.py b/apps/api/app/api/v1/routes/jobs.py index d7c6972a2..ce9c486f7 100644 --- a/apps/api/app/api/v1/routes/jobs.py +++ b/apps/api/app/api/v1/routes/jobs.py @@ -4,658 +4,51 @@ from __future__ import annotations -import os -import uuid -from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Literal, Optional, cast -from urllib.parse import urlparse +from datetime import datetime +from typing import Optional -from app.repositories.job_repository import JobRepository -from app.services.job_document_scope_service import ( - find_active_job_for_document, - is_active_document_job_unique_violation, - raise_document_ingestion_conflict, - resolve_effective_document_scope, +from app.services.document_ingestion import DocumentIngestionService +from app.services.jobs import ( + get_job_result_for_user, + list_jobs_for_user, ) -from app.services.knowledge.kb_orchestrator import KBOrchestrator -from app.services.rate_limit.dependencies import ( - CurrentUser, - enforce_job_creation_capacity, - require_billing_limits, - with_current_user, -) -from app.services.state_machine import JobStateMachine -from fastapi import APIRouter, Depends, Query, Request -from loguru import logger -from sqlalchemy.exc import IntegrityError +from app.api.dependencies.current_user import with_current_user +from app.api.dependencies.job_admission import require_billing_limits +from app.services.rate_limit.data_structures import CurrentUser +from fastapi import APIRouter, Depends, Query from sqlalchemy.ext.asyncio import AsyncSession -from shared.core.billing import MicroDollar -from shared.core.config import settings from shared.core.database import get_db -from shared.core.exceptions.domain_exceptions import ( - ConflictException, - JobOperationException, - NotFoundException, - PermissionDeniedException, - RateLimitException, - UnavailableException, - ValidationException, -) -from shared.core.exceptions.webhook_exceptions import WebhookConfigException -from shared.core.state_machine.states import JobStatus from shared.models.schemas.job import ( ConfirmUploadRequest, JobCreate, JobList, JobResponse, JobResultResponse, - StandardErrorObject, -) -from shared.services.storage.file_upload_service import FileUploadService -from shared.utils.utc_now import utc_now_naive -from shared.utils.url_security import ( - validate_http_url_and_resolve_ip_async, ) -from shared.utils.error_details import normalize_error_details -from shared.utils.url_file_type import resolve_file_extension_async router = APIRouter(tags=["Jobs"]) -JobStatusValue = Literal[ - "pending", "waiting-file", "running", "converting", "done", "failed" -] +_document_ingestion_service = DocumentIngestionService() # ==================== Shared Helpers ==================== -def get_supported_formats() -> str: - """Return the supported file extensions as a comma-separated string.""" - return ", ".join(sorted(settings.get_supported_extensions())) - - -async def transition_to_uploaded( - db: AsyncSession, - job_id: str, - job_type: str, - trigger: str = "manual_upload_completed", -): - """ - Move the job into the uploaded flow. - - Args: - db: Database session. - job_id: Job identifier. - job_type: Job type. - trigger: Transition trigger. - """ - state_machine = JobStateMachine() - - # Once the upload is confirmed, transition the job to pending. - await state_machine.transition( - db, job_id, JobStatus.PENDING.value, trigger, None, "system" - ) - - -async def start_workflow_for_job( - db: AsyncSession, - job_id: str, - job_type: str, - source_type: str, - user_id: str, - file_path: Optional[str] = None, - file_url: Optional[str] = None, -): - """ - Start the workflow for a job. - - Args: - db: Database session. - job_id: Job identifier. - job_type: Job type. - source_type: Source type. - user_id: User identifier. - file_path: File path. - file_url: File URL. - """ - if job_type == "kb_management": - orchestrator = KBOrchestrator() - await orchestrator.start_workflow( - db=db, - job_id=job_id, - source_type=source_type, - file_path=file_path, - file_url=file_url, - user_id=user_id, - ) - else: - raise ValidationException( - user_message="Unsupported job type", - violations=[ - { - "field": "job_type", - "description": f"Job type '{job_type}' is not supported", - } - ], - ) - - -def check_job_permission(job, user_id: str, job_id: str) -> None: - """ - Verify that the job belongs to the current user. - - Args: - job: Job object. - user_id: Current user ID. - job_id: Requested job ID. - - Raises: - HTTPException: Raised when the user does not own the job. - """ - if not job: - raise NotFoundException( - resource="Job", resource_id=job_id, internal_message="Job not found" - ) - - if str(job.user_id) != user_id: - raise PermissionDeniedException( - user_message="You don't have permission to access this job", - ) - - -def _build_error_response( - job: Any, job_metadata: Optional[dict] = None -) -> Optional[StandardErrorObject]: - """ - Build StandardErrorObject for embedded error pattern. - - Args: - job: Job object with job_id, error_code, and error_message - job_metadata: Job metadata dict that may contain error_details - - Returns: - StandardErrorObject or None - """ - if not job.error_message: - return None - - # Extract error_details from job_metadata if present - error_details = None - if job_metadata and isinstance(job_metadata, dict): - error_details = normalize_error_details(job_metadata.get("error_details")) - - return StandardErrorObject( - code=job.error_code or "UNKNOWN", - message=job.error_message, - request_id=job.job_id, - details=error_details, - ) - - -def create_job_response( - job_id: str, - job, - source_type: str, - data_id: Optional[str], - namespace: Optional[str] = None, - document_id: Optional[str] = None, - upload_url: Optional[str] = None, - upload_headers: Optional[dict] = None, - expires_in: Optional[int] = None, -) -> JobResponse: - """ - Build a JobResponse object. - - Args: - job_id: Job identifier. - job: Job object. - source_type: Source type. - data_id: Data identifier. - upload_url: Upload URL in file mode. - upload_headers: Upload headers in file mode. - expires_in: Upload expiry in file mode. - - Returns: - JobResponse: Serialized job response payload. - """ - return JobResponse( - job_id=job_id, - status=job.status, - source_type=source_type, - data_id=data_id, - namespace=namespace, - document_id=document_id, - created_at=job.created_at, - upload_url=upload_url, - upload_headers=upload_headers, - expires_in=expires_in, - ) - - -def resolve_public_document_id(job) -> Optional[str]: - """Expose document_id only after it is published in the persisted job result.""" - job_result = getattr(job, "job_result", None) - published_document_id = getattr(job_result, "document_id", None) - if isinstance(published_document_id, str) and published_document_id: - return published_document_id - - return None - - -def validate_file_type(file_name: str) -> bool: - """ - Return whether the file extension is supported. - - Args: - file_name: File name. - - Returns: - bool: Whether the file type is supported. - """ - if not file_name: - return False - - file_extension = os.path.splitext(file_name)[1].lower() - - return file_extension in settings.get_supported_extensions() - - -def ensure_utc(dt: Optional[datetime]) -> Optional[datetime]: - """Normalize a datetime to UTC.""" - if not dt: - return None - if dt.tzinfo: - return dt.astimezone(timezone.utc) - return dt.replace(tzinfo=timezone.utc) - - -def normalize_naive_utc_filter_datetime(dt: Optional[datetime]) -> Optional[datetime]: - """Convert a query datetime into the naive UTC form used by database columns.""" - if dt is None: - return None - if dt.tzinfo is None or dt.utcoffset() is None: - return dt - return dt.astimezone(timezone.utc).replace(tzinfo=None) - - -def require_utc(dt: Optional[datetime], *, field_name: str) -> datetime: - """Normalize a required datetime to UTC.""" - normalized_dt = ensure_utc(dt) - if normalized_dt is None: - raise JobOperationException( - internal_message=f"Job is missing required datetime field: {field_name}" - ) - return normalized_dt - - -def to_job_status_value(status: str) -> JobStatusValue: - """Cast a persisted job status into the response literal type.""" - return cast(JobStatusValue, status) - - @router.post("", response_model=JobResponse, summary="Create a parsing job") @router.post("/", include_in_schema=False) async def create_job( # pyright: ignore[reportGeneralTypeIssues] payload: JobCreate, - http_request: Request, current_user: CurrentUser = Depends(require_billing_limits), db: AsyncSession = Depends(get_db), ): """ Create a parsing job. """ - try: - job_id = f"job_{uuid.uuid4().hex[:12]}" - # Validate input parameters. - if payload.source_type == "file" and not payload.file_name: - raise ValidationException( - user_message="file_name is required when source_type is 'file'", - violations=[ - { - "field": "file_name", - "description": "Required for file source type", - } - ], - ) - if payload.source_type == "url" and not payload.source_url: - raise ValidationException( - user_message="source_url is required when source_type is 'url'", - violations=[ - { - "field": "source_url", - "description": "Required for url source type", - } - ], - ) - - # Validate webhook config if present - if payload.webhook: - # Check for URL validity - if payload.webhook.url: - validation_result = await validate_http_url_and_resolve_ip_async( - payload.webhook.url, - ) - if not validation_result.is_valid: - raise WebhookConfigException( - user_message="Invalid webhook URL", - internal_message=f"Webhook validation failed: {validation_result.error_message}", - ) - - # Validate the source file type. - if ( - payload.source_type == "file" - and payload.file_name - and not validate_file_type(payload.file_name) - ): - supported_formats = get_supported_formats() - raise ValidationException( - user_message=f"Unsupported file type. Supported formats: {supported_formats}", - violations=[ - {"field": "file_name", "description": "File type not supported"} - ], - ) - elif payload.source_type == "url": - assert payload.source_url is not None - # Resolve file type from URL path or Content-Type header - file_ext = await resolve_file_extension_async(payload.source_url) - if not file_ext: - supported_formats = get_supported_formats() - raise ValidationException( - user_message=f"Unsupported URL file type. Supported formats: {supported_formats}", - violations=[ - { - "field": "source_url", - "description": "URL file type not supported", - } - ], - ) - - job_type = "kb_management" - - # Keep job creation lightweight. - from shared.services.redis import RedisServiceFactory - - redis_service = RedisServiceFactory.get_service() - - # Build job_metadata without embedding user_config. - from shared.models.schemas.job_metadata import JobMetadataHelper - - job_metadata = JobMetadataHelper.create_from_request(payload) - requested_document_id = cast(Optional[str], job_metadata.get("document_id")) - if requested_document_id: - active_job = await find_active_job_for_document( - db, - user_id=current_user.user_id, - document_id=requested_document_id, - ) - if active_job is not None: - raise_document_ingestion_conflict( - document_id=requested_document_id, - active_job_id=active_job.job_id, - ) - ( - effective_document_id, - effective_namespace, - ) = await resolve_effective_document_scope( - db, - user_id=current_user.user_id, - document_id=requested_document_id, - requested_namespace=cast(Optional[str], payload.namespace), - ) - if not requested_document_id: - active_job = await find_active_job_for_document( - db, - user_id=current_user.user_id, - document_id=effective_document_id, - ) - if active_job is not None: - raise_document_ingestion_conflict( - document_id=effective_document_id, - active_job_id=active_job.job_id, - ) - job_metadata["document_id"] = effective_document_id - job_metadata["namespace"] = effective_namespace - - # Enforce Layers 2-3 immediately before DB insert so the row lock - # lifetime is limited to capacity check + create_job commit. - await enforce_job_creation_capacity( - request=http_request, - db=db, - current_user=current_user, - ) - - if payload.source_type == "file": - # File-upload mode: reserve the job row first. - assert payload.file_name is not None - file_extension = os.path.splitext(payload.file_name)[1] - s3_key = f"uploads/{job_id}{file_extension}" - job_metadata["source_file_name"] = payload.file_name - job_metadata["source_type"] = "file" - - # Create the waiting-file job row with the final S3 key in one insert. - job_repo = JobRepository() - try: - job = await job_repo.create_job( - db=db, - job_id=job_id, - user_id=current_user.user_id, - job_type=job_type, - source_type="file", - file_path=None, # The file has not been uploaded yet. - webhook_url=payload.webhook.url if payload.webhook else None, - metadata=job_metadata, - initial_state="waiting-file", - s3_key=s3_key, - ) - except IntegrityError as exc: - if is_active_document_job_unique_violation(exc): - raise_document_ingestion_conflict(document_id=effective_document_id) - raise - - if not job: - raise JobOperationException( - internal_message="Failed to create job in database" - ) - - # Generate the presigned upload URL. - upload_service = FileUploadService() - upload_info = await upload_service.generate_upload_url( - job_id, file_extension - ) - - # 3. Cache job_metadata in Redis for two hours. - from shared.services.redis.job_metadata_service import JobMetadataService - - metadata_service = JobMetadataService(redis_service) - await metadata_service.save_metadata(job_id, job_metadata) - - # 4. Cache the basic job info in Redis for two hours. - from datetime import datetime - - from shared.services.redis import JobInfoRedisService - - job_info_service = JobInfoRedisService(redis_service) - job_info = { - "job_id": job_id, - "s3_key": s3_key, - "user_id": current_user.user_id, - "webhook_enabled": bool(payload.webhook and payload.webhook.url), - "job_type": job_type, - "source_type": "file", - "created_at": datetime.now(timezone.utc).isoformat(), - } - await job_info_service.save_job_info(job_id, job_info) - - logger.info( - f"Job {job_id} upload_url returned to client: {upload_info['upload_url']}" - ) - - # Build the response payload. - response = create_job_response( - job_id=job_id, - job=job, - source_type="file", - data_id=payload.data_id, - namespace=effective_namespace, - upload_url=upload_info["upload_url"], - upload_headers=upload_info["upload_headers"], - expires_in=upload_info["expires_in"], - ) - - return response - - else: - # URL mode: create the job first, then download and upload asynchronously. - try: - assert payload.source_url is not None - # Resolve file extension (URL path first, then Content-Type header) - file_extension = await resolve_file_extension_async(payload.source_url) - if not file_extension: - supported_formats = get_supported_formats() - raise ValidationException( - user_message=f"Unsupported URL file type. Supported formats: {supported_formats}", - violations=[ - { - "field": "source_url", - "description": "URL file type not supported", - } - ], - ) - - parsed_url = urlparse(payload.source_url) - url_basename = str(os.path.basename(parsed_url.path)) - # Ensure source_file_name carries the correct extension. - # URLs like arxiv.org/pdf/1706.03762 have no real extension in the path. - if ( - url_basename - and os.path.splitext(url_basename)[1].lower() == file_extension - ): - source_file_name = url_basename - elif url_basename: - source_file_name = f"{url_basename}{file_extension}" - else: - source_file_name = ( - f"url_file_{uuid.uuid4().hex[:8]}{file_extension}" - ) - - s3_key = f"uploads/{job_id}{file_extension}" - - job_metadata.update( - { - "source_file_name": source_file_name, - "source_url": payload.source_url, - "source_type": "url", - } - ) - - # Create the waiting-file job row; the file will arrive asynchronously. - job_repo = JobRepository() - try: - job = await job_repo.create_job( - db=db, - job_id=job_id, - user_id=current_user.user_id, - job_type=job_type, - source_type="url", - file_path=None, - webhook_url=payload.webhook.url if payload.webhook else None, - metadata=job_metadata, - initial_state=JobStatus.WAITING_FILE.value, # Reuse waiting-file for URL uploads. - s3_key=s3_key, # Precomputed target S3 key. - ) - except IntegrityError as exc: - if is_active_document_job_unique_violation(exc): - raise_document_ingestion_conflict( - document_id=effective_document_id - ) - raise - - if not job: - raise JobOperationException( - internal_message="Failed to create URL job in database" - ) - - # Cache job_metadata in Redis for two hours. - from shared.services.redis.job_metadata_service import ( - JobMetadataService, - ) - - metadata_service = JobMetadataService(redis_service) - await metadata_service.save_metadata(job_id, job_metadata) - - # Cache the basic job info in Redis for two hours. - from datetime import datetime - - from shared.services.redis import JobInfoRedisService - - job_info_service = JobInfoRedisService(redis_service) - job_info = { - "job_id": job_id, - "s3_key": s3_key, - "user_id": current_user.user_id, - "webhook_enabled": bool(payload.webhook and payload.webhook.url), - "job_type": job_type, - "source_type": "url", - "created_at": datetime.now(timezone.utc).isoformat(), - } - await job_info_service.save_job_info(job_id, job_info) - - # Start the URL download/upload task asynchronously in the worker. - from shared.core.celery_app import get_celery_app - - celery_app = get_celery_app() - upload_url_file_task = celery_app.signature( - "app.core.tasks.kb_tasks.upload_url_file_task" - ) - upload_url_file_task.apply_async( - args=[job_id, payload.source_url, current_user.user_id], - kwargs={ - "job_type": job_type, - }, - ) - - # Build the response payload. - response = create_job_response( - job_id=job_id, - job=job, - source_type="url", - data_id=payload.data_id, - namespace=effective_namespace, - ) - - return response - - except ValidationException: - raise - except WebhookConfigException: - raise - except ConflictException: - raise - except (RateLimitException, UnavailableException): - raise - except JobOperationException: - raise - except Exception as e: - logger.error(f"Failed to create URL job: {e}") - raise JobOperationException( - internal_message=f"URL job creation failed: {str(e)}" - ) - - except NotFoundException: - raise - except ValidationException: - raise - except ConflictException: - raise - except WebhookConfigException: - raise - except (RateLimitException, UnavailableException): - raise - except JobOperationException: - raise - except Exception as e: - logger.error(f"Failed to create job: {e}") - raise JobOperationException(internal_message=f"Job creation failed: {str(e)}") + return await _document_ingestion_service.create_job( + db, + payload=payload, + current_user=current_user, + ) @router.get("", response_model=JobList, summary="List jobs") @@ -680,189 +73,17 @@ async def list_jobs( """ List jobs for the current user. """ - try: - job_repo = JobRepository() - - if recent_days not in (None, 1, 7, 30): - raise ValidationException( - user_message="recent_days only supports 1, 7, or 30", - violations=[{"field": "recent_days", "description": "Invalid value"}], - ) - created_after: Optional[datetime] = None - if recent_days: - created_after = utc_now_naive() - timedelta(days=recent_days) - - normalized_start_time = normalize_naive_utc_filter_datetime(start_time) - normalized_end_time = normalize_naive_utc_filter_datetime(end_time) - - if ( - normalized_start_time - and normalized_end_time - and normalized_start_time > normalized_end_time - ): - raise ValidationException( - user_message="start_time cannot be later than end_time", - violations=[ - {"field": "start_time", "description": "Must be before end_time"} - ], - ) - # start_time / end_time take priority over recent_days. - if normalized_start_time: - created_after = normalized_start_time - created_before = normalized_end_time - - # Count matching rows. - total_count = await job_repo.count_jobs_by_user( - db=db, - user_id=current_user.user_id, - created_after=created_after, - created_before=created_before, - job_type=job_type, - job_status=job_status, - ) - - # Fetch the matching jobs. - jobs = await job_repo.get_jobs_by_user( - db=db, - user_id=current_user.user_id, - limit=page_size, - offset=(page - 1) * page_size, - created_after=created_after, - created_before=created_before, - job_type=job_type, - job_status=job_status, - ) - - # Build the response payload. - job_responses = [] - upload_service = FileUploadService() - from shared.models.schemas.job_metadata import JobMetadataHelper - from shared.services.redis import RedisServiceFactory - - redis_service = RedisServiceFactory.get_service() - for job in jobs: - # Load job_metadata through the shared access path. - job_metadata = await job_repo.get_job_metadata( - db, job.job_id, redis_service - ) - job_result = job.job_result - status_for_api = to_job_status_value(job.status) - - result_url = None - result = None - result_url_expires_at = job.created_at # Default to created_at. - - if job_result and job_result.result_s3_key: - result_url_info = cast( - Dict[str, Any], - await upload_service.generate_download_url( - job_result.result_s3_key - ), - ) - result_url = result_url_info["download_url"] - - # Read checksum-only data from inline_payload when present. - if job_result.inline_payload: - result = job_result.inline_payload - - # Compute result_url_expires_at when a download URL was issued. - if result_url: - expires_in = int(result_url_info.get("expires_in", 3600)) - result_url_expires_at = utc_now_naive() + timedelta( - seconds=expires_in - ) - - original_request = ( - job_metadata.get("original_request") - if isinstance(job_metadata, dict) - else {} - ) - source_url = ( - original_request.get("source_url") - if isinstance(original_request, dict) - else None - ) - file_name = None - if source_url: - parsed_source = urlparse(source_url) - file_name = os.path.basename(parsed_source.path) or None - if not file_name and isinstance(original_request, dict): - file_name = original_request.get("file_name") - file_extension = None - if file_name: - ext = os.path.splitext(file_name)[1] - file_extension = ext[1:].upper() if ext else None - - parsing_params = {} - if isinstance(original_request, dict): - parsing_params = original_request.get("parsing_params") or {} - if not parsing_params and isinstance(job_metadata, dict): - parsing_params = job_metadata.get("parsing_params") or {} - model = ( - parsing_params.get("model") - if isinstance(parsing_params, dict) - else None - ) - ocr_enabled = ( - parsing_params.get("ocr_enabled") - if isinstance(parsing_params, dict) - else None - ) - - duration_seconds = None - if job.updated_at and job.created_at: - duration_seconds = (job.updated_at - job.created_at).total_seconds() - - job_responses.append( - JobResultResponse( - job_id=job.job_id, - namespace=JobMetadataHelper.get_field(job_metadata, "namespace"), - document_id=resolve_public_document_id(job), - status=status_for_api, - source_type=job.source_type, - data_id=JobMetadataHelper.get_field(job_metadata, "data_id"), - created_at=require_utc(job.created_at, field_name="created_at"), - progress=None, # The list view does not expose detailed progress. - error=_build_error_response(job, job_metadata), - result=result, - result_url=result_url, - result_url_expires_at=require_utc( - result_url_expires_at, - field_name="result_url_expires_at", - ), - file_name=file_name, - file_extension=file_extension, - model=model, - ocr_enabled=ocr_enabled, - duration_seconds=duration_seconds, - credits_spent=( - MicroDollar(job.credits_charged).to_credit() - if hasattr(job, "credits_charged") - else 0 - ), - ) - ) - - # Compute the total page count. - import math - - total_pages = math.ceil(total_count / page_size) if total_count > 0 else 0 - - response = JobList( - jobs=job_responses, - total=total_count, - page=page, - page_size=page_size, - total_pages=total_pages, - ) - - return response - - except Exception as e: - logger.error(f"Failed to list jobs: {e}") - raise JobOperationException( - internal_message=f"Failed to get job list: {str(e)}" - ) + return await list_jobs_for_user( + db, + user_id=current_user.user_id, + page=page, + page_size=page_size, + job_status=job_status, + job_type=job_type, + recent_days=recent_days, + start_time=start_time, + end_time=end_time, + ) @router.get("/{job_id}", response_model=JobResultResponse, summary="Get a job result") @@ -874,138 +95,11 @@ async def get_job_result( """ Return the result payload for one job. """ - try: - job_repo = JobRepository() - - # Load the job and verify access. - job = await job_repo.get_job_by_id(db, job_id) - check_job_permission(job, current_user.user_id, job_id) - assert job is not None - - status_for_api = to_job_status_value(job.status) - - # Load detailed progress from Redis while the job is running. - progress = None - if status_for_api == "running": - # TODO: Load detailed progress from Redis and convert it to the progress schema. - # from shared.services.redis import RedisServiceFactory - # redis_service = RedisServiceFactory.get_service() - # from shared.utils.redis_key_builder import redis_key_builder - - # progress_key = redis_key_builder.task_progress(job_id) - # progress = await redis_service.hgetall(progress_key) - progress = {"total_pages": 10, "processed_pages": 5} - - # Load job_metadata through the shared access path. - from shared.models.schemas.job_metadata import JobMetadataHelper - from shared.services.redis import RedisServiceFactory - - redis_service = RedisServiceFactory.get_service() - job_metadata = await job_repo.get_job_metadata(db, job_id, redis_service) - - # Result delivery fields. - job_result = job.job_result - result_url = None - result = None - result_url_expires_at = job.created_at # Default to created_at. - - if job_result and job_result.result_s3_key: - upload_service = FileUploadService() - result_url_info = cast( - Dict[str, Any], - await upload_service.generate_download_url(job_result.result_s3_key), - ) - result_url = result_url_info["download_url"] - expires_in = int(result_url_info["expires_in"]) - - # Read checksum/statistics data from inline_payload when present. - if job_result.inline_payload: - result = job_result.inline_payload - - # Compute result_url_expires_at when a download URL was issued. - if result_url: - from datetime import datetime, timedelta - - result_url_expires_at = datetime.now() + timedelta(seconds=expires_in) - - original_request = ( - job_metadata.get("original_request") - if isinstance(job_metadata, dict) - else {} - ) - source_url = ( - original_request.get("source_url") - if isinstance(original_request, dict) - else None - ) - file_name = None - if source_url: - parsed_source = urlparse(source_url) - file_name = os.path.basename(parsed_source.path) or None - if not file_name and isinstance(original_request, dict): - file_name = original_request.get("file_name") - file_extension = None - if file_name: - ext = os.path.splitext(file_name)[1] - file_extension = ext[1:].upper() if ext else None - - parsing_params = {} - if isinstance(original_request, dict): - parsing_params = original_request.get("parsing_params") or {} - if not parsing_params and isinstance(job_metadata, dict): - parsing_params = job_metadata.get("parsing_params") or {} - model = ( - parsing_params.get("model") if isinstance(parsing_params, dict) else None - ) - ocr_enabled = ( - parsing_params.get("ocr_enabled") - if isinstance(parsing_params, dict) - else None - ) - - response_data = JobResultResponse( - job_id=job.job_id, - namespace=JobMetadataHelper.get_field(job_metadata, "namespace"), - document_id=resolve_public_document_id(job), - status=status_for_api, - source_type=job.source_type, - data_id=JobMetadataHelper.get_field(job_metadata, "data_id"), - created_at=require_utc(job.created_at, field_name="created_at"), - progress=progress, - error=_build_error_response(job, job_metadata), - result=result, - result_url=result_url, - result_url_expires_at=require_utc( - result_url_expires_at, - field_name="result_url_expires_at", - ), - file_name=file_name, - file_extension=file_extension, - model=model, - ocr_enabled=ocr_enabled, - duration_seconds=( - (job.updated_at - job.created_at).total_seconds() - if job.updated_at and job.created_at - else None - ), - credits_spent=( - MicroDollar(job.credits_charged).to_credit() - if hasattr(job, "credits_charged") - else 0 - ), - ) - - return response_data - - except NotFoundException: - raise - except PermissionDeniedException: - raise - except Exception as e: - logger.error(f"Failed to get job result: {e}") - raise JobOperationException( - internal_message=f"Failed to get job result: {str(e)}" - ) + return await get_job_result_for_user( + db, + job_id=job_id, + user_id=current_user.user_id, + ) @router.post( @@ -1022,63 +116,9 @@ async def confirm_upload( """ Confirm a completed file upload as a fallback path. """ - try: - job_repo = JobRepository() - - # Load the job and verify access. - job = await job_repo.get_job_by_id(db, job_id) - check_job_permission(job, current_user.user_id, job_id) - assert job is not None - - # Check the current job state. - logger.info(f"Confirm upload - Job {job_id} current status: {job.status}") - if job.status not in [JobStatus.PENDING.value, JobStatus.WAITING_FILE.value]: - # If the webhook already advanced the job, return success idempotently. - logger.info(f"Job {job_id} already processed, status: {job.status}") - return {"message": "Job status already updated"} - - # Verify that the S3 object exists. - if not job.s3_key: - raise ValidationException( - user_message="Job is missing S3 key information", - violations=[ - {"field": "s3_key", "description": "S3 key not set for this job"} - ], - ) - - upload_service = FileUploadService() - file_info = await upload_service.verify_s3_file_exists(job.s3_key) - - if not file_info.get("exists"): - raise ValidationException( - user_message="S3 file does not exist, please upload the file first", - violations=[{"field": "file", "description": "File not found in S3"}], - ) - - # Advance the job state. - await transition_to_uploaded( - db, job_id, job.job_type, "manual_upload_completed" - ) - - # Start job processing. - await start_workflow_for_job( - db=db, - job_id=job_id, - job_type=job.job_type, - source_type="file", - user_id=current_user.user_id, - ) - - return {"message": "File upload confirmed; processing started"} - - except NotFoundException: - raise - except PermissionDeniedException: - raise - except ValidationException: - raise - except Exception as e: - logger.error(f"Failed to confirm upload: {e}") - raise JobOperationException( - internal_message=f"Failed to confirm upload: {str(e)}" - ) + return await _document_ingestion_service.confirm_upload( + db, + job_id=job_id, + request_payload=request, + user_id=current_user.user_id, + ) diff --git a/apps/api/app/api/v1/routes/qstash_callbacks.py b/apps/api/app/api/v1/routes/qstash_callbacks.py index 84cae94f2..e2fba04e7 100644 --- a/apps/api/app/api/v1/routes/qstash_callbacks.py +++ b/apps/api/app/api/v1/routes/qstash_callbacks.py @@ -1,262 +1,33 @@ -""" -QStash callback endpoints. - -These endpoints receive delivery status from Upstash QStash after it -delivers (or fails to deliver) a webhook to the customer's endpoint. - -Both endpoints verify the QStash JWT signature before processing. -""" +"""QStash callback endpoints.""" from __future__ import annotations -import json -from datetime import datetime, timezone -from typing import Any, Dict, Optional -from uuid import NAMESPACE_URL, uuid5 - +from app.services.webhook import qstash_callback_service +from app.services.webhook.qstash_callback_service import QStashCallbackOutcome from fastapi import APIRouter, Request, Response -from loguru import logger -from sqlalchemy import select - -from shared.core.config import app_config -from shared.core.database_sync import get_sync_db_context -from shared.models.database.webhook import WebhookEvent, WebhookEventStatus -from shared.models.database.webhook_log import WebhookLog router = APIRouter(tags=["QStash Callbacks"]) -def _get_qstash_verification_url(callback_path: str, request_url: str) -> str: - """Build the URL used for QStash signature verification. - - Prefer the configured public callback URL because ingress/TLS termination - can make ``request.url`` appear as an internal ``http://`` URL. - """ - callback_base_url = app_config.QSTASH_CALLBACK_BASE_URL - - if callback_base_url: - return f"{callback_base_url.rstrip('/')}{callback_path}" - - return request_url - - -def _verify_qstash_signature(raw_body: bytes, signature: str, url: str) -> bool: - """Verify the QStash JWT signature on an inbound callback.""" - current_key = app_config.QSTASH_CURRENT_SIGNING_KEY - next_key = app_config.QSTASH_NEXT_SIGNING_KEY - - if not current_key or not next_key: - logger.error("QStash signing keys not configured — rejecting callback") - return False - - try: - from qstash import Receiver - - receiver = Receiver( - current_signing_key=current_key, - next_signing_key=next_key, - ) - receiver.verify( - body=raw_body.decode("utf-8"), - signature=signature, - url=url, - ) - return True - except Exception as exc: - logger.warning( - "QStash signature verification failed: error_type={error_type}, url={url}", - error_type=type(exc).__name__, - url=url, - ) - return False - - -def _extract_callback_data(body: bytes) -> Dict[str, Any]: - """Parse the QStash callback body.""" - try: - return json.loads(body) - except (json.JSONDecodeError, ValueError): - return {"raw": body.decode("utf-8", errors="replace")} - - -def _normalize_header_value(value: Any) -> Optional[str]: - """Normalize a callback header value to a single string.""" - if isinstance(value, list): - if not value: - return None - first_value = value[0] - return first_value if isinstance(first_value, str) else str(first_value) - - if isinstance(value, str): - return value - - if value is None: - return None - - return str(value) - - -def _find_event_id(data: Dict[str, Any]) -> Optional[str]: - """Extract the Knowhere event ID from QStash sourceHeader.""" - source_header = data.get("sourceHeader", {}) or {} - event_id = _normalize_header_value( - source_header.get("X-Knowhere-Event-Id") - or source_header.get("x-knowhere-event-id") - ) - if not event_id: - for key, value in source_header.items(): - if key.lower() == "x-knowhere-event-id": - event_id = _normalize_header_value(value) - break - return event_id - - -def _build_callback_log_idempotency_key( - qstash_message_id: Optional[str], - event_id: str, -) -> str: - """Build a fixed-width idempotency key for webhook_logs. - - ``webhook_logs.idempotency_key`` is limited to 36 characters. QStash - ``sourceMessageId`` is longer, so store the raw value in - ``qstash_message_id`` and derive a stable UUID from it for the - idempotency key column. - """ - if qstash_message_id: - return str(uuid5(NAMESPACE_URL, qstash_message_id)) - - return event_id - - -def _get_response_status_code(value: Any) -> Optional[int]: - """Return the destination response status reported by QStash.""" - if value is None: - return None - - try: - return int(value) - except (TypeError, ValueError): - return None - - -def _is_success_response_status(status_code: Optional[int]) -> bool: - """Return whether a destination response status is successful.""" - return status_code is not None and 200 <= status_code < 300 - - -def _get_callback_event_status(data: Dict[str, Any]) -> str: - """Map a normal QStash callback to the current webhook event status.""" - response_status = _get_response_status_code(data.get("status")) - if _is_success_response_status(response_status): - return WebhookEventStatus.DELIVERED - - return WebhookEventStatus.DELIVERING - - -def _resolve_event_status(current_status: str, callback_status: str) -> str: - """Apply callback status without downgrading terminal delivery state.""" - if current_status in ( - WebhookEventStatus.DELIVERED, - WebhookEventStatus.FAILED, - WebhookEventStatus.CANCELED, - ): - return current_status - - return callback_status - - -def _process_qstash_callback( - data: Dict[str, Any], - event_id: str, - callback_status: str, - log_label: str, -) -> Response: - """Shared logic for both success and failure QStash callbacks. - - Fetches the WebhookEvent, updates its status, and writes a WebhookLog entry. - """ - response_status_code = _get_response_status_code(data.get("status")) - response_body = data.get("body", "") - qstash_message_id = data.get("sourceMessageId") - retried = data.get("retried", 0) - is_failed_delivery_attempt = ( - callback_status == WebhookEventStatus.FAILED - or ( - callback_status == WebhookEventStatus.DELIVERING - and not _is_success_response_status(response_status_code) - ) - ) - error_message = None - if is_failed_delivery_attempt: - error_message = data.get("error") or response_body - - with get_sync_db_context() as db: - event = db.execute( - select(WebhookEvent).where(WebhookEvent.id == event_id) - ).scalar_one_or_none() - - if not event: - logger.warning(f"QStash {log_label}: event {event_id} not found in DB") - return Response(status_code=200, content="OK (event not found)") - - now = datetime.now(timezone.utc).replace(tzinfo=None) - event_status = _resolve_event_status(event.status, callback_status) - attempt_number = retried + 1 - event.status = event_status - event.attempts = max(event.attempts, attempt_number) - event.updated_at = now - - log = WebhookLog( - job_id=event.job_id, - event_id=event.id, - webhook_url=event.target_url, - attempt_number=attempt_number, - request_payload=event.payload, - signature="", - idempotency_key=_build_callback_log_idempotency_key( - qstash_message_id, event.id - ), - response_status_code=response_status_code, - response_body=response_body[:4096] if response_body else None, - error_message=str(error_message)[:4096] if error_message else None, - duration_ms=0, - qstash_message_id=qstash_message_id, - ) - db.add(log) - db.commit() - - return Response(status_code=200, content="OK") - - @router.post("/qstash/callback") async def handle_qstash_callback(request: Request) -> Response: """Handle QStash success callback after webhook delivery.""" raw_body = await request.body() signature = request.headers.get("upstash-signature", "") - verification_url = _get_qstash_verification_url( + verification_url = qstash_callback_service.get_qstash_verification_url( "/webhooks/qstash/callback", str(request.url), ) - if not _verify_qstash_signature(raw_body, signature, verification_url): + if not qstash_callback_service.verify_qstash_signature( + raw_body, + signature, + verification_url, + ): return Response(status_code=401, content="Invalid signature") - data = _extract_callback_data(raw_body) - event_id = _find_event_id(data) - - if not event_id: - logger.warning("QStash callback: missing event_id, cannot correlate") - return Response(status_code=200, content="OK (no event_id)") - - retried = data.get("retried", 0) - logger.info( - f"QStash callback: event_id={event_id}, status={data.get('status')}, " - f"retried={retried}, qstash_message_id={data.get('sourceMessageId')}" - ) - - event_status = _get_callback_event_status(data) - return _process_qstash_callback( - data, event_id, event_status, "callback" + return _to_response( + qstash_callback_service.handle_qstash_success_callback(raw_body) ) @@ -265,28 +36,27 @@ async def handle_qstash_failure(request: Request) -> Response: """Handle QStash failure callback after all retries exhausted.""" raw_body = await request.body() signature = request.headers.get("upstash-signature", "") - verification_url = _get_qstash_verification_url( + verification_url = qstash_callback_service.get_qstash_verification_url( "/webhooks/qstash/failure", str(request.url), ) - if not _verify_qstash_signature(raw_body, signature, verification_url): + if not qstash_callback_service.verify_qstash_signature( + raw_body, + signature, + verification_url, + ): return Response(status_code=401, content="Invalid signature") - data = _extract_callback_data(raw_body) - event_id = _find_event_id(data) - - if not event_id: - logger.warning("QStash failure callback: missing event_id, cannot correlate") - return Response(status_code=200, content="OK (no event_id)") - - retried = data.get("retried", 0) - max_retries = data.get("maxRetries", 0) - logger.warning( - f"QStash failure: event_id={event_id}, status={data.get('status')}, " - f"retried={retried}/{max_retries}, qstash_message_id={data.get('sourceMessageId')}" + return _to_response( + qstash_callback_service.handle_qstash_failure_callback(raw_body) ) - return _process_qstash_callback( - data, event_id, WebhookEventStatus.FAILED, "failure" - ) + +def _to_response(outcome: QStashCallbackOutcome) -> Response: + content_by_kind = { + "processed": "OK", + "missing_event_id": "OK (no event_id)", + "event_not_found": "OK (event not found)", + } + return Response(status_code=200, content=content_by_kind[outcome.kind]) diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index 1ef332e4b..5db937130 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -4,13 +4,15 @@ from typing import Literal -from app.services.rate_limit.dependencies import CurrentUser, with_current_user +from app.api.dependencies.current_user import with_current_user +from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends from pydantic import BaseModel, Field, field_validator from sqlalchemy.ext.asyncio import AsyncSession from shared.core.database import get_db -from shared.services.retrieval import run_retrieval_query +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace +from shared.services.retrieval.app_service import run_retrieval_query router = APIRouter(tags=["Retrieval"]) @@ -22,7 +24,9 @@ class ExcludeSection(BaseModel): class RetrievalQueryRequest(BaseModel): namespace: str | None = Field( - None, description="Effective namespace; defaults to default" + None, + max_length=255, + description="Effective namespace; defaults to default", ) query: str top_k: int = 10 @@ -66,6 +70,11 @@ def validate_channels(cls, v: list[str]) -> list[str]: raise ValueError(f"Invalid channel: {ch}. Must be one of {valid}") return v + @field_validator("namespace") + @classmethod + def normalize_namespace(cls, namespace: str | None) -> str: + return normalize_retrieval_namespace(namespace) + class RetrievalQueryResponse(BaseModel): namespace: str @@ -85,7 +94,7 @@ async def query_retrieval( return await run_retrieval_query( db=db, user_id=current_user.user_id, - namespace=payload.namespace or "default", + namespace=normalize_retrieval_namespace(payload.namespace), query=payload.query, top_k=payload.top_k, exclude_document_ids=payload.exclude_document_ids, diff --git a/apps/api/app/api/v1/routes/s3_events.py b/apps/api/app/api/v1/routes/s3_events.py index ee92dec6a..ea3feff24 100644 --- a/apps/api/app/api/v1/routes/s3_events.py +++ b/apps/api/app/api/v1/routes/s3_events.py @@ -1,130 +1,14 @@ -""" -S3 event webhook routes. -""" +"""S3 event webhook routes.""" -import base64 -import json -import os -from typing import Any, Dict - -from app.repositories.job_repository import JobRepository -from app.services.knowledge.kb_orchestrator import KBOrchestrator -from app.services.state_machine import JobStateMachine +from app.services.s3_events.service import safely_handle_s3_event_post +from app.services.s3_events.intake_outcome import sanitize_storage_event_headers from fastapi import APIRouter, Header, Request from loguru import logger -from shared.core.database import get_db_context from shared.core.logging import LogEvent -from shared.core.state_machine.states import JobStatus -from shared.models.schemas.oss_event import OSSEvent -from shared.models.schemas.s3_event import S3Event -from shared.utils.pinned_outbound_http import ( - send_pinned_outbound_request, -) -from shared.utils.url_security import ( - validate_http_url_and_resolve_ip_async, -) router = APIRouter(tags=["Internal"]) -SNS_SUBSCRIPTION_TIMEOUT_SECONDS = 10 - - -def verify_sns_signature(request_body: bytes, signature: str, message: str) -> bool: - """ - Validate an SNS message signature. - - Args: - request_body: Request body. - signature: Signature header value. - message: Message payload. - - Returns: - bool: Whether validation succeeded. - """ - try: - # This is intentionally simplified. Production code should use the AWS SDK. - return True - except Exception as e: - logger.error(f"SNS signature verification failed: {e}") - return False - - -def verify_minio_signature(auth_token: str, expected_token: str) -> bool: - """ - Validate the MinIO webhook token. - - Args: - auth_token: Token supplied by the request. - expected_token: Token configured on the server. - - Returns: - bool: Whether validation succeeded. - """ - if not expected_token: - return True # Skip verification when no token is configured. - - return auth_token == expected_token - - -def verify_oss_signature(request_body: bytes, headers: Dict[str, str]) -> bool: - """ - Validate an OSS callback signature. - - Args: - request_body: Request body. - headers: Request headers. - - Returns: - bool: Whether validation succeeded. - """ - try: - from shared.core.config import settings - - # Allow an opt-out for local development and controlled environments. - if not getattr(settings, "OSS_EVENT_VERIFY_SIGNATURE", True): - return True - - # OSS callback verification is simplified here. Production code should - # follow the official OSS callback verification flow. - callback_key = getattr(settings, "OSS_EVENT_CALLBACK_KEY", "") - if not callback_key: - logger.warning( - "OSS_EVENT_CALLBACK_KEY is not configured; skipping signature verification" - ) - return True - - # TODO: Implement OSS callback signature verification. - # Expected steps: - # 1. Read the signature metadata from the headers. - # 2. Compute the signature with callback_key. - # 3. Compare the computed and provided signatures. - - return True - except Exception as e: - logger.error(f"OSS signature verification failed: {e}") - return False - - -def extract_job_id_from_s3_key(s3_key: str) -> str | None: - """ - Extract the job_id from an S3 object key. - - Args: - s3_key: S3 key in the format uploads/{job_id}.ext. - - Returns: - str: Job identifier. - """ - if not s3_key.startswith("uploads/"): - return None - - # Strip the uploads/ prefix and remove the file extension. - filename = s3_key[8:] # Remove the "uploads/" prefix. - job_id = os.path.splitext(filename)[0] - - return job_id - @router.get("/s3-events", response_model=dict, summary="Handle S3 webhook GET requests") async def handle_s3_events_get( @@ -132,16 +16,13 @@ async def handle_s3_events_get( x_amz_sns_message_type: str = Header(None, alias="x-amz-sns-message-type"), x_minio_auth_token: str = Header(None, alias="x-minio-auth-token"), authorization: str = Header(None), -): - """ - Handle S3-event GET requests, primarily for SNS subscription confirmation. - """ +) -> dict[str, str]: + """Handle S3-event GET requests, primarily for SNS subscription confirmation.""" logger.info("======== S3 event GET request ========") logger.info(f"Headers: {dict(request.headers)}") if request.client: logger.info(f"Client IP: {request.client.host}") - # Handle SNS subscription confirmation requests. if x_amz_sns_message_type == "SubscriptionConfirmation": logger.info("Received an SNS subscription confirmation request") return {"message": "SNS subscription confirmed"} @@ -157,411 +38,17 @@ async def handle_s3_events( x_amz_sns_message_type: str = Header(None, alias="x-amz-sns-message-type"), x_minio_auth_token: str = Header(None, alias="x-minio-auth-token"), authorization: str = Header(None), -): - """ - Handle S3-event POST requests from AWS SNS, MinIO, or OSS. - """ +) -> dict[str, str]: + """Handle S3-event POST requests from AWS SNS, MinIO, OSS, or tests.""" logger.bind(event=LogEvent.S3_WEBHOOK_EVENT).info( - f"S3 event Headers: {dict(request.headers)}" + f"S3 event Headers: {sanitize_storage_event_headers(dict(request.headers))}" ) if request.client: logger.info(f"Client IP: {request.client.host}") - try: - # Read the request body. - body = await request.body() - headers = dict(request.headers) - - # Determine the event source. - if x_amz_sns_message_type: - # AWS SNS event. - result = await handle_sns_event(body) - if result: - return result - elif _is_oss_event(headers): - # OSS event, including Aliyun MNS proxy notifications. - await handle_oss_event(body, headers) - elif x_minio_auth_token: - # MinIO event, identified by the dedicated x-minio-auth-token header. - await handle_minio_event(body, x_minio_auth_token) - else: - # Direct S3 event payload used in tests. - await handle_direct_s3_event(body) - - return {"message": "Event handled successfully"} - - except Exception as e: - logger.error(f"Failed to handle S3 event: {e}") - # Return 200 even on failure so the upstream storage service does not retry blindly. - return {"message": "Event handling completed"} - - -async def handle_sns_event(body: bytes): - """ - Handle an AWS SNS event payload. - """ - try: - # Parse the SNS message envelope. - sns_message = json.loads(body.decode("utf-8")) - - # Branch on the SNS message type. - message_type = sns_message.get("Type") - logger.info(f"SNS message type: {message_type}") - - if message_type == "SubscriptionConfirmation": - # Handle subscription confirmation. - logger.info("Received an SNS subscription confirmation request") - subscribe_url = sns_message.get("SubscribeURL") - if subscribe_url: - logger.info(f"SNS subscription confirmation URL: {subscribe_url}") - # Visit the URL to confirm the subscription. - return await confirm_sns_subscription(subscribe_url) - else: - logger.warning( - "SNS subscription confirmation did not include SubscribeURL" - ) - return {"message": "SNS subscription confirmation failed"} - - elif message_type == "Notification": - # Handle notification messages. - logger.info("Received an SNS notification") - logger.info(f"SNS message payload: {sns_message}") - - # Parse the embedded S3 event. - try: - s3_event_data = json.loads(sns_message["Message"]) - logger.info(f"S3 event payload: {s3_event_data}") - - # Skip S3 test events — AWS/LocalStack sends these when - # bucket notification configuration is first applied. - # They lack the standard Records[] structure. - if ( - isinstance(s3_event_data, dict) - and s3_event_data.get("Event") == "s3:TestEvent" - ): - logger.info("Skip S3 test event") - return {"message": "S3 test event confirmed and skipped"} - s3_event = S3Event(**s3_event_data) - - # Process the upload events. - await process_upload_events(s3_event) - except Exception as e: - logger.error(f"Failed to parse the S3 event payload: {e}") - logger.error(f"SNS payload: {sns_message}") - # Fall back to treating the SNS payload itself as an S3 event. - try: - s3_event = S3Event(**sns_message) - await process_upload_events(s3_event) - except Exception as e2: - logger.error( - f"Fallback parsing of the SNS payload as an S3 event also failed: {e2}" - ) - raise - else: - logger.warning(f"Unknown SNS message type: {message_type}") - return {"message": f"Unknown SNS message type: {message_type}"} - - except Exception as e: - logger.error(f"Failed to handle SNS event: {e}") - raise - - -async def confirm_sns_subscription(subscribe_url: str) -> dict[str, str]: - """Confirm an SNS subscription after SSRF validation and IP pinning.""" - validation = await validate_http_url_and_resolve_ip_async( - subscribe_url, + return await safely_handle_s3_event_post( + body=await request.body(), + headers=dict(request.headers), + sns_message_type=x_amz_sns_message_type, + minio_auth_token=x_minio_auth_token, ) - if not validation.is_valid: - logger.warning( - f"SNS subscription confirmation URL failed validation: {validation.error_message}" - ) - return {"message": "SNS subscription confirmation failed"} - - if not validation.validated_ip: - logger.warning("SNS subscription confirmation URL validation returned no IP") - return {"message": "SNS subscription confirmation failed"} - - try: - response = await send_pinned_outbound_request( - method="GET", - url=subscribe_url, - pinned_ip=validation.validated_ip, - timeout_seconds=SNS_SUBSCRIPTION_TIMEOUT_SECONDS, - ) - if response.status == 200: - logger.info("SNS subscription confirmed successfully") - return {"message": "SNS subscription confirmed"} - - if 300 <= response.status < 400: - logger.warning( - f"SNS subscription confirmation redirect blocked, status={response.status}" - ) - else: - logger.error( - f"SNS subscription confirmation failed, status={response.status}" - ) - return {"message": "SNS subscription confirmation failed"} - except Exception as e: - logger.error(f"Failed to reach the SNS confirmation URL: {e}") - return {"message": "SNS subscription confirmation failed"} - - -async def handle_minio_event(body: bytes, auth_token: str): - """ - Handle a MinIO webhook event. - """ - try: - # Validate the webhook token. - from shared.core.config import settings - - expected_token = getattr(settings, "S3_WEBHOOK_AUTH_TOKEN", "") - - if not verify_minio_signature(auth_token, expected_token): - logger.warning("MinIO webhook authentication failed") - return - - # Parse the S3 event payload. - s3_event_data = json.loads(body.decode("utf-8")) - s3_event = S3Event(**s3_event_data) - - # Process the upload events. - await process_upload_events(s3_event) - - except Exception as e: - logger.error(f"Failed to handle MinIO event: {e}") - - -async def handle_direct_s3_event(body: bytes): - """ - Handle a direct S3 event payload used in tests. - """ - try: - # Parse the S3 event payload. - s3_event_data = json.loads(body.decode("utf-8")) - s3_event = S3Event(**s3_event_data) - - # Process the upload events. - await process_upload_events(s3_event) - - except Exception as e: - logger.error(f"Failed to handle direct S3 event: {e}") - - -def _is_oss_event(headers: Dict[str, str]) -> bool: - """ - Return whether the incoming request looks like an OSS event. - - Args: - headers: Request headers. - - Returns: - bool: Whether the request matches OSS event heuristics. - """ - # Identify OSS events by storage type, known headers, or request shape. - - storage_type = os.getenv("S3_TYPE", "s3").lower() - if storage_type == "oss": - return True - - # Also look for OSS-specific headers such as x-oss-pub-key-url. - if "x-oss-pub-key-url" in headers: - return True - - # Recognize Aliyun MNS proxy headers and user agents. - if "x-mns-version" in headers or "x-mns-signing-cert-url" in headers: - return True - user_agent = headers.get("user-agent") or headers.get("User-Agent") - if user_agent and "Aliyun Notification Service Agent" in user_agent: - return True - - return False - - -async def handle_oss_event(body: bytes, headers: Dict[str, str]): - """ - Handle an OSS event payload. - """ - try: - # Verify the callback signature. - if not verify_oss_signature(body, headers): - logger.warning("OSS event signature verification failed") - return - - # Parse the OSS payload, including MNS wrapper envelopes. - event_data = json.loads(body.decode("utf-8")) - logger.info(f"OSS event payload: {event_data}") - # MNS may place the real event inside Message as base64 or raw JSON. - if isinstance(event_data, dict) and "Message" in event_data: - inner = event_data.get("Message") - if isinstance(inner, str): - decoded = None - # Prefer base64 decoding first. - try: - decoded_bytes = base64.b64decode(inner, validate=True) - decoded_str = decoded_bytes.decode("utf-8") - decoded = json.loads(decoded_str) - except Exception: - decoded = None - - if decoded is None: - # Fall back to parsing the raw JSON string directly. - try: - decoded = json.loads(inner) - except Exception: - decoded = None - - if decoded is not None: - event_data = decoded - logger.info(f"Decoded MNS Message payload: {event_data}") - elif isinstance(inner, dict): - event_data = inner - - # Detect the payload shape. - if "events" in event_data: - # Standard OSS event format. - oss_event = OSSEvent(**event_data) - elif "Records" in event_data: - # Compatibility path for S3-like payloads emitted by OSS. - oss_event = _convert_s3_format_to_oss(event_data) - else: - logger.error(f"Unknown OSS event format: {event_data}") - return - - # Convert to S3Event so the existing upload flow can be reused. - s3_event = oss_event.to_s3_event() - - # Process the upload events. - await process_upload_events(s3_event) - - except Exception as e: - logger.error(f"Failed to handle OSS event: {e}") - raise - - -def _convert_s3_format_to_oss(event_data: Dict[str, Any]) -> OSSEvent: - """ - Convert an S3-style event payload into an OSS event payload. - - Args: - event_data: S3-format event data. - - Returns: - OSSEvent: OSS event object. - """ - from shared.models.schemas.oss_event import OSSEventRecord - - # Convert each S3-style record into the OSS schema. - records = event_data.get("Records", []) - oss_records = [] - - for record in records: - oss_record = OSSEventRecord( - eventName=record.get("eventName", "").replace("s3:", ""), - eventSource="acs:oss", - eventTime=record.get("eventTime", ""), - region=record.get("awsRegion", ""), - oss={ - "bucket": record.get("s3", {}).get("bucket", {}), - "object": record.get("s3", {}).get("object", {}), - }, - ) - oss_records.append(oss_record) - - return OSSEvent(events=oss_records) - - -async def process_upload_events(s3_event: S3Event): - """ - Process upload events delivered by S3-compatible storage. - - Args: - s3_event: S3 event object. - """ - try: - # Gather only upload-related records. - upload_events = s3_event.get_upload_events() - - # Instantiate services once outside the loop - job_repo = JobRepository() - - for event in upload_events: - # Read the object key from the event record. - s3_key = event.object_key or event.s3.get("object", {}).get("key") - if not s3_key: - continue - - # Extract the job_id from the object key. - job_id = extract_job_id_from_s3_key(s3_key) - if not job_id: - logger.warning(f"Could not extract job_id from S3 key: {s3_key}") - continue - - logger.info(f"Processing S3 upload event: {s3_key} -> job_id={job_id}") - - # Load the matching job. - async with get_db_context() as db: - job = await job_repo.get_job_by_id(db, job_id) - - if not job: - logger.warning(f"No job found for upload event: {job_id}") - continue - - # Only react while the job is still waiting for file upload. - if job.status != "waiting-file": - logger.info( - f"Job {job_id} is not in waiting-file status: {job.status}" - ) - continue - - # Check if upload window has expired (race-condition safe via optimistic lock) - from shared.core.config import settings - from shared.core.state_machine.states import is_job_expired - - if is_job_expired(job.updated_at, settings.JOB_WAITING_EXPIRE_SECONDS): - logger.warning(f"Job {job_id} upload expired, marking failed") - state_machine = JobStateMachine() - await state_machine.mark_failed( - db, - job_id, - "Upload expired: file was not uploaded within the allowed time window", - error_code="UPLOAD_EXPIRED", - ) - continue - - # Skip S3 file verification — we are processing the upload - # notification itself, so the file is guaranteed to exist. - - # Advance the job state. - state_machine = JobStateMachine() - - # Once upload is complete, move the job to pending. - await state_machine.transition( - db, - job_id, - JobStatus.PENDING.value, - "s3_upload_completed", - None, - "system", - ) - - # Start job processing. - if job.job_type == "kb_management": - orchestrator = KBOrchestrator() - await orchestrator.start_workflow( - db=db, - job_id=job_id, - source_type="file", - file_path=None, - file_url=None, - user_id=str(job.user_id), - ) - else: - logger.warning( - f"Unsupported job type for upload event: {job.job_type}, job_id={job_id}" - ) - - logger.info(f"Triggered processing for job {job_id}") - - except Exception as e: - logger.error(f"Failed to process upload events: {e}") - raise diff --git a/apps/api/app/api/v1/routes/webhook.py b/apps/api/app/api/v1/routes/webhook.py index 18ba4abc8..19eb84f30 100644 --- a/apps/api/app/api/v1/routes/webhook.py +++ b/apps/api/app/api/v1/routes/webhook.py @@ -9,7 +9,8 @@ from app.repositories.job_repository import JobRepository from app.repositories.webhook_repository import WebhookRepository -from app.services.rate_limit.dependencies import CurrentUser, with_current_user +from app.api.dependencies.current_user import with_current_user +from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends, Query from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession @@ -160,24 +161,17 @@ async def trigger_webhook( ), ) - # Use dispatcher to send synchronously dispatcher = get_webhook_dispatcher() - # Pass db session and is_manual=True to handle logging internally - ( - success, - status_code, - duration_ms, - error_message, - ) = await dispatcher._send_webhook(db=db, event=event, is_manual=True) + delivery_result = await dispatcher.send_manual_webhook(db=db, event=event) # 5. Return response return WebhookTriggerResponse( - success=success, - status_code=status_code, + success=delivery_result.success, + status_code=delivery_result.status_code, response_body=None, # Dispatcher doesn't return response body - duration_ms=duration_ms, + duration_ms=delivery_result.duration_ms, delivery_id=None, # Manual trigger doesn't create delivery log - error_message=error_message, + error_message=delivery_result.error_message, ) except KnowhereException: diff --git a/apps/api/app/api/v1/routes/webhook_secrets.py b/apps/api/app/api/v1/routes/webhook_secrets.py index c0f63bf44..0fb4f9bf2 100644 --- a/apps/api/app/api/v1/routes/webhook_secrets.py +++ b/apps/api/app/api/v1/routes/webhook_secrets.py @@ -6,7 +6,8 @@ from typing import List, Optional -from app.services.rate_limit.dependencies import CurrentUser, with_current_user +from app.api.dependencies.current_user import with_current_user +from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends from pydantic import BaseModel, Field from sqlalchemy.ext.asyncio import AsyncSession diff --git a/apps/api/app/core/dependencies.py b/apps/api/app/core/dependencies.py deleted file mode 100644 index bf56d39c8..000000000 --- a/apps/api/app/core/dependencies.py +++ /dev/null @@ -1,158 +0,0 @@ -import threading -from datetime import timedelta -from typing import Any - -import jwt -from app.services.auth.api_key_service import APIKeyService -from fastapi import Depends, Header, Request -from jwt import PyJWKClient -from loguru import logger -from sqlalchemy import select -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.config import settings -from shared.core.database import get_db -from shared.core.exceptions.domain_exceptions import ( - AuthException, -) -from shared.models.database.user import User -from shared.utils.api_keys import is_api_key_token - -# Standard JWKS endpoint path (fixed, following OpenID Connect convention) -JWKS_ENDPOINT_PATH = "/api/auth/jwks" - -# Cache settings: 1 hour in seconds -JWKS_CACHE_TTL_SECONDS = 60 * 60 # 3600 seconds - -# Cached PyJWKClient instance -_jwks_client: PyJWKClient | None = None -_jwks_client_lock = threading.Lock() - - -def _get_jwks_client() -> PyJWKClient: - """ - Get or create a cached PyJWKClient instance. - - The JWKS endpoint is constructed from INTERNAL_DASHBOARD_ENDPOINT + fixed path. - PyJWKClient caches the JWKS response for 1 hour. - """ - global _jwks_client - - if _jwks_client is None: - with _jwks_client_lock: - if _jwks_client is None: - jwks_url = f"{settings.INTERNAL_DASHBOARD_ENDPOINT}{JWKS_ENDPOINT_PATH}" - _jwks_client = PyJWKClient( - jwks_url, - cache_jwk_set=True, - lifespan=JWKS_CACHE_TTL_SECONDS, - timeout=30, - ) - logger.info(f"Initialized JWKS client with endpoint: {jwks_url}") - - return _jwks_client - - -def _get_verification_key(token: str) -> Any: - """ - Get the verification key for the JWT from the JWKS endpoint. - - Uses PyJWKClient's built-in cache with 1 hour TTL. - """ - try: - jwks_client = _get_jwks_client() - signing_key = jwks_client.get_signing_key_from_jwt(token) - return signing_key.key - except jwt.PyJWKClientError as e: - logger.error(f"Failed to fetch JWKS: {e}") - raise AuthException( - internal_message=f"Failed to fetch verification key from JWKS endpoint: {e}" - ) - except jwt.PyJWKSetError as e: - logger.error(f"Invalid JWKS format: {e}") - raise AuthException(internal_message=f"Invalid JWKS format: {e}") - - -def decode_jwt_token(token: str) -> str: - """ - Decode and validate JWT token using JWKS. - - Expected JWT claims: - - id/sub: User ID - - exp: Expiration - """ - try: - key = _get_verification_key(token) - - # Decode and Verify - payload = jwt.decode( - token, - key, - algorithms=["HS256", "RS256", "EdDSA"], - leeway=timedelta(seconds=30), - options={"verify_aud": False}, - ) - - # Extract user_id from 'id' claim - user_id = payload.get("id") - - if not user_id: - raise AuthException(user_message="Token missing 'id' claim") - - return user_id - - except jwt.ExpiredSignatureError: - raise AuthException(user_message="Token has expired") - except jwt.InvalidTokenError as e: - logger.warning(f"Invalid JWT token: {e}") - raise AuthException(user_message="Invalid token") - - -async def _ensure_authenticated_user_exists( - db: AsyncSession, - user_id: str, -) -> None: - result = await db.execute(select(User.id).where(User.id == user_id).limit(1)) - if result.scalar_one_or_none() is not None: - return - - raise AuthException( - user_message="Invalid authentication credentials", - internal_message=( - "Authenticated user id is not present in the user table: " - f"user_id={user_id}" - ), - ) - - -async def get_current_user_id( - request: Request, - authorization: str | None = Header( - default=None, description="Bearer OR internal signature auth" - ), - db: AsyncSession = Depends(get_db), -) -> str: - """Authenticate the caller and return user_id.""" - if not authorization: - raise AuthException( - user_message="Authentication required. Provide Authorization header." - ) - - # Parse Authorization header - scheme, _, token = authorization.partition(" ") - if scheme.lower() != "bearer" or not token: - raise AuthException(user_message="Invalid Authorization header format") - - # Mode 1: API Key verification (for external clients) - if is_api_key_token(token): - api_key_service = APIKeyService.get_instance() - user_id = await api_key_service.validate_api_key(db, token) - if user_id: - return user_id - - raise AuthException(user_message="Invalid API Key") - - # Mode 2: JWT verification (for Dashboard/Internal) - user_id = decode_jwt_token(token) - await _ensure_authenticated_user_exists(db, user_id) - return user_id diff --git a/apps/api/app/core/image_cli.py b/apps/api/app/core/image_cli.py deleted file mode 100644 index 522ca99bb..000000000 --- a/apps/api/app/core/image_cli.py +++ /dev/null @@ -1,20 +0,0 @@ -from typing import Optional - -import httpx - - -class ImageCli: - """Image-processing client.""" - - http_client: Optional[httpx.AsyncClient] = None - - -http_client: Optional[httpx.AsyncClient] = None - - -def get_http_client() -> httpx.AsyncClient: - if http_client is None: - raise RuntimeError( - "HTTP client has not been initialized. Is it in the lifespan manager?" - ) - return http_client diff --git a/apps/api/app/core/response/ResponseCode.py b/apps/api/app/core/response/ResponseCode.py index 55bb13358..85ee75d9c 100644 --- a/apps/api/app/core/response/ResponseCode.py +++ b/apps/api/app/core/response/ResponseCode.py @@ -28,34 +28,3 @@ def get_all_as_dict(cls) -> Dict[int, str]: """Return all response codes and messages as a dictionary.""" return {member.code: member.msg for member in cls} - - -IS_TEST_MODE = False - - -if __name__ == "__main__": - if IS_TEST_MODE: - # 1. Access an enum member. - success_code = ResponseCode.SUCCESS - print(f"Member: {success_code}") - # Output: Member: ResponseCode.SUCCESS - - # 2. Access member attributes (code and msg). - print(f"Code: {success_code.code}, Message: {success_code.msg}") - # Output: Code: 200, Message: Operation succeeded - - fail_code = ResponseCode.FAIL - print(f"Code: {fail_code.code}, Message: {fail_code.msg}") - # Output: Code: 1, Message: Operation failed - - # 3. Iterate over all enum members. - print("\n--- All Response Codes ---") - for member in ResponseCode: - print(f"{member.name}: code={member.code}, msg='{member.msg}'") - - # 4. Call the class method to get a dictionary. - all_messages = ResponseCode.get_all_as_dict() - print("\n--- Dictionary Form ---") - import json - - print(json.dumps(all_messages, indent=2, ensure_ascii=False)) diff --git a/apps/api/app/data/demo_documents/tsla-q4-2025/chunks.json b/apps/api/app/data/demo_documents/tsla-q4-2025/chunks.json index 7e7147254..177ae25c9 100644 --- a/apps/api/app/data/demo_documents/tsla-q4-2025/chunks.json +++ b/apps/api/app/data/demo_documents/tsla-q4-2025/chunks.json @@ -25,7 +25,7 @@ "chunk_id": "68a6be7d-c587-5c73-abf2-56f4686e28e6", "type": "text", "content": "[tables/table-0 Tesla 2025 Results.html]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->HIGHLIGHTS", + "path": "TSLA-Q4-2025-Update.pdf-->HIGHLIGHTS", "metadata": { "length": 83, "summary": "", @@ -138,7 +138,7 @@ "chunk_id": "60109008-6261-51e6-b202-093d904eb881", "type": "text", "content": "FINANCIAL SUMMARY\n(Unaudited)\n\n[tables/table-1 Q4 2025 Financials.html]\n\n(1) As a result of the adoption of the new crypto assets standard, the previously reported quarterly periods in 2024 have been recast.\n(2) Beginning in Q1'25, Adjusted EBITDA (non-GAAP) is presented net of digital assets gains and losses and all prior periods have been adjusted.\n(3) Beginning in Q1'25, Net income attributable to common stockholders (non-GAAP) is presented net of digital assets gains and losses and all prior periods have been adjusted.\n(4) Beginning in Q1'25, Capital expenditures is presented inclusive of purchases of energy generation and storage systems and all prior periods have been adjusted.\nFINANCIAL SUMMARY\n(Unaudited)\n\n[tables/table-2 Financial Data 2021-25.html]\n\n(1) Beginning in Q1'25, Adjusted EBITDA (non-GAAP) is presented net of digital assets gains and losses and all prior periods have been adjusted.\n(2) Beginning in Q1'25, Net income attributable to common stockholders (non-GAAP) is presented net of digital assets gains and losses and all prior periods have been adjusted.\n(3) Beginning in Q1'25, Capital expenditures is presented inclusive of purchases of energy generation and storage systems and all prior periods have been adjusted.\nOPERATIONAL SUMMARY\n(Unaudited)\n\n[tables/table-3 Tesla Q4-2025 Data.html]\n\nOPERATIONAL SUMMARY\n(Unaudited)\n\n[tables/table-4 Tesla 2021-2025 Data.html]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY", "metadata": { "length": 1517, "summary": "The document presents unaudited financial and operational summaries for a company, covering quarterly data from Q4 2024 through Q4 2025 and annual data from 2021 to 2025. Key notes indicate significant accounting changes effective Q1 2025: Adjusted EBITDA and Net income attributable to common stockholders are now presented net of digital assets gains and losses, with all prior periods adjusted accordingly. Additionally, Capital expenditures now include purchases of energy generation and storage systems, requiring restatement of previous periods. The content references multiple tables detailing these metrics but does not display the specific numerical values.", @@ -240,7 +240,7 @@ "chunk_id": "60339310-5480-5ae2-8791-e6017bafb730", "type": "text", "content": "While automotive sales declined sequentially, gross margin (even when excluding the impact of regulatory credits) improved. The APAC region continued to show strength across multiple markets and set a record for deliveries in the quarter. We continued the rollout of Model Y variants across markets in Q4, including the standard and performance versions.\nPreparations continue in North America for the production ramps of Tesla Semi and Cybercab, both commencing 1H26, and production of the next-generation Roadster.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->Automotive", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->Automotive", "metadata": { "length": 516, "summary": "", @@ -304,7 +304,7 @@ "chunk_id": "eb18e33e-b093-532e-b4df-bd05d63b0294", "type": "text", "content": "We achieved our highest quarterly energy storage deployments, driven by record Megapack deployments. Total gross profit rose, both sequentially and year-over-year, to a record \\$1.1 billion, marking the fifth consecutive record quarter. We plan to begin Megapack 3 and Megablock production at Megafactory Houston in 2026. In 2025, our global Powerwall network supported more than 89,000 Virtual Power Plant events across over 1 million installed units, allowing homeowners to save over \\$1 billion in electricity bills as Virtual Power Plant participation continues to scale rapidly.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->Energy generation and storage", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->Energy generation and storage", "metadata": { "length": 583, "summary": "", @@ -395,7 +395,7 @@ "chunk_id": "60518074-bcbe-5b48-89a9-dfeb5d1b531b", "type": "text", "content": "We made further progress on the Optimus program in 2025. In Q1 of this year, we plan to unveil the Gen 3 version of Optimus, which will include major upgrades from version 2.5, including our latest hand design. The Gen 3 is our first design meant for mass production. Preparations are underway for the first production line, including supply chain readiness, with start of production planned before the end of 2026 and eventual planned capacity of 1 million robots per year.\nInstalled Annual Manufacturing Capacity\n\n[tables/table-5 Tesla Production.html]\n\nInstalled capacity ≠ current production rate and there may be limitations discovered as production rates approach capacity. Production rates depend on a variety of factors, including equipment uptime, component supply, downtime related to factory upgrades, regulatory considerations and other factors. Construction includes factory and infrastructure buildout as well as tool installation.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->Robotics", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->Robotics", "metadata": { "length": 959, "summary": "", @@ -491,7 +491,7 @@ "chunk_id": "c49883a3-5834-5537-be53-35e58da70bf7", "type": "text", "content": "We are currently building Cortex 2 at Gigafactory Texas to further increase our AI training compute capacity. In the first half of 2026, we plan to more than double the size of onsite compute in Texas (in terms of H100 equivalents). We aim to maximize capital efficiency by scaling training compute judiciously, including when the training backlog gets too long or in anticipation of greater demand from our engineers to support our AI-related offerings.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->AI Training Compute", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->AI Training Compute", "metadata": { "length": 454, "summary": "", @@ -545,7 +545,7 @@ "chunk_id": "fcbc3b02-dec6-5e91-849e-fccd3f22320e", "type": "text", "content": "Our lithium refinery commenced pilot production and is the first spodumene to lithium hydroxide refinery in North America, leveraging a simpler, cheaper and more environmentally friendly process. This refinery enables us to domestically produce critical minerals in support of energy storage, battery manufacturing and ultimately for EV growth.\nWe have begun to produce battery packs for certain Model Ys with our 4680 cells, unlocking an additional vector of supply to help navigate increasingly complex supply chain challenges caused by trade barriers and tariff risks. We now produce dry-electrode for 4680 cells with both anode and cathode made in Austin. We expect both domestic cathode material in Texas and LFP lines in Nevada to begin production in 2026.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->Battery", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->Battery", "metadata": { "length": 762, "summary": "", @@ -667,7 +667,7 @@ "chunk_id": "34777c71-d716-5c39-b18a-bec0d0b64706", "type": "text", "content": "We continue to efficiently utilize our existing physical footprint in North America, with targeted augmentation to support the rollout of Robotaxi. While in the short-term, operational workstreams such as charging, cleaning and maintenance can be managed through our existing charging network, service centers and sales and delivery locations, we will have to add more capacity as the service expands. We added over 3,800 net new Supercharging stalls, growing the network 19% year-over-year.\nInstalled Annual Capacity\n\n[tables/table-6 Facility Status.html]\n\nInstalled capacity ≠ current production rate and there may be limitations discovered as production rates approach capacity. Production rates depend on a variety of factors, including equipment uptime, component supply, downtime related to factory upgrades, regulatory considerations and other factors. Construction includes factory and infrastructure buildout as well as tool installation.\n\n## Other Supporting Infrastructure Tesla AI Training Capacity Ramp (H100 equivalent GPUs)0\n[images/image-1-Capacity Growth Projection.jpg]\n\nTesla AI Training Capacity Ramp (H100 equivalent GPUs)", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->Other Supporting Infrastructure", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->Other Supporting Infrastructure", "metadata": { "length": 1142, "summary": "", @@ -786,7 +786,7 @@ "chunk_id": "778a63c2-7955-514c-84d0-9c2cfd99a489", "type": "text", "content": "We continue to enhance FSD (Supervised) $^{1}$ via our end-to-end foundation model trained on both customer and Robotaxi real-world data with our latest version, v14. FSD (Supervised) $^{1}$ increasingly provides safety and convenience functionality that can relieve drivers of many tedious and potentially dangerous aspects of road travel, including giving access to personal transport for those who otherwise have difficulty driving. V14 offers unparalleled driver assistance to safely drive to the customer's destination, find a free parking spot and park at that location. Our global fleet can collect the equivalent of over 500 years of continuous driving data per day $^{2}$ , allowing us to safely deploy and scale capabilities that can handle the long and fat tail of corner cases across diverse geographies and driving environments.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->AI Software", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->AI Software", "metadata": { "length": 841, "summary": "", @@ -876,7 +876,7 @@ "chunk_id": "164babde-7a64-5b81-9a5a-41443b221c40", "type": "text", "content": "Development of our in-house, custom designed AI5 and AI6 inference chips for autonomy progressed during the quarter, with production planned for 2027 and 2028, respectively. We are targeting a 50x improvement in performance for AI5 relative to AI4 thanks to 10x raw compute, 9x memory capacity and 5x hardened block quantization and softmax function (the latter enabling efficient low-precision computing without sacrificing model accuracy).", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->AI Inference Compute", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->AI Inference Compute", "metadata": { "length": 441, "summary": "", @@ -970,7 +970,7 @@ "chunk_id": "32bfae6f-b2d4-51d3-94d9-82b9640828dd", "type": "text", "content": "The Robotaxi iOS app no longer has a waitlist in the areas we serve. Our vehicles keep getting better with our over-the-air updates, including: Grok (an AI companion) which now supports navigation commands (allowing users to find, add and edit navigation destinations hands-free); Tesla Photobooth which enables users to take photos in their car and download or share via the Tesla mobile app; Supercharger Site Maps which displays Supercharger layouts, nearby businesses and live availability details; Automatic HOV Lane Routing based on interior camera occupancy detection; Phone Left Behind Chime and SpaceX ISS Docking Simulator Game.\n\nWe continue to enhance FSD (Supervised) $^{1}$ via our end-to-end foundation model trained on both customer and Robotaxi real-world data with our latest version, v14. FSD (Supervised) $^{1}$ increasingly provides safety and convenience functionality that can relieve drivers of many tedious and potentially dangerous aspects of road travel, including giving access to personal transport for those who otherwise have difficulty driving. V14 offers unparalleled driver assistance to safely drive to the customer's destination, find a free parking spot and park at that location. Our global fleet can collect the equivalent of over 500 years of continuous driving data per day $^{2}$ , allowing us to safely deploy and scale capabilities that can handle the long and fat tail of corner cases across diverse geographies and driving environments. Cumulative Miles Driven with FSD (Supervised) $^{1}$ (billions)0\n[images/image-2-FSD Mileage Growth.jpg]\n\nCumulative Miles Driven with FSD (Supervised) $^{1}$ (billions)\n\nDevelopment of our in-house, custom designed AI5 and AI6 inference chips for autonomy progressed during the quarter, with production planned for 2027 and 2028, respectively. We are targeting a 50x improvement in performance for AI5 relative to AI4 thanks to 10x raw compute, 9x memory capacity and 5x hardened block quantization and softmax function (the latter enabling efficient low-precision computing without sacrificing model accuracy). Targeting Step-Function Improvement for our Next-Generation Inference Chip, AI50\n[images/image-3-Tesla Silicon Optimization.jpg]\n\nTargeting Step-Function Improvement for our Next-Generation Inference Chip, AI5\n(1) Active driver supervision required; does not make the vehicle autonomous\n(2) Calculated based on continuous hours of driving at an average of 30 miles per hour", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->Automotive and Other Software", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->Automotive and Other Software", "metadata": { "length": 2444, "summary": "The Robotaxi iOS app has removed its waitlist in served areas and now includes features like Grok navigation, Tesla Photobooth, Supercharger maps, automatic HOV routing, phone left-behind chimes, and a SpaceX game. FSD (Supervised) v14 uses an end-to-end foundation model trained on vast real-world data to assist drivers with navigation, parking, and safety, though active supervision remains required. The company is developing custom AI5 and AI6 inference chips for 2027 and 2028, targeting significant performance improvements over previous generations to handle complex driving scenarios globally.", @@ -1205,7 +1205,7 @@ "chunk_id": "41533301-58f0-554f-916d-8254cc2700df", "type": "text", "content": "We began testing driverless Robotaxis in Austin in December and began removing the safety monitor from customer rides in January on a limited basis, which will unlock further expansion of our Robotaxi fleet and coverage area in the Austin-metro. Our Bay Area ride-hailing service began serving the San Jose Airport in October, with plans to expand to other major airports in the Bay Area upon receiving required permitting.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->Robotaxi", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->Robotaxi", "metadata": { "length": 423, "summary": "", @@ -1263,7 +1263,7 @@ "chunk_id": "3968a25f-933e-5a7f-a118-07db2e581153", "type": "text", "content": "We launched FSD (Supervised) $^{1}$ in South Korea, where customers drove over 1 million kilometers using the software in just one month. While we continue to pursue regulatory approval in China and Europe, we began offering ride-along experiences to consumers in Italy, Germany, France and Switzerland.\nMonthly subscriptions to FSD (Supervised) $^{1}$ continued to grow sequentially and more than doubled in 2025. Starting this quarter, we are transitioning access to FSD (Supervised) $^{1}$ to monthly subscriptions only as we begin to sunset the up-front payment option.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->FSD (Supervised) $^{1}$", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->FSD (Supervised) $^{1}$", "metadata": { "length": 573, "summary": "", @@ -1364,7 +1364,7 @@ "chunk_id": "adf89909-d366-51e1-b68c-459f9024191c", "type": "text", "content": "Services and Other gross profit of approximately \\$300 million was partly driven by Part Sales and Supercharging. We now offer Tesla Insurance in Florida, as we continue to expand our insurance product to new states. In certain states, customers receive a discount on their insurance premiums when using FSD (Supervised) $^{1}$ . The more you drive with FSD (Supervised) $^{1}$ enabled, the bigger the discount is on your insurance premium – helping, in certain cases, to completely offset the monthly subscription cost for FSD (Supervised) $^{1}$ .\n\n## FSD (Supervised) $^{1}$ Cumulative Paid Robotaxi Miles0\n[images/image-4-Growth Trend 2025.jpg]\n\nCumulative Paid Robotaxi Miles\n\n[tables/table-7 Autonomous Driving Status.html]\n\nPlanned Robotaxi Coverage", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->Automotive Services", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->Automotive Services", "metadata": { "length": 742, "summary": "", @@ -1450,7 +1450,7 @@ "chunk_id": "cf91377e-176a-5768-855a-b859092f4695", "type": "text", "content": "On January 16, 2026, Tesla entered into an agreement to invest approximately \\$2 billion to acquire shares of Series E Preferred Stock of xAI as part of their recent publicly-disclosed financing round. Tesla’s investment was made on market terms consistent with those previously agreed to by other investors in the financing round. As set forth in Master Plan Part IV, Tesla is building products and services that bring AI into the physical world. Meanwhile, xAI is developing leading digital AI products and services, such as its large language model (Grok).\nIn that context, and as part of Tesla's broader strategy under Master Plan Part IV, Tesla and xAI also entered into a framework agreement in connection with the investment. Among other things, the framework agreement builds upon the existing relationship between Tesla and xAI by providing a framework for evaluating potential AI collaborations between the companies. Together, the investment and the related framework agreement are intended to enhance Tesla's ability to develop and deploy AI products and services into the physical world at scale. This investment is subject to customary regulatory conditions with the expectation to close in Q1'2026.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OTHER UPDATES", + "path": "TSLA-Q4-2025-Update.pdf-->OTHER UPDATES", "metadata": { "length": 1213, "summary": "", @@ -1552,7 +1552,7 @@ "chunk_id": "79b23d73-b353-5998-a646-93d6739ec465", "type": "text", "content": "", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK", "metadata": { "length": 0, "summary": "", @@ -1569,7 +1569,7 @@ "chunk_id": "8d4c49a6-c531-5fa6-a77d-ad028537fa52", "type": "text", "content": "We are focused on maximum capacity utilization at our factories. Deliveries and deployments will be impacted by aggregate demand for our products, supply chain readiness and allocation decisions between sale to customers or use for our owned and operated fleet.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Volume", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Volume", "metadata": { "length": 261, "summary": "", @@ -1609,7 +1609,7 @@ "chunk_id": "426e5119-3ef6-5176-9119-fc044eb90be7", "type": "text", "content": "We will manage the businesses such that we ensure a strong balance sheet, maintaining sufficient liquidity to fund our product roadmap, long-term capacity expansion plans – including further vertical integration – and other expenses.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Cash", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Cash", "metadata": { "length": 233, "summary": "", @@ -1649,7 +1649,7 @@ "chunk_id": "99f54247-8409-5df2-8299-39184863cd09", "type": "text", "content": "While we continue to execute on innovations to reduce the cost of manufacturing and operations, over time, we expect our hardware-related profits to be accompanied by an acceleration of AI, software and fleet-based profits.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Profit", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Profit", "metadata": { "length": 223, "summary": "", @@ -1686,7 +1686,7 @@ "chunk_id": "29e89187-3dbe-5298-87c3-c38d3cbe6883", "type": "text", "content": "We continue to evolve and augment our product lineup with a focus on cost, scale and future monetization opportunities via services powered by our AI software. We remain focused on growing our sales volumes through a differentiated and efficiently managed product portfolio, which includes leveraging and optimizing our existing production capacity before building new factories and production lines.\nCybercab, Tesla Semi and Megapack 3 are on schedule for volume production starting in 2026. First generation production lines for Optimus are being installed in anticipation of volume production.\nPHOTOS & CHARTS", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product", "metadata": { "length": 612, "summary": "", @@ -1772,7 +1772,7 @@ "chunk_id": "bbe9d89e-eebd-56b1-b7c0-7eda12dd3377", "type": "text", "content": "Cybercab, Tesla Semi and Megapack 3 are on schedule for volume production starting in 2026. First generation production lines for Optimus are being installed in anticipation of volume production. 0\n[images/image-5-Tesla Model Y Driving.jpg]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->MODEL Y - 2025 BEST IN CLASS EURO NCAP $^{(1)}$ SMALL SUV", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->MODEL Y - 2025 BEST IN CLASS EURO NCAP $^{(1)}$ SMALL SUV", "metadata": { "length": 245, "summary": "", @@ -1835,7 +1835,7 @@ "chunk_id": "02f7e756-798f-5675-9dd9-6ce9123f92d4", "type": "text", "content": " 0\n[images/image-6-Red Tesla on Coastal Road.jpg]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->MODEL 3 - 2025 BEST IN CLASS EURO NCAP $^{(1)}$ LARGE FAMILY CAR", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->MODEL 3 - 2025 BEST IN CLASS EURO NCAP $^{(1)}$ LARGE FAMILY CAR", "metadata": { "length": 66, "summary": "", @@ -1884,7 +1884,7 @@ "chunk_id": "494118dd-5788-526b-a139-407d1cfb0d20", "type": "text", "content": " 0\n[images/image-7-Tesla Interior Interface.jpg]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->FSD (SUPERVISED) $^{1}$ – V14 OFFERS UNPARALLELED DRIVER ASSISTANCE", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->FSD (SUPERVISED) $^{1}$ – V14 OFFERS UNPARALLELED DRIVER ASSISTANCE", "metadata": { "length": 66, "summary": "", @@ -1933,7 +1933,7 @@ "chunk_id": "6ac7765a-4da5-573a-963b-f35ab3796f6f", "type": "text", "content": " 0\n[images/image-8-Tesla Interior.jpg]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->DRIVERLESS ROBOTAXI - TESTING IN AUSTIN", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->DRIVERLESS ROBOTAXI - TESTING IN AUSTIN", "metadata": { "length": 66, "summary": "", @@ -2016,7 +2016,7 @@ "chunk_id": "4cfe7ce9-5f93-5816-9505-ba562683ee42", "type": "text", "content": " 0\n[images/image-9-Tesla Cybertruck in Snow.jpg]\n\n\n### DRIVERLESS ROBOTAXI - TESTING IN AUSTIN 0\n[images/image-10-Tesla Semi Trucks.jpg]\n\nTESLA SEMI - MEGACHARGER NETWORK PLANNED SITES FOR 2026\n\n### CYBERCAB - COLD WEATHER TESTING IN ALASKA 0\n[images/image-11-US Lightning Map.jpg]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->CYBERCAB - COLD WEATHER TESTING IN ALASKA", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->CYBERCAB - COLD WEATHER TESTING IN ALASKA", "metadata": { "length": 318, "summary": "", @@ -2102,7 +2102,7 @@ "chunk_id": "07535539-a2f9-5b79-ac6b-d41cf785ec88", "type": "text", "content": " 0\n[images/image-12-Tesla Factory Milestone.jpg]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->GIGAFACTORY SHANGHAI - 9 MILLIONTH VEHICLE PRODUCED (GLOBALLY)", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->GIGAFACTORY SHANGHAI - 9 MILLIONTH VEHICLE PRODUCED (GLOBALLY)", "metadata": { "length": 67, "summary": "", @@ -2202,7 +2202,7 @@ "chunk_id": "702bff82-1736-584d-956e-e7bf7eac4039", "type": "text", "content": " 0\n[images/image-13-Tesla Factory Milestone.jpg]\n\n\n### GIGAFACTORY SHANGHAI - 9 MILLIONTH VEHICLE PRODUCED (GLOBALLY) 0\n[images/image-14-Vehicle Delivery Trends.jpg]\n\n\n 0\n[images/image-15-Quarterly Cash Flow.jpg]\n\n\n 0\n[images/image-16-Financial Performance Chart.jpg]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->GIGAFACTORY NEVADA - 6 MILLIONTH DRIVE UNIT PRODUCED", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->GIGAFACTORY NEVADA - 6 MILLIONTH DRIVE UNIT PRODUCED", "metadata": { "length": 327, "summary": "", @@ -2319,7 +2319,7 @@ "chunk_id": "5e17588f-ea71-56df-be53-135da93fb3a0", "type": "text", "content": " 0\n[images/image-17-Projected Vehicle Deliveries.jpg]\n\n\n 0\n[images/image-18-Cash Flow Trends.jpg]\n\n\n 0\n[images/image-19-Financial Performance Forecast.jpg]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)", "metadata": { "length": 207, "summary": "", @@ -2369,7 +2369,7 @@ "chunk_id": "c82a6fed-f085-5ec3-b10a-0381723fa3c9", "type": "text", "content": "Total quarterly revenue decreased 3% YoY to \\$24.9B. YoY, revenue was impacted by the following items $^{(1)}$ :\n- decrease in vehicle deliveries\n- lower regulatory credit revenue\n+ growth in Energy Generation and Storage\n+ growth in Services and Other\n+ positive FX impact of \\$0.3B $^{1}$\n\\+ growth in other automotive ancillary sales, partly driven by an increase in FSD subscriptions\n\\+ higher vehicle average selling price (ASP) (excl. FX impact $^{1}$ ), inclusive of mix impact", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)-->Revenue", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)-->Revenue", "metadata": { "length": 484, "summary": "", @@ -2428,7 +2428,7 @@ "chunk_id": "946fabf6-000f-557a-9d81-0b8c1a25aafb", "type": "text", "content": "Our quarterly operating income decreased 11% YoY to \\$1.4B, resulting in a 5.7% operating margin. YoY, operating income was primarily impacted by the following items $^{(1)}$ :\n- increase in SBC and Restructuring and Other charges\n- increase in operating expenses (excl. SBC and Restructuring and Other) driven by AI and other R&D projects and SG&A\n- higher average cost per vehicle due to lower fixed cost absorption for certain models and an increase in tariffs\n- decrease in vehicle deliveries\n- lower regulatory credit revenue\n+ higher vehicle average gross profit due to mix and pricing impacts\n+ growth in Energy Generation and Storage gross profit\n+ growth in Services and Other gross profit\n+ growth in other automotive ancillary sales, partly driven by an increase in FSD subscriptions", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)-->Profitability", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)-->Profitability", "metadata": { "length": 794, "summary": "", @@ -2502,7 +2502,7 @@ "chunk_id": "001d4b8a-b264-5cae-86a5-8c419eabbec0", "type": "text", "content": "Quarter-end cash, cash equivalents and investments was \\$44.1B. The sequential increase of \\$2.4B was primarily the result of positive free cash flow.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)-->Cash", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)-->Cash", "metadata": { "length": 150, "summary": "", @@ -2683,7 +2683,7 @@ "chunk_id": "15c18264-8e1e-5db6-b3c9-70cd181d1f39", "type": "text", "content": "STATEMENT OF OPERATIONS\n(Unaudited)\n\n[tables/table-8 Q4 2024-Q4 2025 Rev.html]\n\nBALANCE SHEET\n(Unaudited)\n\n[tables/table-9 Balance Sheet 2024-25.html]\n\nSTATEMENT OF CASH FLOWS\n(Unaudited)\n\n[tables/table-10 Cash Flow Q4-25.html]\n\nRECONCILIATION OF GAAP TO NON-GAAP FINANCIAL INFORMATION (Unaudited)\n\n[tables/table-11 Q4 2024-Q4 2025.html]\n\nRECONCILIATION OF GAAP TO NON-GAAP FINANCIAL INFORMATION\n(Unaudited)\n\n[tables/table-12 Financial Metrics 2021-25.html]\n\nRECONCILIATION OF GAAP TO NON-GAAP FINANCIAL INFORMATION\n(Unaudited)\n\n[tables/table-13 Financial Data 2022-25.html]\n\n\n[tables/table-14 Financial Metrics 2023-25.html]\n\nTTM = Trailing twelve months\n(1) Beginning in Q1'25, Capital expenditures is presented inclusive of purchases of energy generation and storage systems and all prior periods have been adjusted.\n(2) As a result of the adoption of the new crypto assets standard, the previously reported quarterly periods in 2024 have been recast.\n(3) Beginning in Q1'25, Adjusted EBITDA (non-GAAP) is presented net of digital assets gains and losses and all prior periods have been adjusted.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->FINANCIAL STATEMENTS", + "path": "TSLA-Q4-2025-Update.pdf-->FINANCIAL STATEMENTS", "metadata": { "length": 1418, "summary": "", @@ -2816,7 +2816,7 @@ "chunk_id": "31babb3c-1abf-5f64-90c7-e8e8a5384bbd", "type": "text", "content": "Tesla will provide a live webcast of its fourth quarter 2025 financial results conference call beginning at 4:30 p.m. CT on January 28, 2026 at ir.tesla.com. This webcast will also be available for replay for approximately one year thereafter.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->WEBCAST INFORMATION", + "path": "TSLA-Q4-2025-Update.pdf-->WEBCAST INFORMATION", "metadata": { "length": 243, "summary": "", @@ -2857,7 +2857,7 @@ "chunk_id": "19e3f236-d15a-5b2a-8589-43d28c5316fe", "type": "text", "content": "When used in this update, certain terms have the following meanings. Our vehicle deliveries include only vehicles that have been transferred to end customers with all paperwork correctly completed. Our energy product deployment volume includes both customer units when installed and equipment sales at time of delivery. \"Net income attributable to common stockholders (non-GAAP)\" is equal to (i) net income attributable to common stockholders before (ii)(a) stock-based compensation expense, net of tax, (b) digital assets (gain) loss, net of tax and (c) release of valuation allowance on deferred tax assets. \"Adjusted EBITDA (non-GAAP)\" is equal to (i) net income attributable to common stockholders before (ii)(a) interest expense, (b) provision for (benefit from) income taxes, (c) depreciation, amortization and impairment, (d) stock-based compensation expense and (e) digital assets loss (gain), net. \"Free cash flow\" is operating cash flow less capital expenditures. Average cost per vehicle is cost of automotive sales divided by new vehicle deliveries (excluding operating leases). \"Days sales outstanding\" is equal to (i) average accounts receivable, net for the period divided by (ii) total revenues and multiplied by (iii) the number of days in the period. \"Days payable outstanding\" is equal to (i) average accounts payable for the period divided by (ii) total cost of revenues and multiplied by (iii) the number of days in the period. \"Days of supply\" is calculated by dividing new car ending inventory by the relevant period's deliveries and using trading days. Constant currency impacts are calculated by comparing actuals against current results converted into USD using average exchange rates from the prior period.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->CERTAIN TERMS", + "path": "TSLA-Q4-2025-Update.pdf-->CERTAIN TERMS", "metadata": { "length": 1733, "summary": "This passage defines key financial and operational terms used in a specific update. It clarifies that vehicle deliveries refer to units transferred to end customers with completed paperwork, while energy product deployment includes installed customer units and equipment sales. The text details non-GAAP measures: 'Net income attributable to common stockholders' adjusts for stock-based compensation, digital asset gains/losses, and valuation allowances; 'Adjusted EBITDA' excludes interest, taxes, depreciation, amortization, impairment, stock-based compensation, and digital asset impacts. 'Free cash flow' is defined as operating cash flow minus capital expenditures. Operational metrics include 'Average cost per vehicle' (automotive sales cost divided by new deliveries excluding leases), 'Days sales outstanding' (average receivables divided by revenue times days in period), 'Days payable outstanding' (average payables divided by cost of revenues times days in period), and 'Days of supply' (ending inventory divided by deliveries using trading days). Finally, constant currency impacts are calculated by comparing actuals against results converted to USD using prior period average exchange rates.", @@ -2982,7 +2982,7 @@ "chunk_id": "332edac0-294c-5e5b-8599-a0e52b25e53a", "type": "text", "content": "Consolidated financial information has been presented in accordance with GAAP as well as on a non-GAAP basis to supplement our consolidated financial results. Our non-GAAP financial measures include non-GAAP net income (loss) attributable to common stockholders, non-GAAP net income (loss) attributable to common stockholders on a diluted per share basis (calculated using weighted average shares for GAAP diluted net income (loss) attributable to common stockholders), Adjusted EBITDA margin, non-GAAP automotive gross margin and free cash flow. These non-GAAP financial measures also facilitate management's internal comparisons to Tesla's historical performance as well as comparisons to the operating results of other companies. Management believes that it is useful to supplement its GAAP financial statements with this non-GAAP information because management uses such information internally for its operating, budgeting and financial planning purposes. Management also believes that presentation of the non-GAAP financial measures provides useful information to our investors regarding our financial condition and results of operations, so that investors can see through the eyes of Tesla management regarding important financial metrics that Tesla uses to run the business and allowing investors to better understand Tesla's performance. Non-GAAP information is not prepared under a comprehensive set of accounting rules and therefore, should only be read in conjunction with financial information reported under U.S. GAAP when understanding Tesla's operating performance. A reconciliation between GAAP and non-GAAP financial information is provided above.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->NON-GAAP FINANCIAL INFORMATION", + "path": "TSLA-Q4-2025-Update.pdf-->NON-GAAP FINANCIAL INFORMATION", "metadata": { "length": 1664, "summary": "Tesla presents consolidated financial information under both GAAP and non-GAAP standards to supplement its results. Key non-GAAP measures include net income attributable to common stockholders, diluted per share figures, Adjusted EBITDA margin, automotive gross margin, and free cash flow. These metrics aid internal management comparisons with historical performance and other companies, supporting operating, budgeting, and planning activities. Management believes these non-GAAP figures provide investors with a clearer view of Tesla's financial condition and operational results by reflecting the metrics used to run the business. However, since non-GAAP data is not prepared under comprehensive accounting rules, it should be read alongside U.S. GAAP information for a complete understanding of Tesla's performance. A reconciliation between GAAP and non-GAAP data is provided elsewhere.", @@ -3077,7 +3077,7 @@ "chunk_id": "34263854-525e-5a28-b9fd-a88d4813756e", "type": "text", "content": "Certain statements in this update, including, but not limited to, statements in the “Outlook” section; statements relating to the development, strategy, ramp, production and capacity, demand and market growth, cost, pricing and profitability, investment, deliveries, deployment, availability and other features and improvements and timing of existing and future Tesla products and services and supporting infrastructure; statements regarding operating margin, operating profits, spending and liquidity; and statements regarding expansion, improvements and/or ramp and related timing at our facilities are “forward-looking statements” within the meaning of the Private Securities Litigation Reform Act of 1995. Forward-looking statements are based on assumptions and management’s current expectations, involve certain risks and uncertainties, and are not guarantees. Future results may differ materially from those expressed in any forward-looking statement. The following important factors, without limitation, could cause actual results to differ materially from those in the forward-looking statements: the risk of delays in launching and/or manufacturing our products, services and features cost-effectively; our ability to build and/or grow our products and services, sales, delivery, installation, servicing and charging capabilities and effectively manage this growth; our ability to successfully and timely develop, introduce and scale, as well as our consumers’ demand for, products and services based on artificial intelligence, robotics and automation, electric vehicles, advanced driver assistance systems, and ride-hailing services generally and our vehicles and services specifically; the ability of suppliers to deliver components according to schedules, prices, quality and volumes acceptable to us, and our ability to manage such components effectively; any issues with lithium-ion cells or other components manufactured at our factories; our ability to ramp our factories in accordance with our plans; our ability to procure supply of battery cells, including through our own manufacturing; risks relating to international operations and expansion, including unfavorable and uncertain regulatory, political, economic, tax, tariff, export controls and labor conditions; any failures by Tesla products to perform as expected or if product recalls occur; the risk of product liability claims; competition in the automotive, transportation and energy product and services and robotics markets; our ability to maintain public credibility and confidence in our long-term business prospects; our ability to manage risks relating to our various product financing programs; the status of government and economic incentives for electric vehicles and energy products; our ability to attract, hire and retain key employees and qualified personnel; our ability to maintain the security of our information and production and product systems; our compliance with various regulations and laws applicable to our operations and products, which may evolve from time to time; risks relating to our indebtedness and financing strategies; and adverse foreign exchange movements. More information on potential factors that could affect our financial results is included from time to time in our Securities and Exchange Commission filings and reports, including the risks identified under the section captioned “Risk Factors” in our annual report on Form 10-K filed with the SEC on January 30, 2025 and subsequent quarterly reports on Form 10-Q. Tesla disclaims any obligation to update information contained in these forward-looking statements whether as a result of new information, future events or otherwise.\nTESLA", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->FORWARD-LOOKING STATEMENTS", + "path": "TSLA-Q4-2025-Update.pdf-->FORWARD-LOOKING STATEMENTS", "metadata": { "length": 3711, "summary": "This passage from Tesla outlines that various statements in the update, particularly regarding future outlooks, product development, production capacity, financial metrics, and facility expansions, constitute forward-looking statements under the Private Securities Litigation Reform Act of 1995. These statements reflect management's current expectations based on assumptions and are subject to risks and uncertainties, meaning actual results may differ materially. The text lists numerous specific risk factors that could impact outcomes, including manufacturing delays, supply chain challenges, regulatory hurdles, competition, product performance issues, employee retention, and foreign exchange fluctuations. Tesla advises investors to consult SEC filings, specifically the Form 10-K filed on January 30, 2025, for detailed risk disclosures. The company explicitly disclaims any obligation to update these forward-looking statements due to new information or future events.", diff --git a/apps/api/app/data/demo_documents/tsla-q4-2025/doc_nav.json b/apps/api/app/data/demo_documents/tsla-q4-2025/doc_nav.json index d27f1ba4c..d8a4cdc7e 100644 --- a/apps/api/app/data/demo_documents/tsla-q4-2025/doc_nav.json +++ b/apps/api/app/data/demo_documents/tsla-q4-2025/doc_nav.json @@ -6,15 +6,299 @@ "text_chunks": 36, "image_chunks": 19, "table_chunks": 15, - "max_depth": 1 + "max_depth": 3 }, "sections": [ { - "title": "Root", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->HIGHLIGHTS", + "title": "HIGHLIGHTS", + "path": "TSLA-Q4-2025-Update.pdf/HIGHLIGHTS", "level": 1, "summary": "[tables/table-0 Tesla 2025 Results.html]", - "chunk_count": 36, + "chunk_count": 1, + "children": [] + }, + { + "title": "SUMMARY", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY", + "level": 1, + "summary": "The document presents unaudited financial and operational summaries for a company, covering quarterly data from Q4 2024 through Q4 2025 and annual data from 2021 to 2025. Key notes indicate significant accounting changes effective Q1 2025: Adjusted EBITDA and Net income attributable to common stockholders are now presented net of digital assets gains and losses, with all prior periods adjusted accordingly. Additionally, Capital expenditures now include purchases of energy generation and storage systems, requiring restatement of previous periods. The content references multiple tables detailing these metrics but does not display the specific numerical values.", + "chunk_count": 13, + "children": [ + { + "title": "Automotive", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/Automotive", + "level": 2, + "summary": "While automotive sales declined sequentially, gross margin (even when excluding the impact of regulatory credits) improved. The APAC region continued to show strength across multiple markets and set a record for deliveries in the quarter. We continued the rollout of Model Y variants across markets in Q4, including the standard and performance versions. Preparations continue in North America for the production ramps of Tesla Semi and Cybercab, both commencing 1H26, and production of the next-generation Roadster.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Energy generation and storage", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/Energy generation and storage", + "level": 2, + "summary": "We achieved our highest quarterly energy storage deployments, driven by record Megapack deployments. Total gross profit rose, both sequentially and year-over-year, to a record \\$1.1 billion, marking the fifth consecutive record quarter. We plan to begin Megapack 3 and Megablock production at Megafactory Houston in 2026. In 2025, our global Powerwall network supported more than 89,000 Virtual Power Plant events across over 1 million installed units, allowing homeowners to save over \\$1 billion in electricity bills as Virtual Power Plant participation continues to scale rapidly.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Robotics", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/Robotics", + "level": 2, + "summary": "We made further progress on the Optimus program in 2025. In Q1 of this year, we plan to unveil the Gen 3 version of Optimus, which will include major upgrades from version 2.5, including our latest hand design. The Gen 3 is our first design meant for mass production. Preparations are underway for the first production line, including supply chain readiness, with start of production planned before the end of 2026 and eventual planned capacity of 1 million robots per year. Installed Annual Manufacturing Capacity [tables/table-5 Tesla Production.html] Installed capacity ≠ current production rate and there may be limitations discovered as production rates approach capacity. Production rates depend on a variety of factors, including equipment uptime, component supply, downtime related to factory upgrades, regulatory considerations and other factors. Construction includes factory and infrastructure buildout as well as tool installation.", + "chunk_count": 1, + "children": [] + }, + { + "title": "AI Training Compute", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/AI Training Compute", + "level": 2, + "summary": "We are currently building Cortex 2 at Gigafactory Texas to further increase our AI training compute capacity. In the first half of 2026, we plan to more than double the size of onsite compute in Texas (in terms of H100 equivalents). We aim to maximize capital efficiency by scaling training compute judiciously, including when the training backlog gets too long or in anticipation of greater demand from our engineers to support our AI-related offerings.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Battery", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/Battery", + "level": 2, + "summary": "Our lithium refinery commenced pilot production and is the first spodumene to lithium hydroxide refinery in North America, leveraging a simpler, cheaper and more environmentally friendly process. This refinery enables us to domestically produce critical minerals in support of energy storage, battery manufacturing and ultimately for EV growth. We have begun to produce battery packs for certain Model Ys with our 4680 cells, unlocking an additional vector of supply to help navigate increasingly complex supply chain challenges caused by trade barriers and tariff risks. We now produce dry-electrode for 4680 cells with both anode and cathode made in Austin. We expect both domestic cathode material in Texas and LFP lines in Nevada to begin production in 2026.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Other Supporting Infrastructure", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/Other Supporting Infrastructure", + "level": 2, + "summary": "We continue to efficiently utilize our existing physical footprint in North America, with targeted augmentation to support the rollout of Robotaxi. While in the short-term, operational workstreams such as charging, cleaning and maintenance can be managed through our existing charging network, service centers and sales and delivery locations, we will have to add more capacity as the service expands. We added over 3,800 net new Supercharging stalls, growing the network 19% year-over-year. Installed Annual Capacity [tables/table-6 Facility Status.html] Installed capacity ≠ current production rate and there may be limitations discovered as production rates approach capacity. Production rates depend on a variety of factors, including equipment uptime, component supply, downtime related to factory upgrades, regulatory considerations and other factors. Construction includes factory and infrastructure buildout as well as tool installation. ## Other Supporting Infrastructure Tesla AI Training Capacity Ramp (H100 equivalent GPUs)0 [images/image-1-Capacity Growth Projection.jpg] Tesla AI Training Capacity Ramp (H100 equivalent GPUs)", + "chunk_count": 1, + "children": [] + }, + { + "title": "AI Software", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/AI Software", + "level": 2, + "summary": "We continue to enhance FSD (Supervised) $^{1}$ via our end-to-end foundation model trained on both customer and Robotaxi real-world data with our latest version, v14. FSD (Supervised) $^{1}$ increasingly provides safety and convenience functionality that can relieve drivers of many tedious and potentially dangerous aspects of road travel, including giving access to personal transport for those who otherwise have difficulty driving. V14 offers unparalleled driver assistance to safely drive to the customer's destination, find a free parking spot and park at that location. Our global fleet can collect the equivalent of over 500 years of continuous driving data per day $^{2}$ , allowing us to safely deploy and scale capabilities that can handle the long and fat tail of corner cases across diverse geographies and driving environments.", + "chunk_count": 1, + "children": [] + }, + { + "title": "AI Inference Compute", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/AI Inference Compute", + "level": 2, + "summary": "Development of our in-house, custom designed AI5 and AI6 inference chips for autonomy progressed during the quarter, with production planned for 2027 and 2028, respectively. We are targeting a 50x improvement in performance for AI5 relative to AI4 thanks to 10x raw compute, 9x memory capacity and 5x hardened block quantization and softmax function (the latter enabling efficient low-precision computing without sacrificing model accuracy).", + "chunk_count": 1, + "children": [] + }, + { + "title": "Automotive and Other Software", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/Automotive and Other Software", + "level": 2, + "summary": "The Robotaxi iOS app has removed its waitlist in served areas and now includes features like Grok navigation, Tesla Photobooth, Supercharger maps, automatic HOV routing, phone left-behind chimes, and a SpaceX game. FSD (Supervised) v14 uses an end-to-end foundation model trained on vast real-world data to assist drivers with navigation, parking, and safety, though active supervision remains required. The company is developing custom AI5 and AI6 inference chips for 2027 and 2028, targeting significant performance improvements over previous generations to handle complex driving scenarios globally.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Robotaxi", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/Robotaxi", + "level": 2, + "summary": "We began testing driverless Robotaxis in Austin in December and began removing the safety monitor from customer rides in January on a limited basis, which will unlock further expansion of our Robotaxi fleet and coverage area in the Austin-metro. Our Bay Area ride-hailing service began serving the San Jose Airport in October, with plans to expand to other major airports in the Bay Area upon receiving required permitting.", + "chunk_count": 1, + "children": [] + }, + { + "title": "FSD (Supervised) $^{1}$", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/FSD (Supervised) $^{1}$", + "level": 2, + "summary": "We launched FSD (Supervised) $^{1}$ in South Korea, where customers drove over 1 million kilometers using the software in just one month. While we continue to pursue regulatory approval in China and Europe, we began offering ride-along experiences to consumers in Italy, Germany, France and Switzerland. Monthly subscriptions to FSD (Supervised) $^{1}$ continued to grow sequentially and more than doubled in 2025. Starting this quarter, we are transitioning access to FSD (Supervised) $^{1}$ to monthly subscriptions only as we begin to sunset the up-front payment option.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Automotive Services", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/Automotive Services", + "level": 2, + "summary": "Services and Other gross profit of approximately \\$300 million was partly driven by Part Sales and Supercharging. We now offer Tesla Insurance in Florida, as we continue to expand our insurance product to new states. In certain states, customers receive a discount on their insurance premiums when using FSD (Supervised) $^{1}$ . The more you drive with FSD (Supervised) $^{1}$ enabled, the bigger the discount is on your insurance premium – helping, in certain cases, to completely offset the monthly subscription cost for FSD (Supervised) $^{1}$ . ## FSD (Supervised) $^{1}$ Cumulative Paid Robotaxi Miles0 [images/image-4-Growth Trend 2025.jpg] Cumulative Paid Robotaxi Miles [tables/table-7 Autonomous Driving Status.html] Planned Robotaxi Coverage", + "chunk_count": 1, + "children": [] + } + ] + }, + { + "title": "OTHER UPDATES", + "path": "TSLA-Q4-2025-Update.pdf/OTHER UPDATES", + "level": 1, + "summary": "On January 16, 2026, Tesla entered into an agreement to invest approximately \\$2 billion to acquire shares of Series E Preferred Stock of xAI as part of their recent publicly-disclosed financing round. Tesla’s investment was made on market terms consistent with those previously agreed to by other investors in the financing round. As set forth in Master Plan Part IV, Tesla is building products and services that bring AI into the physical world. Meanwhile, xAI is developing leading digital AI products and services, such as its large language model (Grok). In that context, and as part of Tesla's broader strategy under Master Plan Part IV, Tesla and xAI also entered into a framework agreement in connection with the investment. Among other things, the framework agreement builds upon the existing relationship between Tesla and xAI by providing a framework for evaluating potential AI collaborations between the companies. Together, the investment and the related framework agreement are intended to enhance Tesla's ability to develop and deploy AI products and services into the physical world at scale. This investment is subject to customary regulatory conditions with the expectation to close in Q1'2026.", + "chunk_count": 1, + "children": [] + }, + { + "title": "OUTLOOK", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK", + "level": 1, + "summary": "", + "chunk_count": 16, + "children": [ + { + "title": "Volume", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Volume", + "level": 2, + "summary": "We are focused on maximum capacity utilization at our factories. Deliveries and deployments will be impacted by aggregate demand for our products, supply chain readiness and allocation decisions between sale to customers or use for our owned and operated fleet.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Cash", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Cash", + "level": 2, + "summary": "We will manage the businesses such that we ensure a strong balance sheet, maintaining sufficient liquidity to fund our product roadmap, long-term capacity expansion plans – including further vertical integration – and other expenses.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Profit", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Profit", + "level": 2, + "summary": "While we continue to execute on innovations to reduce the cost of manufacturing and operations, over time, we expect our hardware-related profits to be accompanied by an acceleration of AI, software and fleet-based profits.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Product", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Product", + "level": 2, + "summary": "We continue to evolve and augment our product lineup with a focus on cost, scale and future monetization opportunities via services powered by our AI software. We remain focused on growing our sales volumes through a differentiated and efficiently managed product portfolio, which includes leveraging and optimizing our existing production capacity before building new factories and production lines. Cybercab, Tesla Semi and Megapack 3 are on schedule for volume production starting in 2026. First generation production lines for Optimus are being installed in anticipation of volume production. PHOTOS & CHARTS", + "chunk_count": 8, + "children": [ + { + "title": "MODEL Y - 2025 BEST IN CLASS EURO NCAP $^{(1)}$ SMALL SUV", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Product/MODEL Y - 2025 BEST IN CLASS EURO NCAP $^{(1)}$ SMALL SUV", + "level": 3, + "summary": "Cybercab, Tesla Semi and Megapack 3 are on schedule for volume production starting in 2026. First generation production lines for Optimus are being installed in anticipation of volume production. 0 [images/image-5-Tesla Model Y Driving.jpg]", + "chunk_count": 1, + "children": [] + }, + { + "title": "MODEL 3 - 2025 BEST IN CLASS EURO NCAP $^{(1)}$ LARGE FAMILY CAR", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Product/MODEL 3 - 2025 BEST IN CLASS EURO NCAP $^{(1)}$ LARGE FAMILY CAR", + "level": 3, + "summary": " 0 [images/image-6-Red Tesla on Coastal Road.jpg]", + "chunk_count": 1, + "children": [] + }, + { + "title": "FSD (SUPERVISED) $^{1}$ – V14 OFFERS UNPARALLELED DRIVER ASSISTANCE", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Product/FSD (SUPERVISED) $^{1}$ – V14 OFFERS UNPARALLELED DRIVER ASSISTANCE", + "level": 3, + "summary": " 0 [images/image-7-Tesla Interior Interface.jpg]", + "chunk_count": 1, + "children": [] + }, + { + "title": "DRIVERLESS ROBOTAXI - TESTING IN AUSTIN", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Product/DRIVERLESS ROBOTAXI - TESTING IN AUSTIN", + "level": 3, + "summary": " 0 [images/image-8-Tesla Interior.jpg]", + "chunk_count": 1, + "children": [] + }, + { + "title": "CYBERCAB - COLD WEATHER TESTING IN ALASKA", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Product/CYBERCAB - COLD WEATHER TESTING IN ALASKA", + "level": 3, + "summary": " 0 [images/image-9-Tesla Cybertruck in Snow.jpg] ### DRIVERLESS ROBOTAXI - TESTING IN AUSTIN 0 [images/image-10-Tesla Semi Trucks.jpg] TESLA SEMI - MEGACHARGER NETWORK PLANNED SITES FOR 2026 ### CYBERCAB - COLD WEATHER TESTING IN ALASKA 0 [images/image-11-US Lightning Map.jpg]", + "chunk_count": 1, + "children": [] + }, + { + "title": "GIGAFACTORY SHANGHAI - 9 MILLIONTH VEHICLE PRODUCED (GLOBALLY)", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Product/GIGAFACTORY SHANGHAI - 9 MILLIONTH VEHICLE PRODUCED (GLOBALLY)", + "level": 3, + "summary": " 0 [images/image-12-Tesla Factory Milestone.jpg]", + "chunk_count": 1, + "children": [] + }, + { + "title": "GIGAFACTORY NEVADA - 6 MILLIONTH DRIVE UNIT PRODUCED", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Product/GIGAFACTORY NEVADA - 6 MILLIONTH DRIVE UNIT PRODUCED", + "level": 3, + "summary": " 0 [images/image-13-Tesla Factory Milestone.jpg] ### GIGAFACTORY SHANGHAI - 9 MILLIONTH VEHICLE PRODUCED (GLOBALLY) 0 [images/image-14-Vehicle Delivery Trends.jpg] 0 [images/image-15-Quarterly Cash Flow.jpg] 0 [images/image-16-Financial Performance Chart.jpg]", + "chunk_count": 1, + "children": [] + } + ] + }, + { + "title": "KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)", + "level": 2, + "summary": " 0 [images/image-17-Projected Vehicle Deliveries.jpg] 0 [images/image-18-Cash Flow Trends.jpg] 0 [images/image-19-Financial Performance Forecast.jpg]", + "chunk_count": 4, + "children": [ + { + "title": "Revenue", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)/Revenue", + "level": 3, + "summary": "Total quarterly revenue decreased 3% YoY to \\$24.9B. YoY, revenue was impacted by the following items $^{(1)}$ : - decrease in vehicle deliveries - lower regulatory credit revenue + growth in Energy Generation and Storage + growth in Services and Other + positive FX impact of \\$0.3B $^{1}$ \\+ growth in other automotive ancillary sales, partly driven by an increase in FSD subscriptions \\+ higher vehicle average selling price (ASP) (excl. FX impact $^{1}$ ), inclusive of mix impact", + "chunk_count": 1, + "children": [] + }, + { + "title": "Profitability", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)/Profitability", + "level": 3, + "summary": "Our quarterly operating income decreased 11% YoY to \\$1.4B, resulting in a 5.7% operating margin. YoY, operating income was primarily impacted by the following items $^{(1)}$ : - increase in SBC and Restructuring and Other charges - increase in operating expenses (excl. SBC and Restructuring and Other) driven by AI and other R&D projects and SG&A - higher average cost per vehicle due to lower fixed cost absorption for certain models and an increase in tariffs - decrease in vehicle deliveries - lower regulatory credit revenue + higher vehicle average gross profit due to mix and pricing impacts + growth in Energy Generation and Storage gross profit + growth in Services and Other gross profit + growth in other automotive ancillary sales, partly driven by an increase in FSD subscriptions", + "chunk_count": 1, + "children": [] + }, + { + "title": "Cash", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)/Cash", + "level": 3, + "summary": "Quarter-end cash, cash equivalents and investments was \\$44.1B. The sequential increase of \\$2.4B was primarily the result of positive free cash flow.", + "chunk_count": 1, + "children": [] + } + ] + } + ] + }, + { + "title": "FINANCIAL STATEMENTS", + "path": "TSLA-Q4-2025-Update.pdf/FINANCIAL STATEMENTS", + "level": 1, + "summary": "STATEMENT OF OPERATIONS (Unaudited) [tables/table-8 Q4 2024-Q4 2025 Rev.html] BALANCE SHEET (Unaudited) [tables/table-9 Balance Sheet 2024-25.html] STATEMENT OF CASH FLOWS (Unaudited) [tables/table-10 Cash Flow Q4-25.html] RECONCILIATION OF GAAP TO NON-GAAP FINANCIAL INFORMATION (Unaudited) [tables/table-11 Q4 2024-Q4 2025.html] RECONCILIATION OF GAAP TO NON-GAAP FINANCIAL INFORMATION (Unaudited) [tables/table-12 Financial Metrics 2021-25.html] RECONCILIATION OF GAAP TO NON-GAAP FINANCIAL INFORMATION (Unaudited) [tables/table-13 Financial Data 2022-25.html] [tables/table-14 Financial Metrics 2023-25.html] TTM = Trailing twelve months (1) Beginning in Q1'25, Capital expenditures is presented inclusive of purchases of energy generation and storage systems and all prior periods have been adjusted. (2) As a result of the adoption of the new crypto assets standard, the previously reported quarterly periods in 2024 have been recast. (3) Beginning in Q1'25, Adjusted EBITDA (non-GAAP) is presented net of digital assets gains and losses and all prior periods have been adjusted.", + "chunk_count": 1, + "children": [] + }, + { + "title": "WEBCAST INFORMATION", + "path": "TSLA-Q4-2025-Update.pdf/WEBCAST INFORMATION", + "level": 1, + "summary": "Tesla will provide a live webcast of its fourth quarter 2025 financial results conference call beginning at 4:30 p.m. CT on January 28, 2026 at ir.tesla.com. This webcast will also be available for replay for approximately one year thereafter.", + "chunk_count": 1, + "children": [] + }, + { + "title": "CERTAIN TERMS", + "path": "TSLA-Q4-2025-Update.pdf/CERTAIN TERMS", + "level": 1, + "summary": "This passage defines key financial and operational terms used in a specific update. It clarifies that vehicle deliveries refer to units transferred to end customers with completed paperwork, while energy product deployment includes installed customer units and equipment sales. The text details non-GAAP measures: 'Net income attributable to common stockholders' adjusts for stock-based compensation, digital asset gains/losses, and valuation allowances; 'Adjusted EBITDA' excludes interest, taxes, depreciation, amortization, impairment, stock-based compensation, and digital asset impacts. 'Free cash flow' is defined as operating cash flow minus capital expenditures. Operational metrics include 'Average cost per vehicle' (automotive sales cost divided by new deliveries excluding leases), 'Days sales outstanding' (average receivables divided by revenue times days in period), 'Days payable outstanding' (average payables divided by cost of revenues times days in period), and 'Days of supply' (ending inventory divided by deliveries using trading days). Finally, constant currency impacts are calculated by comparing actuals against results converted to USD using prior period average exchange rates.", + "chunk_count": 1, + "children": [] + }, + { + "title": "NON-GAAP FINANCIAL INFORMATION", + "path": "TSLA-Q4-2025-Update.pdf/NON-GAAP FINANCIAL INFORMATION", + "level": 1, + "summary": "Tesla presents consolidated financial information under both GAAP and non-GAAP standards to supplement its results. Key non-GAAP measures include net income attributable to common stockholders, diluted per share figures, Adjusted EBITDA margin, automotive gross margin, and free cash flow. These metrics aid internal management comparisons with historical performance and other companies, supporting operating, budgeting, and planning activities. Management believes these non-GAAP figures provide investors with a clearer view of Tesla's financial condition and operational results by reflecting the metrics used to run the business. However, since non-GAAP data is not prepared under comprehensive accounting rules, it should be read alongside U.S. GAAP information for a complete understanding of Tesla's performance. A reconciliation between GAAP and non-GAAP data is provided elsewhere.", + "chunk_count": 1, + "children": [] + }, + { + "title": "FORWARD-LOOKING STATEMENTS", + "path": "TSLA-Q4-2025-Update.pdf/FORWARD-LOOKING STATEMENTS", + "level": 1, + "summary": "This passage from Tesla outlines that various statements in the update, particularly regarding future outlooks, product development, production capacity, financial metrics, and facility expansions, constitute forward-looking statements under the Private Securities Litigation Reform Act of 1995. These statements reflect management's current expectations based on assumptions and are subject to risks and uncertainties, meaning actual results may differ materially. The text lists numerous specific risk factors that could impact outcomes, including manufacturing delays, supply chain challenges, regulatory hurdles, competition, product performance issues, employee retention, and foreign exchange fluctuations. Tesla advises investors to consult SEC filings, specifically the Form 10-K filed on January 30, 2025, for detailed risk disclosures. The company explicitly disclaims any obligation to update these forward-looking statements due to new information or future events.", + "chunk_count": 1, "children": [] } ], @@ -160,4 +444,4 @@ } ] } -} \ No newline at end of file +} diff --git a/apps/api/app/mcp/retrieval_server.py b/apps/api/app/mcp/retrieval_server.py index 2121cdb4e..e1bff80e8 100644 --- a/apps/api/app/mcp/retrieval_server.py +++ b/apps/api/app/mcp/retrieval_server.py @@ -2,18 +2,20 @@ from typing import Annotated, Any, AsyncContextManager, Callable -from app.core.dependencies import get_current_user_id +from app.services.auth.current_user_authentication_service import ( + get_current_user_authentication_service, +) from mcp.server.fastmcp import Context, FastMCP from mcp.server.transport_security import TransportSecuritySettings from pydantic import Field from sqlalchemy.ext.asyncio import AsyncSession from shared.core.database import get_db_context -from shared.services.retrieval import run_retrieval_query +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace +from shared.services.retrieval.app_service import run_retrieval_query DbFactory = Callable[[], AsyncContextManager[AsyncSession]] KNOWHERE_NAMESPACE_HEADER = "x-knowhere-namespace" -DEFAULT_NAMESPACE = "default" def create_public_mcp_transport_security() -> TransportSecuritySettings: @@ -42,10 +44,10 @@ def resolve_mcp_namespace(*, ctx: Context | None) -> str: try: request = get_mcp_request(ctx) except (RuntimeError, ValueError): - return DEFAULT_NAMESPACE + return normalize_retrieval_namespace(None) headers = getattr(request, "headers", {}) or {} - namespace = str(get_header(headers, KNOWHERE_NAMESPACE_HEADER) or "").strip() - return namespace or DEFAULT_NAMESPACE + namespace = get_header(headers, KNOWHERE_NAMESPACE_HEADER) + return normalize_retrieval_namespace(namespace) def to_mcp_query_response(response: dict[str, Any]) -> dict[str, Any]: @@ -81,10 +83,9 @@ async def resolve_mcp_user_id(*, ctx: Context | None, db: AsyncSession) -> str: request = get_mcp_request(ctx) headers = getattr(request, "headers", {}) or {} authorization = get_header(headers, "authorization") - return await get_current_user_id( - request=request, + return await get_current_user_authentication_service().authenticate_authorization_header( + db, authorization=authorization, - db=db, ) @@ -96,7 +97,7 @@ def create_retrieval_mcp_server( server = FastMCP( "knowhere-retrieval", instructions=( - "Use this server to search knowledge. " + "Use this server to search published documents. " "If you need information before answering, try searching with this tool." ), streamable_http_path=streamable_http_path, @@ -105,10 +106,10 @@ def create_retrieval_mcp_server( ) @server.tool( - name="kb.query", - description="Search for information and return relevant knowledge snippets.", + name="retrieval.query", + description="Search published documents and return relevant snippets.", ) - async def kb_query( + async def query_documents( query: Annotated[str, Field(description="What you want to search for.")], top_k: Annotated[ int, Field(description="Maximum number of results to return.") diff --git a/apps/api/app/repositories/job_repository.py b/apps/api/app/repositories/job_repository.py index fd9deafe3..c3750352b 100644 --- a/apps/api/app/repositories/job_repository.py +++ b/apps/api/app/repositories/job_repository.py @@ -3,13 +3,13 @@ from datetime import datetime from typing import Any, Dict, Optional, Sequence -from app.services.state_machine import JobStateMachine from loguru import logger from sqlalchemy import and_, desc, select, update from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import selectinload +from shared.core.state_machine.service import AsyncStateMachineService from shared.models.database.job import Job from shared.models.database.job_state_history import JobStateHistory @@ -18,7 +18,7 @@ class JobRepository: """Repository for Job persistence operations.""" def __init__(self): - self.state_machine = JobStateMachine() + self.state_machine = AsyncStateMachineService() async def create_job( self, diff --git a/apps/api/app/services/__init__.py b/apps/api/app/services/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/apps/api/app/services/auth/api_key_authentication_service.py b/apps/api/app/services/auth/api_key_authentication_service.py new file mode 100644 index 000000000..67e677247 --- /dev/null +++ b/apps/api/app/services/auth/api_key_authentication_service.py @@ -0,0 +1,178 @@ +"""API key authentication workflow.""" + +from __future__ import annotations + +import asyncio +import json +from datetime import datetime, timezone + +from app.repositories.api_key_repository import APIKeyRepository +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.config import redis_pool_manager +from shared.core.database import get_db_context +from shared.services.redis.redis_service import RedisService +from shared.utils.api_keys import hash_api_key + +_API_KEY_USER_CACHE_TTL_SECONDS: int = 3600 + + +class APIKeyAuthenticationService: + """Validate API keys and maintain the API-key auth cache.""" + + def __init__( + self, + *, + repository: APIKeyRepository | None = None, + ) -> None: + self._repository = repository or APIKeyRepository() + + async def validate_api_key( + self, + session: AsyncSession, + api_key: str, + ) -> str | None: + """Validate an API key and return the owning user ID.""" + key_hash = hash_api_key(api_key) + redis_service = redis_pool_manager.get_redis_service() + cached_user_id = await self._get_cached_user_id(redis_service, key_hash) + if cached_user_id is not None: + return cached_user_id + + api_key_record = await self._repository.get_by_key_hash(session, key_hash) + if not api_key_record or not api_key_record.is_valid(): + return None + + self._schedule_last_used_update(str(api_key_record.id)) + user_id = str(api_key_record.user_id) + await self._set_cached_user_id( + redis_service, + key_hash, + user_id, + self._resolve_api_key_cache_ttl_seconds(api_key_record.expires_at), + ) + return user_id + + async def invalidate_api_key_user_cache( + self, + *, + user_id: str, + api_key_hash: str, + ) -> None: + """Remove one API-key auth cache entry.""" + await self._invalidate_cached_api_key_user_id( + redis_pool_manager.get_redis_service(), + user_id, + api_key_hash, + ) + + @staticmethod + def _get_user_id_key(api_key_hash: str) -> str: + return f"api-key:user-id:{api_key_hash}" + + @staticmethod + def _get_user_api_keys_key(user_id: str) -> str: + return f"api-key:user-hashes:{user_id}" + + async def _get_cached_user_id( + self, + redis_service: RedisService, + api_key_hash: str, + ) -> str | None: + try: + raw_user_id = await redis_service.get(self._get_user_id_key(api_key_hash)) + return self._coerce_user_id(raw_user_id) + except Exception: + logger.warning("api_key_authentication: failed to read API-key user cache") + return None + + async def _set_cached_user_id( + self, + redis_service: RedisService, + api_key_hash: str, + user_id: str, + ttl_seconds: int, + ) -> None: + effective_ttl_seconds = min(_API_KEY_USER_CACHE_TTL_SECONDS, ttl_seconds) + user_id_key = self._get_user_id_key(api_key_hash) + user_api_keys_key = self._get_user_api_keys_key(user_id) + + try: + await redis_service.set(user_id_key, user_id, ttl=effective_ttl_seconds) + await redis_service.sadd(user_api_keys_key, api_key_hash) + reverse_ttl_seconds = await redis_service.ttl(user_api_keys_key) + if ( + reverse_ttl_seconds in (-2, -1) + or reverse_ttl_seconds < effective_ttl_seconds + ): + await redis_service.expire(user_api_keys_key, effective_ttl_seconds) + except Exception: + logger.warning( + "api_key_authentication: failed to write API-key user cache for user_id={}", + user_id, + ) + + async def _invalidate_cached_api_key_user_id( + self, + redis_service: RedisService, + user_id: str, + api_key_hash: str, + ) -> None: + try: + await redis_service.delete(self._get_user_id_key(api_key_hash)) + await redis_service.srem(self._get_user_api_keys_key(user_id), api_key_hash) + except Exception: + logger.warning( + "api_key_authentication: failed to invalidate API-key cache for user_id={}", + user_id, + ) + + def _coerce_user_id(self, raw_user_id: object) -> str | None: + if isinstance(raw_user_id, str): + try: + parsed_user_id: object = json.loads(raw_user_id) + except json.JSONDecodeError: + return raw_user_id + else: + parsed_user_id = raw_user_id + + if isinstance(parsed_user_id, str): + return parsed_user_id + + if isinstance(parsed_user_id, dict): + legacy_user_id = parsed_user_id.get("user_id") + if isinstance(legacy_user_id, str): + return legacy_user_id + + return None + + def _resolve_api_key_cache_ttl_seconds(self, expires_at: datetime | None) -> int: + if expires_at is None: + return _API_KEY_USER_CACHE_TTL_SECONDS + + expires_at_utc = expires_at + if expires_at_utc.tzinfo is None: + expires_at_utc = expires_at_utc.replace(tzinfo=timezone.utc) + + now = datetime.now(timezone.utc) + remaining_seconds = int((expires_at_utc - now).total_seconds()) + return max(1, min(_API_KEY_USER_CACHE_TTL_SECONDS, remaining_seconds)) + + def _schedule_last_used_update(self, api_key_id: str) -> None: + try: + asyncio.create_task( + self._update_last_used_best_effort(api_key_id), + name=f"api_key_last_used:{api_key_id}", + ) + except Exception as exc: + logger.warning( + f"Failed to schedule API key last-used update (ignored): {exc}" + ) + + async def _update_last_used_best_effort(self, api_key_id: str) -> None: + try: + async with get_db_context() as db: + await self._repository.update_last_used(db, api_key_id) + except Exception as exc: + logger.warning(f"Failed to update API key last-used time (ignored): {exc}") diff --git a/apps/api/app/services/auth/api_key_management_service.py b/apps/api/app/services/auth/api_key_management_service.py new file mode 100644 index 000000000..0adbfee94 --- /dev/null +++ b/apps/api/app/services/auth/api_key_management_service.py @@ -0,0 +1,217 @@ +"""API key management workflow.""" + +from __future__ import annotations + +from datetime import datetime +from typing import TypedDict + +from app.repositories.api_key_repository import APIKeyRepository +from app.services.auth.api_key_authentication_service import ( + APIKeyAuthenticationService, +) +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + APIKeyOperationException, + KnowhereException, + NotFoundException, + ValidationException, +) +from shared.models.database.api_key import APIKey +from shared.utils.api_keys import generate_api_key, hash_api_key, mask_api_key + + +class APIKeyListItem(TypedDict): + id: str + name: str + api_key: str + enabled_modules: list[str] | None + is_active: bool + created_at: datetime + last_used_at: datetime | None + expires_at: datetime | None + + +class APIKeyManagementService: + """Create, list, read, revoke, and toggle API keys.""" + + def __init__( + self, + *, + repository: APIKeyRepository | None = None, + authentication_service: APIKeyAuthenticationService | None = None, + ) -> None: + self._repository = repository or APIKeyRepository() + self._authentication_service = ( + authentication_service or APIKeyAuthenticationService() + ) + + async def create_api_key( + self, + session: AsyncSession, + *, + user_id: str, + name: str, + enabled_modules: list[str] | None = None, + expires_at: datetime | None = None, + ) -> str: + key_count = await self._repository.count_by_user(session, user_id) + if key_count >= 10: + raise ValidationException( + user_message="Maximum API Key limit reached (10)", + violations=[ + { + "field": "api_keys", + "description": "User has reached the maximum API Key limit", + } + ], + ) + + existing_key = await self._repository.get_by_user_and_name( + session, + user_id, + name, + ) + if existing_key: + raise ValidationException( + user_message="API Key name already exists", + violations=[ + { + "field": "name", + "description": f"An API Key with name '{name}' already exists", + } + ], + ) + + api_key = generate_api_key() + api_key_record = APIKey( + user_id=user_id, + key_hash=hash_api_key(api_key), + key_mask=mask_api_key(api_key), + name=name, + enabled_modules=enabled_modules or ["all"], + expires_at=expires_at, + ) + await self._repository.create(session, api_key_record) + return api_key + + async def revoke_api_key( + self, + session: AsyncSession, + *, + api_key_id: str, + user_id: str, + ) -> bool: + logger.info(f"Revoking API key: api_key_id={api_key_id}, user_id={user_id}") + api_key = await self._repository.get_by_id(session, api_key_id) + + if not api_key: + logger.warning("API key does not exist") + raise NotFoundException( + resource="APIKey", + resource_id=api_key_id, + internal_message="API Key not found", + ) + + if str(api_key.user_id) != user_id: + logger.warning( + f"User ID mismatch: api_key.user_id={api_key.user_id}, user_id={user_id}" + ) + raise NotFoundException( + resource="APIKey", + resource_id=api_key_id, + internal_message="API Key not found or does not belong to user", + ) + + success = await self._repository.delete_by_id(session, api_key_id) + logger.info(f"Delete result: {success}") + + if success: + await session.commit() + logger.info("Transaction committed") + await self._authentication_service.invalidate_api_key_user_cache( + user_id=user_id, + api_key_hash=api_key.key_hash, + ) + + return success + + async def list_user_api_keys( + self, + session: AsyncSession, + *, + user_id: str, + ) -> list[APIKeyListItem]: + api_keys = await self._repository.get_unexpired_by_user_id(session, user_id) + return [ + { + "id": str(api_key.id), + "name": api_key.name, + "api_key": api_key.key_mask + or f"sk_{api_key.id[:8]}••••••••••••••••••••••••••••••••••••••••", + "enabled_modules": api_key.enabled_modules, + "is_active": api_key.is_active, + "created_at": api_key.created_at, + "last_used_at": api_key.last_used_at, + "expires_at": api_key.expires_at, + } + for api_key in api_keys + ] + + async def get_api_key( + self, + session: AsyncSession, + *, + user_id: str, + api_key_id: str, + ) -> APIKey | None: + try: + api_key = await self._repository.get(session, api_key_id) + if api_key and str(api_key.user_id) == user_id: + return api_key + return None + except KnowhereException: + raise + except Exception as exc: + logger.error(f"Failed to get API key: {exc}") + raise APIKeyOperationException( + internal_message=f"Failed to get API key: {str(exc)}", + original_exception=exc, + ) + + async def toggle_api_key( + self, + session: AsyncSession, + *, + user_id: str, + api_key_id: str, + ) -> bool: + try: + api_key = await self._repository.get(session, api_key_id) + if not api_key or str(api_key.user_id) != user_id: + return False + + api_key.is_active = not api_key.is_active + await session.commit() + await session.refresh(api_key) + + if not api_key.is_active: + await self._authentication_service.invalidate_api_key_user_cache( + user_id=user_id, + api_key_hash=api_key.key_hash, + ) + + logger.info( + f"API key status toggled successfully: {api_key_id}, new_status={api_key.is_active}" + ) + return True + except KnowhereException: + raise + except Exception as exc: + logger.error(f"Failed to toggle API key status: {exc}") + await session.rollback() + raise APIKeyOperationException( + internal_message=f"Failed to toggle API key status: {str(exc)}", + original_exception=exc, + ) diff --git a/apps/api/app/services/auth/api_key_service.py b/apps/api/app/services/auth/api_key_service.py deleted file mode 100644 index f2b22e4da..000000000 --- a/apps/api/app/services/auth/api_key_service.py +++ /dev/null @@ -1,362 +0,0 @@ -"""API key management service.""" - -from __future__ import annotations - -import asyncio -import json -from datetime import datetime, timezone -from typing import List, Optional - -from app.repositories.api_key_repository import APIKeyRepository -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.config import redis_pool_manager -from shared.core.database import get_db_context -from shared.core.exceptions.domain_exceptions import ( - APIKeyOperationException, - KnowhereException, - NotFoundException, - ValidationException, -) -from shared.models.database.api_key import APIKey -from shared.utils.api_keys import generate_api_key, hash_api_key, mask_api_key - -from shared.services.redis.redis_service import RedisService - -_API_KEY_USER_CACHE_TTL_SECONDS: int = 3600 - - -class APIKeyService: - """API key management service.""" - - _instance: "APIKeyService | None" = None - - def __new__(cls) -> "APIKeyService": - """Return the singleton API-key service object.""" - if cls._instance is None: - cls._instance = super().__new__(cls) - return cls._instance - - def __init__(self) -> None: - if hasattr(self, "repository"): - return - self.repository = APIKeyRepository() - - @classmethod - def get_instance(cls) -> "APIKeyService": - """Return the singleton API-key service instance.""" - return cls() - - def _mask_api_key(self, api_key: str) -> str: - """Mask an API key, exposing only the first 8 and last 4 characters.""" - return mask_api_key(api_key) - - async def create_api_key( - self, - session: AsyncSession, - user_id: str, - name: str, - enabled_modules: Optional[List[str]] = None, - expires_at: Optional[datetime] = None, - ) -> str: - """Create an API key.""" - key_count = await self.repository.count_by_user(session, user_id) - if key_count >= 10: - raise ValidationException( - user_message="Maximum API Key limit reached (10)", - violations=[ - { - "field": "api_keys", - "description": "User has reached the maximum API Key limit", - } - ], - ) - - existing_key = await self.repository.get_by_user_and_name( - session, user_id, name - ) - if existing_key: - raise ValidationException( - user_message="API Key name already exists", - violations=[ - { - "field": "name", - "description": f"An API Key with name '{name}' already exists", - } - ], - ) - - api_key = generate_api_key() - key_hash = hash_api_key(api_key) - key_mask = mask_api_key(api_key) - - api_key_record = APIKey( - user_id=user_id, - key_hash=key_hash, - key_mask=key_mask, - name=name, - enabled_modules=enabled_modules or ["all"], - expires_at=expires_at, - ) - - await self.repository.create(session, api_key_record) - - return api_key - - async def validate_api_key( - self, session: AsyncSession, api_key: str - ) -> Optional[str]: - """Validate API key against DB, return user_id or None.""" - key_hash: str = hash_api_key(api_key) - cached_user_id = await self._get_cached_user_id( - redis_pool_manager.get_redis_service(), - key_hash, - ) - if cached_user_id is not None: - return cached_user_id - - api_key_record = await self.repository.get_by_key_hash(session, key_hash) - if not api_key_record or not api_key_record.is_valid(): - return None - - self._schedule_last_used_update(str(api_key_record.id)) - user_id = str(api_key_record.user_id) - await self._set_cached_user_id( - redis_pool_manager.get_redis_service(), - key_hash, - user_id, - self._resolve_api_key_cache_ttl_seconds(api_key_record.expires_at), - ) - return user_id - - @staticmethod - def _get_user_id_key(api_key_hash: str) -> str: - """Return the Redis key for an API-key hash to user ID lookup.""" - return f"api-key:user-id:{api_key_hash}" - - @staticmethod - def _get_user_api_keys_key(user_id: str) -> str: - """Return the Redis reverse-index key for a user's API-key hashes.""" - return f"api-key:user-hashes:{user_id}" - - async def _get_cached_user_id( - self, - redis_service: RedisService, - api_key_hash: str, - ) -> str | None: - """Return cached API-key user ID or None on miss/cache failure.""" - try: - raw_user_id = await redis_service.get(self._get_user_id_key(api_key_hash)) - return self._coerce_user_id(raw_user_id) - except Exception: - logger.warning("api_key_service: failed to read API-key user cache") - return None - - async def _set_cached_user_id( - self, - redis_service: RedisService, - api_key_hash: str, - user_id: str, - ttl_seconds: int, - ) -> None: - """Cache a validated API-key to user ID lookup.""" - effective_ttl_seconds = min(_API_KEY_USER_CACHE_TTL_SECONDS, ttl_seconds) - user_id_key = self._get_user_id_key(api_key_hash) - user_api_keys_key = self._get_user_api_keys_key(user_id) - - try: - await redis_service.set(user_id_key, user_id, ttl=effective_ttl_seconds) - await redis_service.sadd(user_api_keys_key, api_key_hash) - reverse_ttl_seconds = await redis_service.ttl(user_api_keys_key) - if ( - reverse_ttl_seconds in (-2, -1) - or reverse_ttl_seconds < effective_ttl_seconds - ): - await redis_service.expire(user_api_keys_key, effective_ttl_seconds) - except Exception: - logger.warning( - "api_key_service: failed to write API-key user cache for user_id={}", - user_id, - ) - - async def _invalidate_cached_api_key_user_id( - self, - redis_service: RedisService, - user_id: str, - api_key_hash: str, - ) -> None: - """Delete one API-key to user ID cache entry.""" - try: - await redis_service.delete(self._get_user_id_key(api_key_hash)) - await redis_service.srem(self._get_user_api_keys_key(user_id), api_key_hash) - except Exception: - logger.warning( - "api_key_service: failed to invalidate API-key cache for user_id={}", - user_id, - ) - - def _coerce_user_id(self, raw_user_id: object) -> str | None: - """Return a typed user ID from current or legacy Redis values.""" - if isinstance(raw_user_id, str): - try: - parsed_user_id: object = json.loads(raw_user_id) - except json.JSONDecodeError: - return raw_user_id - else: - parsed_user_id = raw_user_id - - if isinstance(parsed_user_id, str): - return parsed_user_id - - if isinstance(parsed_user_id, dict): - legacy_user_id = parsed_user_id.get("user_id") - if isinstance(legacy_user_id, str): - return legacy_user_id - - return None - - def _resolve_api_key_cache_ttl_seconds(self, expires_at: datetime | None) -> int: - """Resolve cache TTL for an API-key lookup without exceeding key expiry.""" - if expires_at is None: - return _API_KEY_USER_CACHE_TTL_SECONDS - - expires_at_utc = expires_at - if expires_at_utc.tzinfo is None: - expires_at_utc = expires_at_utc.replace(tzinfo=timezone.utc) - - now = datetime.now(timezone.utc) - remaining_seconds = int((expires_at_utc - now).total_seconds()) - return max(1, min(_API_KEY_USER_CACHE_TTL_SECONDS, remaining_seconds)) - - async def revoke_api_key( - self, session: AsyncSession, api_key_id: str, user_id: str - ) -> bool: - """Revoke an API key by deleting it directly.""" - logger.info(f"Revoking API key: api_key_id={api_key_id}, user_id={user_id}") - - api_key = await self.repository.get_by_id(session, api_key_id) - - if not api_key: - logger.warning("API key does not exist") - raise NotFoundException( - resource="APIKey", - resource_id=api_key_id, - internal_message="API Key not found", - ) - - if str(api_key.user_id) != user_id: - logger.warning( - f"User ID mismatch: api_key.user_id={api_key.user_id}, user_id={user_id}" - ) - raise NotFoundException( - resource="APIKey", - resource_id=api_key_id, - internal_message="API Key not found or does not belong to user", - ) - - success = await self.repository.delete_by_id(session, api_key_id) - logger.info(f"Delete result: {success}") - - if success: - await session.commit() - logger.info("Transaction committed") - await self._invalidate_cached_api_key_user_id( - redis_pool_manager.get_redis_service(), - user_id, - api_key.key_hash, - ) - - return success - - async def list_user_api_keys( - self, session: AsyncSession, user_id: str - ) -> List[dict]: - """List a user's API keys, including disabled ones that are still valid.""" - api_keys = await self.repository.get_unexpired_by_user_id(session, user_id) - return [ - { - "id": str(api_key.id), - "name": api_key.name, - "api_key": api_key.key_mask - or f"sk_{api_key.id[:8]}••••••••••••••••••••••••••••••••••••••••", - "enabled_modules": api_key.enabled_modules, - "is_active": api_key.is_active, - "created_at": api_key.created_at, - "last_used_at": api_key.last_used_at, - "expires_at": api_key.expires_at, - } - for api_key in api_keys - ] - - def _schedule_last_used_update(self, api_key_id: str) -> None: - """Schedule a best-effort background update for api_keys.last_used_at.""" - try: - asyncio.create_task( - self._update_last_used_best_effort(api_key_id), - name=f"api_key_last_used:{api_key_id}", - ) - except Exception as e: - logger.warning( - f"Failed to schedule API key last-used update (ignored): {e}" - ) - - async def _update_last_used_best_effort(self, api_key_id: str) -> None: - """Best-effort async update; failures are logged but never propagated.""" - try: - async with get_db_context() as db: - await self.repository.update_last_used(db, api_key_id) - except Exception as e: - logger.warning(f"Failed to update API key last-used time (ignored): {e}") - - async def get_api_key( - self, session: AsyncSession, user_id: str, api_key_id: str - ) -> Optional[APIKey]: - """Get a single API key for a user.""" - try: - api_key = await self.repository.get(session, api_key_id) - if api_key and api_key.user_id == user_id: - return api_key - return None - except KnowhereException: - raise - except Exception as e: - logger.error(f"Failed to get API key: {e}") - raise APIKeyOperationException( - internal_message=f"Failed to get API key: {str(e)}", - original_exception=e, - ) - - async def toggle_api_key( - self, session: AsyncSession, user_id: str, api_key_id: str - ) -> bool: - """Enable or disable an API key.""" - try: - api_key = await self.repository.get(session, api_key_id) - if not api_key or str(api_key.user_id) != user_id: - return False - - api_key.is_active = not api_key.is_active - await session.commit() - await session.refresh(api_key) - - if not api_key.is_active: - await self._invalidate_cached_api_key_user_id( - redis_pool_manager.get_redis_service(), - user_id, - api_key.key_hash, - ) - - logger.info( - f"API key status toggled successfully: {api_key_id}, new_status={api_key.is_active}" - ) - return True - except KnowhereException: - raise - except Exception as e: - logger.error(f"Failed to toggle API key status: {e}") - await session.rollback() - raise APIKeyOperationException( - internal_message=f"Failed to toggle API key status: {str(e)}", - original_exception=e, - ) diff --git a/apps/api/app/services/auth/current_user_authentication_service.py b/apps/api/app/services/auth/current_user_authentication_service.py new file mode 100644 index 000000000..a32c02068 --- /dev/null +++ b/apps/api/app/services/auth/current_user_authentication_service.py @@ -0,0 +1,96 @@ +"""Current-user authentication workflow.""" + +from __future__ import annotations + +from app.services.auth.api_key_authentication_service import ( + APIKeyAuthenticationService, +) +from app.services.auth.dashboard_jwt_authentication_service import ( + DashboardJWTAuthenticationService, + get_dashboard_jwt_authentication_service, +) +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import AuthException +from shared.models.database.user import User +from shared.utils.api_keys import is_api_key_token + + +class CurrentUserAuthenticationService: + """Authenticate API keys or Dashboard JWTs and return the current user ID.""" + + def __init__( + self, + *, + api_key_authentication_service: APIKeyAuthenticationService | None = None, + dashboard_jwt_authentication_service: ( + DashboardJWTAuthenticationService | None + ) = None, + ) -> None: + self._api_key_authentication_service = ( + api_key_authentication_service or APIKeyAuthenticationService() + ) + self._dashboard_jwt_authentication_service = ( + dashboard_jwt_authentication_service + or get_dashboard_jwt_authentication_service() + ) + + async def authenticate_authorization_header( + self, + session: AsyncSession, + authorization: str | None, + ) -> str: + """Authenticate an Authorization header and return the owning user ID.""" + token = self._extract_bearer_token(authorization) + + if is_api_key_token(token): + user_id = await self._api_key_authentication_service.validate_api_key( + session, + token, + ) + if user_id: + return user_id + raise AuthException(user_message="Invalid API Key") + + user_id = self._dashboard_jwt_authentication_service.decode_user_id(token) + await self._ensure_authenticated_user_exists(session, user_id) + return user_id + + @staticmethod + def _extract_bearer_token(authorization: str | None) -> str: + if not authorization: + raise AuthException( + user_message="Authentication required. Provide Authorization header." + ) + + scheme, _, token = authorization.partition(" ") + if scheme.lower() != "bearer" or not token: + raise AuthException(user_message="Invalid Authorization header format") + + return token + + @staticmethod + async def _ensure_authenticated_user_exists( + session: AsyncSession, + user_id: str, + ) -> None: + result = await session.execute(select(User.id).where(User.id == user_id).limit(1)) + if result.scalar_one_or_none() is not None: + return + + raise AuthException( + user_message="Invalid authentication credentials", + internal_message=( + "Authenticated user id is not present in the user table: " + f"user_id={user_id}" + ), + ) + + +_current_user_authentication_service = CurrentUserAuthenticationService() + + +def get_current_user_authentication_service() -> CurrentUserAuthenticationService: + """Return the process-wide current-user authentication service.""" + return _current_user_authentication_service diff --git a/apps/api/app/services/auth/dashboard_jwt_authentication_service.py b/apps/api/app/services/auth/dashboard_jwt_authentication_service.py new file mode 100644 index 000000000..07cc6ef4c --- /dev/null +++ b/apps/api/app/services/auth/dashboard_jwt_authentication_service.py @@ -0,0 +1,91 @@ +"""Dashboard JWT authentication workflow.""" + +from __future__ import annotations + +import threading +from datetime import timedelta +from typing import Any + +import jwt +from jwt import PyJWKClient +from loguru import logger + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import AuthException + +JWKS_ENDPOINT_PATH = "/api/auth/jwks" +JWKS_CACHE_TTL_SECONDS = 60 * 60 + + +class DashboardJWTAuthenticationService: + """Validate Dashboard-issued JWTs through the configured JWKS endpoint.""" + + def __init__(self) -> None: + self._jwks_client: PyJWKClient | None = None + self._jwks_client_lock = threading.Lock() + + def decode_user_id(self, token: str) -> str: + """Decode and validate a JWT, returning its authenticated user ID.""" + try: + key = self._get_verification_key(token) + payload: dict[str, Any] = jwt.decode( + token, + key, + algorithms=["HS256", "RS256", "EdDSA"], + leeway=timedelta(seconds=30), + options={"verify_aud": False}, + ) + user_id = payload.get("id") + if not isinstance(user_id, str) or not user_id: + raise AuthException(user_message="Token missing 'id' claim") + return user_id + except jwt.ExpiredSignatureError: + raise AuthException(user_message="Token has expired") + except jwt.InvalidTokenError as exc: + logger.warning(f"Invalid JWT token: {exc}") + raise AuthException(user_message="Invalid token") + + def _get_verification_key(self, token: str) -> Any: + """Resolve the JWT verification key from the Dashboard JWKS endpoint.""" + try: + jwks_client = self._get_jwks_client() + signing_key = jwks_client.get_signing_key_from_jwt(token) + return signing_key.key + except jwt.PyJWKClientError as exc: + logger.error(f"Failed to fetch JWKS: {exc}") + raise AuthException( + internal_message=( + "Failed to fetch verification key from JWKS endpoint: " + f"{exc}" + ) + ) + except jwt.PyJWKSetError as exc: + logger.error(f"Invalid JWKS format: {exc}") + raise AuthException(internal_message=f"Invalid JWKS format: {exc}") + + def _get_jwks_client(self) -> PyJWKClient: + """Return a cached JWKS client for Dashboard token verification.""" + if self._jwks_client is None: + with self._jwks_client_lock: + if self._jwks_client is None: + jwks_url = ( + f"{settings.INTERNAL_DASHBOARD_ENDPOINT}" + f"{JWKS_ENDPOINT_PATH}" + ) + self._jwks_client = PyJWKClient( + jwks_url, + cache_jwk_set=True, + lifespan=JWKS_CACHE_TTL_SECONDS, + timeout=30, + ) + logger.info(f"Initialized JWKS client with endpoint: {jwks_url}") + + return self._jwks_client + + +_dashboard_jwt_authentication_service = DashboardJWTAuthenticationService() + + +def get_dashboard_jwt_authentication_service() -> DashboardJWTAuthenticationService: + """Return the process-wide Dashboard JWT authentication service.""" + return _dashboard_jwt_authentication_service diff --git a/apps/api/app/services/billing/billing_command_workflow.py b/apps/api/app/services/billing/billing_command_workflow.py new file mode 100644 index 000000000..bd3d62073 --- /dev/null +++ b/apps/api/app/services/billing/billing_command_workflow.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +from app.services.billing.stripe_purchase_service import StripePurchaseService +from app.services.billing.stripe_webhook_service import StripeWebhookService +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.billing import MicroDollar +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import StripeServiceException +from shared.models.database.user import User +from shared.models.schemas.billing import ( + BuyCreditsPackageRequest, + BuyCreditsRequest, + CheckoutSessionResponse, + PaymentIntentResponse, +) + + +class BillingCommandWorkflow: + async def buy_credits( + self, + *, + request: BuyCreditsRequest, + user_id: str, + ) -> PaymentIntentResponse: + stripe_purchase_service = StripePurchaseService() + try: + amount_cny = request.credits_amount * 0.02 + amount_cents = int(amount_cny * 100) + payment_intent = await stripe_purchase_service.create_payment_intent( + user_id=user_id, + amount=amount_cents, + credits_amount=MicroDollar.from_dollars(request.credits_amount).amount, + currency="cny", + ) + + return PaymentIntentResponse( + client_secret=payment_intent["client_secret"], + payment_intent_id=payment_intent["payment_intent_id"], + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to buy credits: {str(exc)}" + ) + + async def buy_credits_package( + self, + db: AsyncSession, + *, + request: BuyCreditsPackageRequest, + user_id: str, + ) -> CheckoutSessionResponse: + stripe_purchase_service = StripePurchaseService() + try: + result = await db.execute(select(User.email).where(User.id == user_id)) + user_email = result.scalar_one_or_none() + + frontend_url = settings.FRONTEND_URL + success_url = f"{frontend_url}/billing?success=true&type=credits_package" + cancel_url = f"{frontend_url}/billing?canceled=true" + + checkout_url = await stripe_purchase_service.create_credits_package_checkout_session( + db=db, + user_id=user_id, + price_id=request.price_id, + success_url=success_url, + cancel_url=cancel_url, + quantity=request.quantity, + email=user_email, + ) + + return CheckoutSessionResponse(checkout_url=checkout_url, session_id="") + except Exception as exc: + raise StripeServiceException( + internal_message=( + "Failed to create credits package purchase: " + f"{str(exc)}" + ) + ) + + async def handle_stripe_webhook( + self, + db: AsyncSession, + *, + payload: bytes, + stripe_signature: str | None, + ) -> dict[str, object]: + stripe_webhook_service = StripeWebhookService() + try: + if not stripe_signature: + raise StripeServiceException( + internal_message="Missing stripe-signature header" + ) + return await stripe_webhook_service.handle_webhook( + db, + payload=payload, + sig_header=stripe_signature, + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to handle webhook: {str(exc)}" + ) diff --git a/apps/api/app/services/billing/billing_read_model.py b/apps/api/app/services/billing/billing_read_model.py new file mode 100644 index 000000000..ccc17fe09 --- /dev/null +++ b/apps/api/app/services/billing/billing_read_model.py @@ -0,0 +1,288 @@ +from __future__ import annotations + +from typing import Optional + +from app.services.billing.price_config_service import PriceConfigService +from pydantic import BaseModel +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.billing import MicroDollar +from shared.core.exceptions.domain_exceptions import StripeServiceException +from shared.models.database.credits_transaction import CreditsTransaction +from shared.models.database.job import Job +from shared.models.database.stripe_price_config import StripePriceConfig +from shared.models.schemas.billing import ( + CreditsBalanceResponse, + TransactionHistoryResponse, + UsageStatsResponse, +) +from shared.services.billing import CreditsService + + +class ParseUsageResponse(BaseModel): + request_total: int + mom_growth: float + credits_used: float + estimated_amount: Optional[float] + success_rate: float + avg_processing_time: float + + +class BillingReadModel: + def __init__( + self, + *, + price_config_service: PriceConfigService | None = None, + credits_service: CreditsService | None = None, + ) -> None: + self._price_config_service = price_config_service or PriceConfigService() + self._credits_service = credits_service or CreditsService() + + async def get_credits_balance( + self, + db: AsyncSession, + *, + user_id: str, + ) -> CreditsBalanceResponse: + try: + await self._credits_service.ensure_user_initialized(db, user_id) + await db.commit() + + balance_micro_dollar = await self._credits_service.get_balance(db, user_id) + return CreditsBalanceResponse( + credits_balance=MicroDollar(balance_micro_dollar).to_credit() + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get credits balance: {str(exc)}" + ) + + async def get_usage_stats( + self, + db: AsyncSession, + *, + user_id: str, + period: str, + ) -> UsageStatsResponse: + try: + stats = await self._credits_service.get_usage_stats(db, user_id, period) + return UsageStatsResponse( + period=stats["period"], + total_credits_used=MicroDollar(stats["total_used"]).to_credit(), + api_calls_count=stats["transaction_count"], + success_rate=95.0, + average_response_time=stats.get("avg_response_time", 0), + top_endpoints=[], + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get usage statistics: {str(exc)}" + ) + + async def get_parse_usage_overview( + self, + db: AsyncSession, + *, + user_id: str, + ) -> ParseUsageResponse: + try: + total_micro_credits_used = await self._load_total_parse_micro_credits_used( + db, + user_id=user_id, + ) + success_rate, avg_processing_time = await self._load_parse_job_usage_stats( + db, + user_id=user_id, + ) + estimated_amount = await self._estimate_parse_usage_amount( + db, + total_micro_credits_used=total_micro_credits_used, + ) + + return ParseUsageResponse( + request_total=0, + mom_growth=0.0, + credits_used=MicroDollar(total_micro_credits_used).to_credit() or 0, + estimated_amount=estimated_amount, + success_rate=round(success_rate, 2), + avg_processing_time=avg_processing_time, + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get parse usage overview: {str(exc)}" + ) + + async def get_transaction_history( + self, + db: AsyncSession, + *, + user_id: str, + limit: int, + ) -> list[TransactionHistoryResponse]: + try: + transactions = await self._credits_service.get_transaction_history( + db, + user_id, + limit, + ) + return [ + TransactionHistoryResponse( + id=transaction.id, + credits_amount=MicroDollar(transaction.credits_amount).to_credit(), + transaction_type=transaction.transaction_type, + description=transaction.description, + created_at=transaction.created_at, + ) + for transaction in transactions + ] + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get transaction history: {str(exc)}" + ) + + async def get_price_configs( + self, + db: AsyncSession, + *, + product_type: str | None, + ) -> dict[str, list[dict[str, object]]]: + try: + if product_type == "subscription": + configs = await self._price_config_service.repository.get_all_active(db) + return { + "subscriptions": [ + _subscription_config_payload(config) + for config in configs + if config.product_type == "subscription" + ], + "credits_packages": [], + } + + if product_type == "credits_package": + credits_configs = await self._price_config_service.get_all_credits_packages( + db + ) + return { + "subscriptions": [], + "credits_packages": [ + _credits_package_config_payload(config) + for config in credits_configs + ], + } + + configs = await self._price_config_service.repository.get_all_active(db) + return { + "subscriptions": [ + _subscription_config_payload(config) + for config in configs + if config.product_type == "subscription" + ], + "credits_packages": [ + _credits_package_config_payload(config) + for config in configs + if config.product_type == "credits_package" + ], + } + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get price configurations: {str(exc)}" + ) + + async def _load_total_parse_micro_credits_used( + self, + db: AsyncSession, + *, + user_id: str, + ) -> int: + credits_row = await db.execute( + select(func.coalesce(func.sum(CreditsTransaction.credits_amount), 0)) + .where(CreditsTransaction.user_id == user_id) + .where(CreditsTransaction.transaction_type.in_(["usage", "refund"])) + ) + return int(abs(credits_row.scalar_one() or 0)) + + async def _load_parse_job_usage_stats( + self, + db: AsyncSession, + *, + user_id: str, + ) -> tuple[float, float]: + job_row = await db.execute( + select( + func.count().filter(Job.status == "done").label("done_cnt"), + func.count() + .filter(Job.status.in_(["done", "failed"])) + .label("terminal_cnt"), + func.avg(func.extract("epoch", Job.updated_at - Job.created_at)) + .filter(Job.status.in_(["done", "failed"])) + .label("avg_secs"), + ).where(Job.user_id == user_id) + ) + job_stats = job_row.first() or (0, 0, 0.0) + done_count = getattr(job_stats, "done_cnt", 0) or 0 + terminal_count = getattr(job_stats, "terminal_cnt", 0) or 0 + success_rate = ( + done_count / terminal_count * 100 if terminal_count > 0 else 0.0 + ) + avg_processing_time = round( + float(getattr(job_stats, "avg_secs", 0.0) or 0.0), + 2, + ) + return success_rate, avg_processing_time + + async def _estimate_parse_usage_amount( + self, + db: AsyncSession, + *, + total_micro_credits_used: int, + ) -> float | None: + price_row = await db.execute( + select(StripePriceConfig) + .where(StripePriceConfig.product_type == "credits_package") + .where(StripePriceConfig.is_active.is_(True)) + .order_by(StripePriceConfig.created_at) + .limit(1) + ) + price_cfg = price_row.scalar_one_or_none() + if not price_cfg or not price_cfg.credits_amount or price_cfg.credits_amount <= 0: + return None + + return round( + price_cfg.amount_cents + * total_micro_credits_used + / (100 * price_cfg.credits_amount), + 4, + ) + + +def _subscription_config_payload(config: StripePriceConfig) -> dict[str, object]: + metadata = config.extra_metadata or {} + return { + "id": config.plan_id, + "plan_id": config.plan_id, + "price_id": config.price_id, + "name": metadata.get("display_name", config.plan_id.upper()), + "description": metadata.get("description", ""), + "features": metadata.get("features", []), + "popular": metadata.get("frontend_config", {}).get("popular", False), + "amount_cents": config.amount_cents, + "currency": config.currency, + "metadata": metadata, + } + + +def _credits_package_config_payload(config: StripePriceConfig) -> dict[str, object]: + metadata = config.extra_metadata or {} + credit_amount = MicroDollar(config.credits_amount).to_credit() + return { + "id": config.plan_id, + "plan_id": config.plan_id, + "price_id": config.price_id, + "name": metadata.get("display_name", f"{credit_amount} Credits"), + "description": metadata.get("description", ""), + "credits_amount": credit_amount, + "amount_cents": config.amount_cents, + "currency": config.currency, + "metadata": metadata, + } diff --git a/apps/api/app/services/billing/billing_workflow_service.py b/apps/api/app/services/billing/billing_workflow_service.py new file mode 100644 index 000000000..1c87939f5 --- /dev/null +++ b/apps/api/app/services/billing/billing_workflow_service.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +from app.services.billing.billing_command_workflow import BillingCommandWorkflow +from app.services.billing.billing_read_model import BillingReadModel, ParseUsageResponse +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.schemas.billing import ( + BuyCreditsPackageRequest, + BuyCreditsRequest, + CheckoutSessionResponse, + CreditsBalanceResponse, + PaymentIntentResponse, + TransactionHistoryResponse, + UsageStatsResponse, +) + +__all__ = ["BillingWorkflowService", "ParseUsageResponse"] + + +class BillingWorkflowService: + def __init__( + self, + *, + command_workflow: BillingCommandWorkflow | None = None, + read_model: BillingReadModel | None = None, + ) -> None: + self._command_workflow = command_workflow or BillingCommandWorkflow() + self._read_model = read_model or BillingReadModel() + + async def buy_credits( + self, + *, + request: BuyCreditsRequest, + user_id: str, + ) -> PaymentIntentResponse: + return await self._command_workflow.buy_credits( + request=request, + user_id=user_id, + ) + + async def get_credits_balance( + self, + db: AsyncSession, + *, + user_id: str, + ) -> CreditsBalanceResponse: + return await self._read_model.get_credits_balance(db, user_id=user_id) + + async def get_usage_stats( + self, + db: AsyncSession, + *, + user_id: str, + period: str, + ) -> UsageStatsResponse: + return await self._read_model.get_usage_stats( + db, + user_id=user_id, + period=period, + ) + + async def get_parse_usage_overview( + self, + db: AsyncSession, + *, + user_id: str, + ) -> ParseUsageResponse: + return await self._read_model.get_parse_usage_overview(db, user_id=user_id) + + async def get_transaction_history( + self, + db: AsyncSession, + *, + user_id: str, + limit: int, + ) -> list[TransactionHistoryResponse]: + return await self._read_model.get_transaction_history( + db, + user_id=user_id, + limit=limit, + ) + + async def get_price_configs( + self, + db: AsyncSession, + *, + product_type: str | None, + ) -> dict[str, list[dict[str, object]]]: + return await self._read_model.get_price_configs( + db, + product_type=product_type, + ) + + async def buy_credits_package( + self, + db: AsyncSession, + *, + request: BuyCreditsPackageRequest, + user_id: str, + ) -> CheckoutSessionResponse: + return await self._command_workflow.buy_credits_package( + db, + request=request, + user_id=user_id, + ) + + async def handle_stripe_webhook( + self, + db: AsyncSession, + *, + payload: bytes, + stripe_signature: str | None, + ) -> dict[str, object]: + return await self._command_workflow.handle_stripe_webhook( + db, + payload=payload, + stripe_signature=stripe_signature, + ) diff --git a/apps/api/app/services/billing/stripe_credits_settlement_service.py b/apps/api/app/services/billing/stripe_credits_settlement_service.py new file mode 100644 index 000000000..9091dd6e6 --- /dev/null +++ b/apps/api/app/services/billing/stripe_credits_settlement_service.py @@ -0,0 +1,283 @@ +from __future__ import annotations + +from typing import Any + +from app.repositories.payment_record_repository import PaymentRecordRepository +from app.services.billing.price_config_service import PriceConfigService +from app.services.rate_limit.tier_service import TierService +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + KnowhereException, + StripeServiceException, +) +from shared.models.database.payment_record import PaymentRecord +from shared.services.billing import CreditsService +from shared.utils.utc_now import utc_now_naive + + +class StripeCreditsSettlementService: + def __init__( + self, + *, + payment_record_repository: PaymentRecordRepository | None = None, + price_config_service: PriceConfigService | None = None, + credits_service: CreditsService | None = None, + ) -> None: + self._payment_record_repository = ( + payment_record_repository or PaymentRecordRepository() + ) + self._price_config_service = price_config_service or PriceConfigService() + self._credits_service = credits_service or CreditsService() + + async def handle_checkout_completed( + self, + db: AsyncSession, + *, + event: dict[str, Any], + ) -> dict[str, object]: + session = event["data"]["object"] + session_id = str(session["id"]) + mode = session.get("mode") + metadata = session.get("metadata", {}) + user_id = metadata.get("user_id") + payment_type = metadata.get("type") + quantity = int(metadata.get("quantity", 1)) + + if not user_id: + logger.warning( + f"Checkout session {session_id} is missing user_id metadata; likely a test event, skipping" + ) + return { + "status": "ignored", + "message": "Missing user_id metadata (likely test event)", + "checkout_session_id": session_id, + "event_type": "checkout.session.completed", + } + + if await self._payment_record_repository.is_processed( + db, + checkout_session_id=session_id, + ): + logger.info(f"Checkout session {session_id} already processed, skipping") + return { + "status": "ignored", + "message": "Already processed", + "checkout_session_id": session_id, + } + + payment_metadata = { + "session_id": session_id, + "stripe_session": session, + } + payment_record = PaymentRecord( + checkout_session_id=session_id, + user_id=user_id, + payment_type=payment_type or "unknown", + amount_cents=session.get("amount_total", 0), + currency=session.get("currency", "cny").upper(), + status="pending", + extra_metadata=payment_metadata, + ) + db.add(payment_record) + await db.flush() + + try: + if mode != "payment" or payment_type != "credits_package": + logger.warning(f"Unknown payment type: mode={mode}, type={payment_type}") + return {"status": "ignored", "message": "Unknown payment type"} + + price_id = metadata.get("price_id") + if not price_id: + logger.error(f"Incomplete Credits pack info: price_id={price_id}") + return {"status": "error", "message": "Missing price_id"} + + price_config = await self._price_config_service.get_price_config(db, price_id) + configured_credits_amount = price_config.credits_amount + if configured_credits_amount is None: + logger.error( + f"Credits amount is not configured for price ID {price_id}" + ) + return { + "status": "error", + "message": "Credits amount not configured", + } + credits_amount = configured_credits_amount * quantity + + product_description = f"Credits pack - {credits_amount} Credits" + if price_config.extra_metadata and price_config.extra_metadata.get( + "description" + ): + product_description = str( + price_config.extra_metadata.get("description") + ) + + payment_record.extra_metadata = { + **payment_metadata, + "product_description": product_description, + "price_id": price_id, + "credits_amount": credits_amount, + "product_metadata": price_config.extra_metadata or {}, + } + await self._credits_service.add_credits( + session=db, + user_id=user_id, + amount=credits_amount, + reason=f"Purchase credits pack: {product_description}", + stripe_payment_id=session.get("payment_intent"), + ) + payment_record.status = "succeeded" + payment_record.credits_amount = credits_amount + payment_record.processed_at = utc_now_naive() + + await TierService.refresh_tier(user_id, db) + await db.commit() + await db.refresh(payment_record) + + logger.info( + f"Credits pack purchase succeeded: user_id={user_id}, credits={credits_amount}, price_id={price_id}" + ) + return { + "status": "success", + "event_type": "checkout.session.completed", + "user_id": user_id, + "credits_amount": credits_amount, + "payment_type": "credits_package", + } + except KnowhereException: + raise + except Exception as exc: + logger.error( + f"Failed to process checkout.session.completed: {exc}", + exc_info=True, + ) + payment_record.status = "failed" + payment_record.extra_metadata = { + **(payment_record.extra_metadata or {}), + "error": str(exc), + } + await db.commit() + raise StripeServiceException( + internal_message=( + "Failed to process checkout.session.completed: " + f"{str(exc)}" + ), + original_exception=exc, + ) + + async def handle_payment_intent_succeeded( + self, + db: AsyncSession, + *, + event: dict[str, Any], + ) -> dict[str, object]: + payment_intent = event["data"]["object"] + payment_intent_id = str(payment_intent["id"]) + metadata = payment_intent.get("metadata", {}) + user_id = metadata.get("user_id") + payment_type = metadata.get("type") + + if payment_type != "credits": + logger.info( + f"PaymentIntent {payment_intent_id} is not a Credits payment, skipping" + ) + return {"status": "ignored", "payment_intent_id": payment_intent_id} + + if not user_id: + logger.warning( + f"PaymentIntent {payment_intent_id} is missing user_id metadata; likely a test event, skipping" + ) + return { + "status": "ignored", + "message": "Missing user_id metadata (likely test event)", + "payment_intent_id": payment_intent_id, + } + + if await self._payment_record_repository.is_processed( + db, + payment_intent_id=payment_intent_id, + ): + logger.info( + f"PaymentIntent {payment_intent_id} already processed, skipping" + ) + return { + "status": "ignored", + "message": "Already processed", + "payment_intent_id": payment_intent_id, + } + + payment_metadata = { + "payment_intent_id": payment_intent_id, + "stripe_payment_intent": payment_intent, + } + payment_record = PaymentRecord( + payment_intent_id=payment_intent_id, + user_id=user_id, + payment_type="credits_package", + amount_cents=payment_intent.get("amount", 0), + currency=payment_intent.get("currency", "cny").upper(), + status="pending", + extra_metadata=payment_metadata, + ) + db.add(payment_record) + await db.flush() + + try: + credits_amount_str = metadata.get("credits_amount") + if not credits_amount_str: + logger.error( + f"PaymentIntent {payment_intent_id} is missing credits_amount" + ) + payment_record.status = "failed" + payment_record.extra_metadata = { + **(payment_record.extra_metadata or {}), + "error": "Missing credits_amount", + } + await db.commit() + return {"status": "error", "message": "Missing credits_amount"} + + credits_amount = int(credits_amount_str) + payment_record.extra_metadata = { + **payment_metadata, + "product_description": f"Credits package - {credits_amount} Credits", + "credits_amount": credits_amount, + "payment_method": "payment_intent", + } + await self._credits_service.add_credits( + session=db, + user_id=user_id, + amount=credits_amount, + reason=f"buy credits - {credits_amount} Credits", + stripe_payment_id=payment_intent_id, + ) + payment_record.status = "succeeded" + payment_record.credits_amount = credits_amount + payment_record.processed_at = utc_now_naive() + + await TierService.refresh_tier(user_id, db) + await db.commit() + await db.refresh(payment_record) + + logger.info( + f"buy credits success: user_id={user_id}, credits={credits_amount}, payment_intent_id={payment_intent_id}" + ) + return { + "status": "success", + "event_type": "payment_intent.succeeded", + "user_id": user_id, + "credits_amount": credits_amount, + "payment_type": "credits_package", + } + except Exception as exc: + logger.error(f"Failed to process Credits purchase: {exc}", exc_info=True) + payment_record.status = "failed" + payment_record.extra_metadata = { + **(payment_record.extra_metadata or {}), + "error": str(exc), + } + await db.commit() + raise StripeServiceException( + internal_message=f"Failed to process Credits purchase: {str(exc)}", + original_exception=exc, + ) diff --git a/apps/api/app/services/billing/stripe_purchase_service.py b/apps/api/app/services/billing/stripe_purchase_service.py new file mode 100644 index 000000000..c5ea9bedd --- /dev/null +++ b/apps/api/app/services/billing/stripe_purchase_service.py @@ -0,0 +1,200 @@ +from __future__ import annotations + +from typing import Any + +import stripe +from app.services.billing.price_config_service import PriceConfigService +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ( + StripeServiceException, + SystemSettingMissingException, + ValidationException, +) +from shared.core.logging import logger +from shared.repositories.credits_repository import CreditsRepository +from shared.services.billing import CreditsService + + +class StripePurchaseService: + def __init__( + self, + *, + price_config_service: PriceConfigService | None = None, + credits_repository: CreditsRepository | None = None, + credits_service: CreditsService | None = None, + ) -> None: + self._configure_stripe_api() + self._price_config_service = price_config_service or PriceConfigService() + self._credits_repository = credits_repository or CreditsRepository() + self._credits_service = credits_service or CreditsService() + + async def create_credits_package_checkout_session( + self, + db: AsyncSession, + *, + user_id: str, + price_id: str, + success_url: str, + cancel_url: str, + quantity: int, + email: str | None = None, + ) -> str: + try: + config = await self._price_config_service.get_price_config(db, price_id) + if not config.is_credits_package(): + raise ValidationException( + user_message="Invalid price configuration", + violations=[ + { + "field": "price_id", + "description": f"Price ID {price_id} is not a credits package", + } + ], + ) + + customer_id = await self._resolve_customer_id( + db, + user_id=user_id, + email=email, + ) + metadata = { + "user_id": str(user_id), + "price_id": str(price_id), + "type": "credits_package", + "credits_amount": ( + str(config.credits_amount) if config.credits_amount else None + ), + "quantity": str(quantity), + } + session_params: dict[str, Any] = { + "customer": customer_id, + "customer_update": {"address": "auto"}, + "client_reference_id": str(user_id), + "line_items": [ + { + "price": price_id, + "quantity": quantity, + } + ], + "mode": "payment", + "success_url": success_url, + "cancel_url": cancel_url, + "metadata": metadata, + "payment_intent_data": {"metadata": metadata}, + "allow_promotion_codes": True, + "adaptive_pricing": {"enabled": False}, + "billing_address_collection": "required", + } + session = stripe.checkout.Session.create(**session_params) + await db.commit() + return str(session.url or "") + except stripe.StripeError as exc: + logger.error(f"Stripe credits checkout session failed: {exc}") + raise StripeServiceException( + internal_message=f"Stripe credits checkout session failed: {exc}" + ) + + async def create_payment_intent( + self, + *, + user_id: str, + amount: int, + credits_amount: int, + currency: str = "usd", + ) -> dict[str, str]: + try: + intent = stripe.PaymentIntent.create( + amount=amount, + currency=currency, + automatic_payment_methods={"enabled": True}, + metadata={ + "user_id": user_id, + "type": "credits", + "credits_amount": str(credits_amount), + }, + ) + return { + "client_secret": str(intent.client_secret or ""), + "payment_intent_id": str(intent.id), + } + except stripe.StripeError as exc: + logger.error(f"Failed to create payment intent: {exc}") + raise StripeServiceException( + internal_message=f"Stripe payment intent creation failed: {exc}" + ) + + async def _resolve_customer_id( + self, + db: AsyncSession, + *, + user_id: str, + email: str | None, + ) -> str: + await self._credits_service.ensure_user_initialized(db, user_id) + user_balance = await self._credits_repository.get_user_balance(db, user_id) + if not user_balance: + raise ValidationException( + user_message="Failed to initialize user balance", + violations=[ + { + "field": "user_id", + "description": f"Failed to initialize user balance for {user_id}", + } + ], + ) + + customer_id = user_balance.stripe_customer_id + if customer_id: + return customer_id + + customer_id = self._find_existing_customer_id(email=email) + if customer_id is None: + customer_id = self._create_customer(user_id=user_id, email=email) + + user_balance.stripe_customer_id = customer_id + return customer_id + + def _find_existing_customer_id( + self, + *, + email: str | None, + ) -> str | None: + if not email: + return None + + existing_customers = stripe.Customer.list(email=email, limit=1) + if existing_customers.data: + return str(existing_customers.data[0].id) + return None + + def _create_customer( + self, + *, + user_id: str, + email: str | None, + ) -> str: + if not email: + raise ValidationException( + user_message="Email required for first-time payment", + violations=[ + { + "field": "email", + "description": "Email is required to create a billing profile", + } + ], + ) + + customer = stripe.Customer.create( + email=email, + metadata={"user_id": str(user_id)}, + ) + return str(customer.id) + + def _configure_stripe_api(self) -> None: + if not settings.STRIPE_SECRET_KEY: + raise SystemSettingMissingException( + internal_message="Stripe API key not configured (STRIPE_SECRET_KEY)" + ) + stripe.api_key = settings.STRIPE_SECRET_KEY diff --git a/apps/api/app/services/billing/stripe_refund_reconciliation_service.py b/apps/api/app/services/billing/stripe_refund_reconciliation_service.py new file mode 100644 index 000000000..d770195e1 --- /dev/null +++ b/apps/api/app/services/billing/stripe_refund_reconciliation_service.py @@ -0,0 +1,234 @@ +from __future__ import annotations + +from typing import Any +from uuid import UUID + +from app.repositories.payment_record_repository import PaymentRecordRepository +from app.services.billing.price_config_service import PriceConfigService +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.logging import logger +from shared.models.database.payment_record import PaymentRecord +from shared.services.billing import CreditsService +from shared.utils.utc_now import utc_now_naive + + +class StripeRefundReconciliationService: + def __init__( + self, + *, + payment_record_repository: PaymentRecordRepository | None = None, + price_config_service: PriceConfigService | None = None, + credits_service: CreditsService | None = None, + ) -> None: + self._payment_record_repository = ( + payment_record_repository or PaymentRecordRepository() + ) + self._price_config_service = price_config_service or PriceConfigService() + self._credits_service = credits_service or CreditsService() + + async def reconcile_charge_refund( + self, + db: AsyncSession, + *, + event: dict[str, Any], + ) -> dict[str, object]: + charge = event["data"]["object"] + charge_id = charge.get("id") + refund_items = (charge.get("refunds", {}) or {}).get("data", []) or [] + latest_refund = refund_items[-1] if refund_items else None + + payment_intent_id = charge.get("payment_intent") + refund_id = latest_refund.get("id") if latest_refund else None + currency = (charge.get("currency") or "cny").upper() + idempotency_key = refund_id or f"{charge_id}-refund" + + original_record = None + if payment_intent_id: + original_record = await self._payment_record_repository.get_by_payment_intent_id( + db, + payment_intent_id, + ) + + metadata = charge.get("metadata") or {} + user_id = metadata.get("user_id") or ( + getattr(original_record, "user_id", None) + ) + payment_type = ( + metadata.get("type") + or getattr(original_record, "payment_type", None) + or "refund" + ) + + if not user_id: + logger.error( + f"Refund event is missing user_id; cannot record refund: charge_id={charge_id}" + ) + return { + "status": "error", + "message": "Missing user_id for refund", + "event_type": "charge.refunded", + } + + normalized_user_id = self._normalize_user_id(user_id) + if normalized_user_id is None: + logger.error(f"Invalid user_id format: {user_id}") + return { + "status": "error", + "message": "Invalid user_id format", + "event_type": "charge.refunded", + } + + user_id_str = str(normalized_user_id) + total_refund_amount_cents = charge.get("amount_refunded") or 0 + origin_total_refund_amount_cents = await self._load_recorded_refund_amount( + db, + payment_intent_id=idempotency_key, + user_id=normalized_user_id, + ) + + refund_amount_cents = ( + total_refund_amount_cents - origin_total_refund_amount_cents + ) + if refund_amount_cents <= 0: + logger.info( + f"Refund already processed, skipping: charge_id={charge_id}, refund_id={refund_id}" + ) + return { + "status": "success", + "event_type": "charge.refunded", + "message": "Already processed", + "user_id": normalized_user_id, + "refund_id": refund_id, + } + + credits_refunded = await self._calculate_refunded_credits( + db, + metadata=metadata, + original_record=original_record, + refund_amount_cents=refund_amount_cents, + ) + + if credits_refunded is not None and credits_refunded < 0: + await self._credits_service.add_credits( + session=db, + user_id=user_id_str, + amount=credits_refunded, + reason="Refund adjustment", + transaction_type="refund", + transaction_metadata={"refund_id": refund_id, "charge_id": charge_id}, + ) + + refund_metadata = { + "refund_id": refund_id, + "charge_id": charge_id, + "original_payment_intent_id": payment_intent_id, + "original_payment_record_id": getattr(original_record, "id", None), + "reason": (latest_refund or {}).get("reason"), + "balance_transaction": (latest_refund or {}).get("balance_transaction"), + } + refund_record = PaymentRecord( + payment_intent_id=idempotency_key, + user_id=normalized_user_id, + payment_type=payment_type, + amount_cents=-abs(refund_amount_cents), + currency=currency, + status="succeeded", + credits_amount=credits_refunded, + plan_id=getattr(original_record, "plan_id", None), + stripe_subscription_id=getattr( + original_record, + "stripe_subscription_id", + None, + ), + processed_at=utc_now_naive(), + extra_metadata=refund_metadata, + ) + db.add(refund_record) + await db.commit() + await db.refresh(refund_record) + + logger.info( + f"Refund record created: user_id={normalized_user_id}, amount_cents={refund_record.amount_cents}, " + f"refund_id={refund_id}, charge_id={charge_id}" + ) + return { + "status": "success", + "event_type": "charge.refunded", + "user_id": normalized_user_id, + "refund_amount_cents": abs(refund_amount_cents), + "payment_intent_id": payment_intent_id, + "refund_id": refund_id, + } + + async def _load_recorded_refund_amount( + self, + db: AsyncSession, + *, + payment_intent_id: str, + user_id: UUID, + ) -> int: + result = await db.execute( + select(func.sum(PaymentRecord.amount_cents)) + .where(PaymentRecord.payment_intent_id == payment_intent_id) + .where(PaymentRecord.user_id == user_id) + .where(PaymentRecord.amount_cents < 0) + ) + return int(abs(result.scalar() or 0)) + + async def _calculate_refunded_credits( + self, + db: AsyncSession, + *, + metadata: dict[str, Any], + original_record: PaymentRecord | None, + refund_amount_cents: int, + ) -> int | None: + credits_refunded: int | None = None + price_id = metadata.get("price_id") or ( + getattr(original_record, "extra_metadata", {}) or {} + ).get("price_id") + if price_id: + try: + price_cfg = await self._price_config_service.get_price_config( + db, + price_id, + ) + if price_cfg and price_cfg.amount_cents and price_cfg.credits_amount: + credits_refunded = -int( + price_cfg.credits_amount + * abs(refund_amount_cents) + / abs(price_cfg.amount_cents) + ) + except Exception as exc: + logger.warning( + f"Failed to calculate refunded Credits, price_id={price_id}: {exc}" + ) + credits_refunded = None + + if ( + credits_refunded is None + and original_record + and original_record.credits_amount + and original_record.amount_cents + ): + credits_refunded = -int( + abs(original_record.credits_amount) + * abs(refund_amount_cents) + / abs(original_record.amount_cents) + ) + + return credits_refunded + + def _normalize_user_id( + self, + user_id: str | UUID, + ) -> UUID | None: + if isinstance(user_id, UUID): + return user_id + + try: + return UUID(user_id) + except ValueError: + return None diff --git a/apps/api/app/services/billing/stripe_service.py b/apps/api/app/services/billing/stripe_service.py deleted file mode 100644 index be1792a2e..000000000 --- a/apps/api/app/services/billing/stripe_service.py +++ /dev/null @@ -1,766 +0,0 @@ -"""Stripe payment service.""" - -from typing import Any, Dict, Optional -from uuid import UUID - -import stripe -from app.repositories.payment_record_repository import PaymentRecordRepository -from app.services.billing.price_config_service import PriceConfigService -from app.services.rate_limit.tier_service import TierService -from sqlalchemy import func, select -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import ( - AuthException, - KnowhereException, - StripeServiceException, - SystemSettingMissingException, - ValidationException, -) -from shared.core.logging import logger -from shared.models.database.payment_record import PaymentRecord -from shared.repositories.credits_repository import CreditsRepository -from shared.services.billing import CreditsService -from shared.utils.utc_now import utc_now_naive - - -class StripeService: - """Stripe payment service.""" - - def __init__(self): - if not settings.STRIPE_SECRET_KEY: - raise SystemSettingMissingException( - internal_message="Stripe API key not configured (STRIPE_SECRET_KEY)" - ) - stripe.api_key = settings.STRIPE_SECRET_KEY - self.credits_repo = CreditsRepository() - self.payment_record_repo = PaymentRecordRepository() - self.price_config_service = PriceConfigService() - self.credits_service = CreditsService() - - async def create_checkout_session( - self, - db: AsyncSession, - user_id: str, - plan_id: str, - success_url: str, - cancel_url: str, - ) -> str: - """Create a Stripe Checkout session for a subscription.""" - try: - # Load the Stripe price ID for the requested plan from the database. - price_id = await self.price_config_service.get_plan_price_id(db, plan_id) - - session = stripe.checkout.Session.create( - line_items=[ - { - "price": price_id, - "quantity": 1, - } - ], - mode="subscription", - success_url=success_url, - cancel_url=cancel_url, - metadata={ - "user_id": user_id, - "plan_id": plan_id, - "type": "subscription", - }, - allow_promotion_codes=True, - # Disable Adaptive Pricing to prevent currency switcher from hiding Alipay. - # Alipay handles USD→CNY conversion internally for customers. - adaptive_pricing={"enabled": False}, - ) - return str(session.url or "") - except stripe.StripeError as e: - logger.error(f"Failed to create subscription checkout session: {e}") - raise StripeServiceException( - internal_message=f"Stripe checkout session creation failed: {e}" - ) - - async def create_checkout_session_for_credits_package( - self, - db: AsyncSession, - user_id: str, - price_id: str, - success_url: str, - cancel_url: str, - quantity: int, - email: Optional[str] = None, - ) -> str: - """Create a Stripe Checkout session for a credits package.""" - try: - # Validate that the selected price configuration exists. - config = await self.price_config_service.get_price_config(db, price_id) - if not config.is_credits_package(): - raise ValidationException( - user_message="Invalid price configuration", - violations=[ - { - "field": "price_id", - "description": f"Price ID {price_id} is not a credits package", - } - ], - ) - - # Ensure user is initialized (UserBalance exists) - await self.credits_service.ensure_user_initialized(db, user_id) - - user_balance = await self.credits_repo.get_user_balance(db, user_id) - if not user_balance: - # Should not happen after ensure_user_initialized - raise ValidationException( - user_message="Failed to initialize user balance", - violations=[ - { - "field": "user_id", - "description": f"Failed to initialize user balance for {user_id}", - } - ], - ) - - customer_id = user_balance.stripe_customer_id - - if not customer_id: - # Reuse an existing Stripe customer when the email already exists. - if email: - existing_customers = stripe.Customer.list(email=email, limit=1) - if existing_customers.data: - customer_id = existing_customers.data[0].id - - if not customer_id: - # Create a new Stripe customer when no existing record matches. - if not email: - # For new customers, we prefer having an email. - # If no email provided, we can't create a good customer record. - # But technically Stripe allows it. - # Better: Require email for new billing profiles. - raise ValidationException( - user_message="Email required for first-time payment", - violations=[ - { - "field": "email", - "description": "Email is required to create a billing profile", - } - ], - ) - - customer_params = { - "email": email, - "metadata": {"user_id": str(user_id)}, - } - # Username is not available without User model, omit it. - - customer = stripe.Customer.create(**customer_params) - customer_id = customer.id - - user_balance.stripe_customer_id = customer_id - - # Keep metadata values as strings so refunds can recover the user ID later. - metadata = { - "user_id": str(user_id), - "price_id": str(price_id), - "type": "credits_package", - "credits_amount": ( - str(config.credits_amount) if config.credits_amount else None - ), - "quantity": str(quantity), - } - - session_params: Dict[str, Any] = { - "customer": customer_id, - "customer_update": {"address": "auto"}, - "client_reference_id": str(user_id), - "line_items": [ - { - "price": price_id, - "quantity": quantity, - } - ], - "mode": "payment", # One-time payment. - "success_url": success_url, - "cancel_url": cancel_url, - "metadata": metadata, - # Copy metadata onto the PaymentIntent and Charge for refund handling. - "payment_intent_data": { - "metadata": metadata, - }, - # Collect more customer information for later reconciliation. - "allow_promotion_codes": True, - # Disable Adaptive Pricing to prevent currency switcher from hiding Alipay. - # Alipay handles USD→CNY conversion internally for customers. - "adaptive_pricing": {"enabled": False}, - # Require a billing address so Checkout syncs it to the customer record. - "billing_address_collection": "required", - } - - session = stripe.checkout.Session.create(**session_params) - - await db.commit() - - return str(session.url or "") - except stripe.StripeError as e: - logger.error(f"Stripe credits checkout session failed: {e}") - raise StripeServiceException( - internal_message=f"Stripe credits checkout session failed: {e}" - ) - - async def create_payment_intent( - self, user_id: str, amount: int, credits_amount: int, currency: str = "usd" - ) -> Dict[str, Any]: - """Create a PaymentIntent for a credits purchase.""" - try: - intent = stripe.PaymentIntent.create( - amount=amount, # amount in cents - currency=currency, - automatic_payment_methods={"enabled": True}, - metadata={ - "user_id": user_id, - "type": "credits", - "credits_amount": str(credits_amount), - }, - ) - return { - "client_secret": intent.client_secret, - "payment_intent_id": intent.id, - } - except stripe.StripeError as e: - logger.error(f"Failed to create payment intent: {e}") - raise StripeServiceException( - internal_message=f"Stripe payment intent creation failed: {e}" - ) - - async def handle_webhook( - self, db: AsyncSession, payload: bytes, sig_header: str - ) -> Dict[str, Any]: - """Handle a Stripe webhook payload.""" - try: - event = stripe.Webhook.construct_event( - payload, sig_header, settings.STRIPE_WEBHOOK_SECRET - ) - return await self._process_webhook_event(db, event) - except ValueError as e: - logger.error(f"Invalid payload: {e}") - raise ValidationException( - user_message="Invalid webhook payload", - violations=[ - {"field": "payload", "description": "Webhook payload is malformed"} - ], - ) - except stripe.SignatureVerificationError as e: - logger.error(f"Invalid signature: {e}") - raise AuthException( - user_message="Invalid webhook signature", - internal_message=f"Webhook signature verification failed: {e}", - ) - - async def _process_webhook_event( - self, db: AsyncSession, event: Dict[str, Any] - ) -> Dict[str, Any]: - """Dispatch an incoming Stripe webhook event.""" - event_type = event["type"] - - if event_type == "checkout.session.completed": - return await self._handle_checkout_completed(db, event) - elif event_type == "payment_intent.succeeded": - return await self._handle_payment_intent_succeeded(db, event) - elif event_type == "invoice.payment_succeeded": - return await self._handle_payment_succeeded(db, event) - elif event_type == "customer.subscription.deleted": - return await self._handle_subscription_deleted(db, event) - elif event_type == "charge.refunded": - return await self._handle_charge_refunded(db, event) - else: - return {"status": "ignored", "event_type": event_type} - - async def _handle_checkout_completed( - self, db: AsyncSession, event: Dict[str, Any] - ) -> Dict[str, Any]: - """Handle a completed Checkout session.""" - session = event["data"]["object"] - session_id = session["id"] - mode = session.get("mode") - metadata = session.get("metadata", {}) - user_id = metadata.get("user_id") - payment_type = metadata.get("type") - quantity = int(metadata.get("quantity", 1)) - - if not user_id: - logger.warning( - f"Checkout session {session_id} is missing user_id metadata; likely a test event, skipping" - ) - return { - "status": "ignored", - "message": "Missing user_id metadata (likely test event)", - "checkout_session_id": session_id, - "event_type": "checkout.session.completed", - } - - # Skip work that was already processed for this Checkout session. - if await self.payment_record_repo.is_processed( - db, checkout_session_id=session_id - ): - logger.info(f"Checkout session {session_id} already processed, skipping") - return { - "status": "ignored", - "message": "Already processed", - "checkout_session_id": session_id, - } - - # Seed audit metadata for the payment record. - payment_metadata = { - "session_id": session_id, - "stripe_session": session, # Full session payload for debugging and audits. - } - - # Create the pending payment record before side effects run. - payment_record = PaymentRecord( - checkout_session_id=session_id, - user_id=user_id, - payment_type=payment_type or "unknown", - amount_cents=session.get("amount_total", 0), - currency=session.get("currency", "cny").upper(), - status="pending", - extra_metadata=payment_metadata, - ) - db.add(payment_record) - await db.flush() # Get the database ID without committing yet. - - try: - if mode == "payment" and payment_type == "credits_package": - # Credits package purchase flow. - price_id = metadata.get("price_id") - - if not price_id: - logger.error(f"Incomplete Credits pack info: price_id={price_id}") - return {"status": "error", "message": "Missing price_id"} - - # Load the credits amount and product metadata from the price config. - price_config = await self.price_config_service.get_price_config( - db, price_id - ) - credits_amount = price_config.credits_amount * quantity - if credits_amount is None: - logger.error( - f"Credits amount is not configured for price ID {price_id}" - ) - return { - "status": "error", - "message": "Credits amount not configured", - } - - # Attach purchased product details to the payment record. - product_description = f"Credits pack - {credits_amount} Credits" - if price_config.extra_metadata and price_config.extra_metadata.get( - "description" - ): - product_description = price_config.extra_metadata.get("description") - - payment_record.extra_metadata = { - **payment_metadata, - "product_description": product_description, - "price_id": price_id, - "credits_amount": credits_amount, - "product_metadata": price_config.extra_metadata - or {}, # Product metadata from the price config. - } - - # Grant the purchased credits to the user balance. - await self.credits_service.add_credits( - session=db, - user_id=user_id, - amount=credits_amount, - reason=f"Purchase credits pack: {product_description}", - stripe_payment_id=session.get("payment_intent"), - ) - - # Mark the payment record as completed. - payment_record.status = "succeeded" - payment_record.credits_amount = credits_amount - payment_record.processed_at = utc_now_naive() - - await TierService.refresh_tier(user_id, db) - await db.commit() - await db.refresh(payment_record) - - logger.info( - f"Credits pack purchase succeeded: user_id={user_id}, credits={credits_amount}, price_id={price_id}" - ) - return { - "status": "success", - "event_type": "checkout.session.completed", - "user_id": user_id, - "credits_amount": credits_amount, - "payment_type": "credits_package", - } - else: - logger.warning( - f"Unknown payment type: mode={mode}, type={payment_type}" - ) - return {"status": "ignored", "message": "Unknown payment type"} - - except KnowhereException: - raise - except Exception as e: - logger.error( - f"Failed to process checkout.session.completed: {e}", exc_info=True - ) - payment_record.status = "failed" - payment_record.extra_metadata = { - **(payment_record.extra_metadata or {}), - "error": str(e), - } - await db.commit() - raise StripeServiceException( - internal_message=f"Failed to process checkout.session.completed: {str(e)}", - original_exception=e, - ) - - async def _handle_payment_intent_succeeded( - self, db: AsyncSession, event: Dict[str, Any] - ) -> Dict[str, Any]: - """Handle a successful PaymentIntent for a credits purchase.""" - payment_intent = event["data"]["object"] - payment_intent_id = payment_intent["id"] - metadata = payment_intent.get("metadata", {}) - user_id = metadata.get("user_id") - payment_type = metadata.get("type") - - if payment_type != "credits": - logger.info( - f"PaymentIntent {payment_intent_id} is not a Credits payment, skipping" - ) - return {"status": "ignored", "payment_intent_id": payment_intent_id} - - if not user_id: - logger.warning( - f"PaymentIntent {payment_intent_id} is missing user_id metadata; likely a test event, skipping" - ) - return { - "status": "ignored", - "message": "Missing user_id metadata (likely test event)", - "payment_intent_id": payment_intent_id, - } - - # Skip work that was already processed for this PaymentIntent. - if await self.payment_record_repo.is_processed( - db, payment_intent_id=payment_intent_id - ): - logger.info( - f"PaymentIntent {payment_intent_id} already processed, skipping" - ) - return { - "status": "ignored", - "message": "Already processed", - "payment_intent_id": payment_intent_id, - } - - # Seed audit metadata for the payment record. - payment_metadata = { - "payment_intent_id": payment_intent_id, - "stripe_payment_intent": payment_intent, # Full PaymentIntent payload for debugging and audits. - } - - # Create the pending payment record before side effects run. - payment_record = PaymentRecord( - payment_intent_id=payment_intent_id, - user_id=user_id, - payment_type="credits_package", - amount_cents=payment_intent.get("amount", 0), - currency=payment_intent.get("currency", "cny").upper(), - status="pending", - extra_metadata=payment_metadata, - ) - db.add(payment_record) - await db.flush() # Get the database ID without committing yet. - - try: - # Read the purchased credits amount from metadata. - credits_amount_str = metadata.get("credits_amount") - if not credits_amount_str: - logger.error( - f"PaymentIntent {payment_intent_id} is missing credits_amount" - ) - payment_record.status = "failed" - payment_record.extra_metadata = { - **(payment_record.extra_metadata or {}), - "error": "Missing credits_amount", - } - await db.commit() - return {"status": "error", "message": "Missing credits_amount"} - - credits_amount = int(credits_amount_str) - - # Attach purchased product details to the payment record. - payment_record.extra_metadata = { - **payment_metadata, - "product_description": f"Credits package - {credits_amount} Credits", - "credits_amount": credits_amount, - "payment_method": "payment_intent", # Marks this purchase as PaymentIntent-based. - } - - # Amount validation can be layered in here if needed later. - payment_intent.get("amount", 0) - - # Grant the purchased credits to the user balance. - await self.credits_service.add_credits( - session=db, - user_id=user_id, - amount=credits_amount, - reason=f"buy credits - {credits_amount} Credits", - stripe_payment_id=payment_intent_id, - ) - - # Mark the payment record as completed. - payment_record.status = "succeeded" - payment_record.credits_amount = credits_amount - payment_record.processed_at = utc_now_naive() - - await TierService.refresh_tier(user_id, db) - await db.commit() - await db.refresh(payment_record) - - logger.info( - f"buy credits success: user_id={user_id}, credits={credits_amount}, payment_intent_id={payment_intent_id}" - ) - return { - "status": "success", - "event_type": "payment_intent.succeeded", - "user_id": user_id, - "credits_amount": credits_amount, - "payment_type": "credits_package", - } - - except Exception as e: - logger.error(f"Failed to process Credits purchase: {e}", exc_info=True) - payment_record.status = "failed" - payment_record.extra_metadata = { - **(payment_record.extra_metadata or {}), - "error": str(e), - } - await db.commit() - raise StripeServiceException( - internal_message=f"Failed to process Credits purchase: {str(e)}", - original_exception=e, - ) - - async def _handle_payment_succeeded( - self, db: AsyncSession, event: Dict[str, Any] - ) -> Dict[str, Any]: - """Handle a successful subscription renewal payment.""" - invoice = event["data"]["object"] - subscription_id = invoice.get("subscription") - - if not subscription_id: - logger.warning("Invoice is missing subscription ID") - return {"status": "ignored", "message": "Missing subscription_id"} - - return {"status": "ignored", "message": "Subscription renewal not implemented"} - - async def _handle_subscription_deleted( - self, db: AsyncSession, event: Dict[str, Any] - ) -> Dict[str, Any]: - """Handle subscription deletion events.""" - subscription = event["data"]["object"] - stripe_subscription_id = subscription["id"] - - try: - # Subscription management not yet implemented - logger.warning( - f"Local subscription record not found: stripe_subscription_id={stripe_subscription_id}" - ) - - return {"status": "success", "subscription_id": stripe_subscription_id} - except KnowhereException: - raise - except Exception as e: - logger.error( - f"Failed to process customer.subscription.deleted: {e}", exc_info=True - ) - raise StripeServiceException( - internal_message=f"Failed to process customer.subscription.deleted: {str(e)}", - original_exception=e, - ) - - async def _handle_charge_refunded( - self, db: AsyncSession, event: Dict[str, Any] - ) -> Dict[str, Any]: - """Handle refund events, including manual refunds from the Stripe dashboard.""" - charge = event["data"]["object"] - charge_id = charge.get("id") - refund_items = (charge.get("refunds", {}) or {}).get("data", []) or [] - latest_refund = refund_items[-1] if refund_items else None - - payment_intent_id = charge.get("payment_intent") - refund_id = latest_refund.get("id") if latest_refund else None - - currency = (charge.get("currency") or "cny").upper() - - # Use a stable idempotency key derived from the refund or charge identifier. - idempotency_key = refund_id or f"{charge_id}-refund" - - # Recover the original payment record to reuse billing context such as user_id. - original_record = None - if payment_intent_id: - original_record = await self.payment_record_repo.get_by_payment_intent_id( - db, payment_intent_id - ) - - metadata = charge.get("metadata") or {} - user_id = metadata.get("user_id") or (getattr(original_record, "user_id", None)) - payment_type = ( - metadata.get("type") - or (getattr(original_record, "payment_type", None)) - or "refund" - ) - - if not user_id: - logger.error( - f"Refund event is missing user_id; cannot record refund: charge_id={charge_id}" - ) - return { - "status": "error", - "message": "Missing user_id for refund", - "event_type": "charge.refunded", - } - - # Normalize user_id to UUID before using it in SQL filters. - if user_id and isinstance(user_id, str): - try: - user_id = UUID(user_id) - except ValueError: - logger.error(f"Invalid user_id format: {user_id}") - return { - "status": "error", - "message": "Invalid user_id format", - "event_type": "charge.refunded", - } - - user_id_str = str(user_id) - - # Compute the incremental refund amount from the cumulative Stripe total. - # total_refund_amount_cents already includes the current refund event. - total_refund_amount_cents = charge.get("amount_refunded") or 0 - - # Load previously recorded refund totals for the same payment flow. - origin_total_refund_amount_cents = 0 - - # Sum historical refund records that use the same synthetic refund key. - query = ( - select(func.sum(PaymentRecord.amount_cents)) - .where(PaymentRecord.payment_intent_id == idempotency_key) - .where(PaymentRecord.user_id == user_id) - .where( - PaymentRecord.amount_cents - < 0 # Refund rows are stored as negative amounts. - ) - ) - result = await db.execute(query) - # Sum negative refund amounts and convert back to a positive total. - origin_total_refund_amount_cents = abs(result.scalar() or 0) - - refund_amount_cents = ( - total_refund_amount_cents - origin_total_refund_amount_cents - ) - if refund_amount_cents <= 0: - # The refund has already been processed; keep this path idempotent. - logger.info( - f"Refund already processed, skipping: charge_id={charge_id}, refund_id={refund_id}" - ) - return { - "status": "success", - "event_type": "charge.refunded", - "message": "Already processed", - "user_id": user_id, - "refund_id": refund_id, - } - - # Translate the refunded cash amount back into credits using price metadata. - credits_refunded = None - price_id = metadata.get("price_id") or ( - getattr(original_record, "extra_metadata", {}) or {} - ).get("price_id") - if price_id: - try: - price_cfg = await self.price_config_service.get_price_config( - db, price_id - ) - if price_cfg and price_cfg.amount_cents: - credits_refunded = -int( - price_cfg.credits_amount - * abs(refund_amount_cents) - / abs(price_cfg.amount_cents) # credits_amount * quantity - ) - except Exception as e: - logger.warning( - f"Failed to calculate refunded Credits, price_id={price_id}: {e}" - ) - credits_refunded = None - - # Fall back to the original payment record ratio when price metadata is unavailable. - if ( - credits_refunded is None - and original_record - and original_record.credits_amount - and original_record.amount_cents - ): - credits_refunded = -int( - abs(original_record.credits_amount) - * abs(refund_amount_cents) - / abs(original_record.amount_cents) - ) - - # Apply the credit adjustment to the user balance when needed. - if credits_refunded is not None and credits_refunded < 0: - await self.credits_service.add_credits( - session=db, - user_id=user_id_str, - amount=credits_refunded, - reason="Refund adjustment", - transaction_type="refund", - transaction_metadata={"refund_id": refund_id, "charge_id": charge_id}, - ) - - refund_metadata = { - "refund_id": refund_id, - "charge_id": charge_id, - "original_payment_intent_id": payment_intent_id, - "original_payment_record_id": getattr(original_record, "id", None), - "reason": (latest_refund or {}).get("reason"), - "balance_transaction": (latest_refund or {}).get("balance_transaction"), - } - - refund_record = PaymentRecord( - payment_intent_id=idempotency_key, - user_id=user_id, - payment_type=payment_type, - amount_cents=-abs(refund_amount_cents), - currency=currency, - status="succeeded", - credits_amount=credits_refunded, - plan_id=getattr(original_record, "plan_id", None), - stripe_subscription_id=getattr( - original_record, "stripe_subscription_id", None - ), - processed_at=utc_now_naive(), - extra_metadata=refund_metadata, - ) - - db.add(refund_record) - await db.commit() - await db.refresh(refund_record) - - logger.info( - f"Refund record created: user_id={user_id}, amount_cents={refund_record.amount_cents}, " - f"refund_id={refund_id}, charge_id={charge_id}" - ) - - return { - "status": "success", - "event_type": "charge.refunded", - "user_id": user_id, - "refund_amount_cents": abs(refund_amount_cents), - "payment_intent_id": payment_intent_id, - "refund_id": refund_id, - } diff --git a/apps/api/app/services/billing/stripe_webhook_service.py b/apps/api/app/services/billing/stripe_webhook_service.py new file mode 100644 index 000000000..41581d287 --- /dev/null +++ b/apps/api/app/services/billing/stripe_webhook_service.py @@ -0,0 +1,196 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from typing import Any, TypeAlias + +import stripe +from app.repositories.payment_record_repository import PaymentRecordRepository +from app.services.billing.stripe_credits_settlement_service import ( + StripeCreditsSettlementService, +) +from app.services.billing.price_config_service import PriceConfigService +from app.services.billing.stripe_refund_reconciliation_service import ( + StripeRefundReconciliationService, +) +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ( + AuthException, + KnowhereException, + StripeServiceException, + SystemSettingMissingException, + ValidationException, +) +from shared.core.logging import logger +from shared.services.billing import CreditsService + +StripeEvent: TypeAlias = dict[str, Any] +StripeWebhookHandler: TypeAlias = Callable[ + [AsyncSession, StripeEvent], Awaitable[dict[str, object]] +] + + +class StripeWebhookService: + def __init__( + self, + *, + payment_record_repository: PaymentRecordRepository | None = None, + price_config_service: PriceConfigService | None = None, + credits_service: CreditsService | None = None, + credits_settlement_service: StripeCreditsSettlementService | None = None, + refund_reconciliation_service: StripeRefundReconciliationService | None = None, + ) -> None: + self._configure_stripe_api() + self._payment_record_repository = ( + payment_record_repository or PaymentRecordRepository() + ) + self._price_config_service = price_config_service or PriceConfigService() + self._credits_service = credits_service or CreditsService() + self._credits_settlement_service = ( + credits_settlement_service + or StripeCreditsSettlementService( + payment_record_repository=self._payment_record_repository, + price_config_service=self._price_config_service, + credits_service=self._credits_service, + ) + ) + self._refund_reconciliation_service = ( + refund_reconciliation_service + or StripeRefundReconciliationService( + payment_record_repository=self._payment_record_repository, + price_config_service=self._price_config_service, + credits_service=self._credits_service, + ) + ) + + async def handle_webhook( + self, + db: AsyncSession, + *, + payload: bytes, + sig_header: str, + ) -> dict[str, object]: + try: + event = stripe.Webhook.construct_event( + payload, + sig_header, + settings.STRIPE_WEBHOOK_SECRET, + ) + return await self._dispatch_event(db, event) + except ValueError as exc: + logger.error(f"Invalid payload: {exc}") + raise ValidationException( + user_message="Invalid webhook payload", + violations=[ + {"field": "payload", "description": "Webhook payload is malformed"} + ], + ) + except stripe.SignatureVerificationError as exc: + logger.error(f"Invalid signature: {exc}") + raise AuthException( + user_message="Invalid webhook signature", + internal_message=f"Webhook signature verification failed: {exc}", + ) + + async def _dispatch_event( + self, + db: AsyncSession, + event: StripeEvent, + ) -> dict[str, object]: + event_type = str(event["type"]) + handler = self._event_handlers().get(event_type) + if handler is None: + return {"status": "ignored", "event_type": event_type} + return await handler(db, event) + + def _event_handlers(self) -> dict[str, StripeWebhookHandler]: + return { + "checkout.session.completed": self._handle_checkout_completed, + "payment_intent.succeeded": self._handle_payment_intent_succeeded, + "invoice.payment_succeeded": self._handle_payment_succeeded, + "customer.subscription.deleted": self._handle_subscription_deleted, + "charge.refunded": self._handle_charge_refunded, + } + + async def _handle_checkout_completed( + self, + db: AsyncSession, + event: StripeEvent, + ) -> dict[str, object]: + return await self._credits_settlement_service.handle_checkout_completed( + db, + event=event, + ) + + async def _handle_payment_intent_succeeded( + self, + db: AsyncSession, + event: StripeEvent, + ) -> dict[str, object]: + return await self._credits_settlement_service.handle_payment_intent_succeeded( + db, + event=event, + ) + + async def _handle_payment_succeeded( + self, + db: AsyncSession, + event: StripeEvent, + ) -> dict[str, object]: + del db + invoice = event["data"]["object"] + subscription_id = invoice.get("subscription") + + if not subscription_id: + logger.warning("Invoice is missing subscription ID") + return {"status": "ignored", "message": "Missing subscription_id"} + + return {"status": "ignored", "message": "Subscription renewal not implemented"} + + async def _handle_subscription_deleted( + self, + db: AsyncSession, + event: StripeEvent, + ) -> dict[str, object]: + del db + subscription = event["data"]["object"] + stripe_subscription_id = str(subscription["id"]) + + try: + logger.warning( + "Local subscription record not found: " + f"stripe_subscription_id={stripe_subscription_id}" + ) + return {"status": "success", "subscription_id": stripe_subscription_id} + except KnowhereException: + raise + except Exception as exc: + logger.error( + f"Failed to process customer.subscription.deleted: {exc}", + exc_info=True, + ) + raise StripeServiceException( + internal_message=( + "Failed to process customer.subscription.deleted: " + f"{str(exc)}" + ), + original_exception=exc, + ) + + async def _handle_charge_refunded( + self, + db: AsyncSession, + event: StripeEvent, + ) -> dict[str, object]: + return await self._refund_reconciliation_service.reconcile_charge_refund( + db, + event=event, + ) + + def _configure_stripe_api(self) -> None: + if not settings.STRIPE_SECRET_KEY: + raise SystemSettingMissingException( + internal_message="Stripe API key not configured (STRIPE_SECRET_KEY)" + ) + stripe.api_key = settings.STRIPE_SECRET_KEY diff --git a/apps/api/app/services/demo/__init__.py b/apps/api/app/services/demo/__init__.py new file mode 100644 index 000000000..a25f9979d --- /dev/null +++ b/apps/api/app/services/demo/__init__.py @@ -0,0 +1 @@ +"""Demo source catalog and materialization services.""" diff --git a/apps/api/app/services/demo/document_service.py b/apps/api/app/services/demo/document_service.py new file mode 100644 index 000000000..f0a63bd1c --- /dev/null +++ b/apps/api/app/services/demo/document_service.py @@ -0,0 +1,87 @@ +"""Demo Source catalog facade.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from app.services.demo.source_catalog import DemoSourceCatalog +from app.services.demo.source_materializer import ( + DemoSourceMaterializer, + MaterializedDemoSource, +) +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.publication_service import RetrievalPublicationService + + +class DemoDocumentService: + """Serves canonical demo data and delegates user-copy side effects.""" + + def __init__( + self, + *, + catalog: DemoSourceCatalog | None = None, + publication_service: RetrievalPublicationService | None = None, + ) -> None: + self._catalog = catalog or DemoSourceCatalog() + self._materializer = DemoSourceMaterializer( + catalog=self._catalog, + publication_service=publication_service, + ) + + def get_catalog(self) -> dict[str, Any]: + return self._catalog.get_catalog() + + def list_chunks( + self, + *, + demo_source_id: str, + page: int, + page_size: int, + ) -> dict[str, Any] | None: + return self._catalog.list_chunks( + demo_source_id=demo_source_id, + page=page, + page_size=page_size, + ) + + def get_chunk( + self, + *, + demo_source_id: str, + demo_chunk_id: str, + ) -> dict[str, Any] | None: + return self._catalog.get_chunk( + demo_source_id=demo_source_id, + demo_chunk_id=demo_chunk_id, + ) + + def get_original_file_path(self, *, demo_source_id: str) -> Path | None: + return self._catalog.get_original_file_path(demo_source_id=demo_source_id) + + def get_asset_file_path( + self, + *, + demo_source_id: str, + asset_path: str, + ) -> Path | None: + return self._catalog.get_asset_file_path( + demo_source_id=demo_source_id, + asset_path=asset_path, + ) + + async def materialize_sources( + self, + db: AsyncSession, + *, + user_id: str, + namespace: str, + demo_source_ids: list[str], + ) -> list[MaterializedDemoSource]: + return await self._materializer.materialize_sources( + db, + user_id=user_id, + namespace=namespace, + demo_source_ids=demo_source_ids, + ) diff --git a/apps/api/app/services/demo/source_catalog.py b/apps/api/app/services/demo/source_catalog.py new file mode 100644 index 000000000..b3bdb71bb --- /dev/null +++ b/apps/api/app/services/demo/source_catalog.py @@ -0,0 +1,293 @@ +"""Canonical Demo Source catalog and file access.""" + +from __future__ import annotations + +import json +import math +from dataclasses import dataclass +from functools import lru_cache +from pathlib import Path +from typing import Any + +from app.services.demo.source_projection import DemoSourceProjection + + +@dataclass(frozen=True) +class DemoCitationDefinition: + section_path: str + description: str + content: str + + +@dataclass(frozen=True) +class DemoExampleDefinition: + id: str + question: str + answer: str + citations: tuple[DemoCitationDefinition, ...] + + +@dataclass(frozen=True) +class DemoSourceDefinition: + demo_source_id: str + canonical_document_id: str + title: str + mime_type: str + size_bytes: int + asset_directory: str + chunk_count: int + examples: tuple[DemoExampleDefinition, ...] + + +_DATA_ROOT = Path(__file__).resolve().parents[2] / "data" / "demo_documents" +_ASSET_DIRECTORY_NAMES = frozenset({"images", "tables"}) +_DEMO_SOURCE_DEFINITIONS: tuple[DemoSourceDefinition, ...] = ( + DemoSourceDefinition( + demo_source_id="demo-tsla-q4-2025", + canonical_document_id="demo-doc-tsla-q4-2025", + title="TSLA-Q4-2025-Update.pdf", + mime_type="application/pdf", + size_bytes=5_648_867, + asset_directory="tsla-q4-2025", + chunk_count=70, + examples=( + DemoExampleDefinition( + id="demo-tsla-q4-2025-xai", + question="What does the document say about Tesla's xAI investment?", + answer=( + "Tesla entered an agreement on January 16, 2026 to invest " + "approximately $2 billion in xAI Series E Preferred Stock.\n\n" + "The document also says Tesla and xAI entered a framework " + "agreement to evaluate AI collaboration, with the investment " + "expected to close in Q1 2026 subject to customary regulatory " + "conditions." + ), + citations=( + DemoCitationDefinition( + section_path=( + "TSLA-Q4-2025-Update.pdf-->OTHER UPDATES" + ), + description="xAI investment", + content=( + "On January 16, 2026, Tesla entered into an agreement " + "to invest approximately" + ), + ), + ), + ), + DemoExampleDefinition( + id="demo-tsla-q4-2025-energy-storage", + question="What does the document say about energy storage?", + answer=( + "Tesla achieved its highest quarterly energy storage " + "deployments, driven by record Megapack deployments.\n\n" + "Energy gross profit reached a record $1.1 billion, marking " + "the fifth consecutive record quarter.\n\n" + "Tesla also plans to begin Megapack 3 and Megablock " + "production at Megafactory Houston in 2026." + ), + citations=( + DemoCitationDefinition( + section_path=( + "TSLA-Q4-2025-Update.pdf-->SUMMARY-->" + "Energy generation and storage" + ), + description="Storage deployment growth", + content=( + "We achieved our highest quarterly energy storage " + "deployments, driven by record Megapack deployments." + ), + ), + ), + ), + DemoExampleDefinition( + id="demo-tsla-q4-2025-production-plans", + question="What production plans does Tesla mention for 2026?", + answer=( + "Tesla says Cybercab, Tesla Semi, and Megapack 3 are on " + "schedule for volume production starting in 2026.\n\n" + "The same product update also notes that first-generation " + "Optimus production lines are being installed before volume " + "production." + ), + citations=( + DemoCitationDefinition( + section_path=( + "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->" + "Product" + ), + description="2026 production plans", + content=( + "Cybercab, Tesla Semi and Megapack 3 are on schedule " + "for volume production starting in 2026." + ), + ), + ), + ), + ), + ), +) + + +class DemoSourceCatalog: + def __init__(self, *, projection: DemoSourceProjection | None = None) -> None: + self._projection = projection or DemoSourceProjection() + + def list_sources(self) -> tuple[DemoSourceDefinition, ...]: + return _DEMO_SOURCE_DEFINITIONS + + def get_catalog(self) -> dict[str, Any]: + return { + "sources": [ + self._projection.source_catalog_payload( + source=source, + chunks=_load_source_chunks(source), + ) + for source in _DEMO_SOURCE_DEFINITIONS + ], + } + + def list_chunks( + self, + *, + demo_source_id: str, + page: int, + page_size: int, + ) -> dict[str, Any] | None: + source = self.get_source(demo_source_id) + if source is None: + return None + + chunks = _load_source_chunks(source) + start = (page - 1) * page_size + page_chunks = chunks[start : start + page_size] + return { + "demo_source_id": source.demo_source_id, + "canonical_document_id": source.canonical_document_id, + "title": source.title, + "mime_type": source.mime_type, + "chunks": [ + self._projection.chunk_payload( + source=source, + chunk=chunk, + sort_order=start + index, + ) + for index, chunk in enumerate(page_chunks) + ], + "pagination": { + "page": page, + "page_size": page_size, + "total": len(chunks), + "total_pages": math.ceil(len(chunks) / page_size) if chunks else 0, + }, + } + + def get_chunk( + self, + *, + demo_source_id: str, + demo_chunk_id: str, + ) -> dict[str, Any] | None: + source = self.get_source(demo_source_id) + if source is None: + return None + + chunks = _load_source_chunks(source) + for sort_order, chunk in enumerate(chunks): + if self._projection.matches_chunk_id( + source=source, + chunk=chunk, + demo_chunk_id=demo_chunk_id, + ): + return { + "demo_source_id": source.demo_source_id, + "canonical_document_id": source.canonical_document_id, + "chunk": self._projection.chunk_payload( + source=source, + chunk=chunk, + sort_order=sort_order, + ), + } + + return None + + def get_original_file_path(self, *, demo_source_id: str) -> Path | None: + source = self.get_source(demo_source_id) + if source is None: + return None + + file_path = self.source_directory(source) / "original.pdf" + return file_path if file_path.is_file() else None + + def get_asset_file_path( + self, + *, + demo_source_id: str, + asset_path: str, + ) -> Path | None: + source = self.get_source(demo_source_id) + if source is None: + return None + + source_directory = self.source_directory(source).resolve() + normalized_asset_path = _normalize_asset_path(asset_path) + if normalized_asset_path is None: + return None + + candidate = (source_directory / normalized_asset_path).resolve() + if not candidate.is_relative_to(source_directory): + return None + + return candidate if candidate.is_file() else None + + def require_source(self, demo_source_id: str) -> DemoSourceDefinition: + source = self.get_source(demo_source_id) + if source is None: + raise KeyError(demo_source_id) + return source + + def get_source(self, demo_source_id: str) -> DemoSourceDefinition | None: + return next( + ( + source + for source in _DEMO_SOURCE_DEFINITIONS + if source.demo_source_id == demo_source_id + ), + None, + ) + + def source_directory(self, source: DemoSourceDefinition) -> Path: + return _DATA_ROOT / source.asset_directory + + def publication_chunks(self, source: DemoSourceDefinition) -> list[dict[str, Any]]: + return self._projection.publication_chunks( + source=source, + chunks=_load_source_chunks(source), + ) + + +def _normalize_asset_path(asset_path: str) -> Path | None: + normalized = str(asset_path or "").strip().replace("\\", "/").lstrip("/") + parts = [part for part in normalized.split("/") if part and part != "."] + if not parts or parts[0] not in _ASSET_DIRECTORY_NAMES: + return None + if any(part == ".." or part.startswith(".") for part in parts): + return None + return Path(*parts) + + +@lru_cache(maxsize=8) +def _load_source_chunks(source: DemoSourceDefinition) -> tuple[dict[str, Any], ...]: + chunks_path = (_DATA_ROOT / source.asset_directory) / "chunks.json" + with chunks_path.open("r", encoding="utf-8") as file: + payload = json.load(file) + + chunks = payload.get("chunks") if isinstance(payload, dict) else None + if not isinstance(chunks, list): + return () + + return tuple( + dict(chunk) + for chunk in chunks + if isinstance(chunk, dict) and isinstance(chunk.get("chunk_id"), str) + ) diff --git a/apps/api/app/services/demo/source_materializer.py b/apps/api/app/services/demo/source_materializer.py new file mode 100644 index 000000000..df9e2f0c8 --- /dev/null +++ b/apps/api/app/services/demo/source_materializer.py @@ -0,0 +1,326 @@ +"""Demo Source Materialization workflow.""" + +from __future__ import annotations + +import shutil +import tempfile +from dataclasses import dataclass +from datetime import datetime, timezone +from hashlib import blake2b +from pathlib import Path +from uuid import uuid4 + +from app.services.demo.source_catalog import DemoSourceCatalog, DemoSourceDefinition +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ValidationException +from shared.models.database.demo_materialization import DemoMaterialization +from shared.models.database.document import Document +from shared.models.database.job import Job +from shared.models.database.job_result import JobResult +from shared.services.retrieval.cache_service import invalidate_retrieval_cache_namespaces +from shared.services.retrieval.publication_service import RetrievalPublicationService +from shared.services.storage.result_storage import get_result_storage + + +@dataclass(frozen=True) +class MaterializedDemoSource: + """User-owned copy of one canonical demo source.""" + + demo_source_id: str + document_id: str + status: str + title: str + mime_type: str + size_bytes: int + chunk_count: int + + +class DemoSourceMaterializer: + """Copies canonical Demo Sources into user-owned retrieval state.""" + + def __init__( + self, + *, + catalog: DemoSourceCatalog, + publication_service: RetrievalPublicationService | None = None, + ) -> None: + self._catalog = catalog + self._publication_service = publication_service or RetrievalPublicationService() + + async def materialize_sources( + self, + db: AsyncSession, + *, + user_id: str, + namespace: str, + demo_source_ids: list[str], + ) -> list[MaterializedDemoSource]: + selected_demo_source_ids = _deduplicate_source_ids(demo_source_ids) + if not selected_demo_source_ids: + raise ValidationException( + user_message="At least one demo source must be selected.", + violations=[ + { + "field": "demo_source_ids", + "description": "Select one or more demo source IDs.", + } + ], + ) + + selected_sources = [ + self._catalog.require_source(demo_source_id) + for demo_source_id in selected_demo_source_ids + ] + results: list[MaterializedDemoSource] = [] + for source in selected_sources: + result = await self._materialize_source( + db, + user_id=user_id, + namespace=namespace, + source=source, + ) + results.append(result) + + await db.commit() + await invalidate_retrieval_cache_namespaces( + user_id=user_id, + namespaces=[namespace], + ) + return results + + async def _materialize_source( + self, + db: AsyncSession, + *, + user_id: str, + namespace: str, + source: DemoSourceDefinition, + ) -> MaterializedDemoSource: + await _lock_materialization_scope( + db, + user_id=user_id, + namespace=namespace, + demo_source_id=source.demo_source_id, + ) + existing = await self._get_existing_materialization( + db, + user_id=user_id, + namespace=namespace, + demo_source_id=source.demo_source_id, + ) + if existing is not None and await self._is_active_document( + db, + document_id=existing.document_id, + ): + return _materialized_source_payload( + source=source, + document_id=existing.document_id, + status="existing", + ) + + document_id = f"doc_{uuid4().hex[:12]}" + job_id = f"job_demo_{uuid4().hex[:12]}" + job_result_id = str(uuid4()) + timestamp = _utc_now() + result_bundle = _upload_demo_result_bundle( + job_id=job_id, + source_directory=self._catalog.source_directory(source), + ) + + db.add( + Job( + job_id=job_id, + user_id=user_id, + job_type="demo_materialization", + status="done", + source_type="demo", + webhook_enabled=False, + job_metadata={ + "document_id": document_id, + "namespace": namespace, + "source_type": "demo", + "source_file_name": source.title, + "demo_source_id": source.demo_source_id, + }, + version=0, + created_at=timestamp, + updated_at=timestamp, + credits_charged=0, + billing_status="skipped", + ) + ) + db.add( + JobResult( + id=job_result_id, + job_id=job_id, + delivery_mode="inline", + document_metadata={ + "source_file_name": source.title, + "demo_source_id": source.demo_source_id, + }, + inline_payload={"source": "canonical_demo"}, + result_s3_key=result_bundle["zip_key"], + result_size=result_bundle["zip_size"], + created_at=timestamp, + updated_at=timestamp, + ) + ) + await db.flush() + chunks = self._catalog.publication_chunks(source) + await db.run_sync( + lambda sync_db: self._publication_service.publish_document_state( + sync_db, + job_id=job_id, + job_result_id=job_result_id, + chunks=[dict(chunk) for chunk in chunks], + ) + ) + await db.run_sync( + lambda sync_db: self._publication_service.publish_document_graph( + sync_db, + job_id=job_id, + job_result_id=job_result_id, + ) + ) + await db.flush() + + if existing is None: + db.add( + DemoMaterialization( + user_id=user_id, + namespace=namespace, + demo_source_id=source.demo_source_id, + document_id=document_id, + created_at=timestamp, + updated_at=timestamp, + ) + ) + else: + existing.document_id = document_id + existing.updated_at = timestamp + await db.flush() + return _materialized_source_payload( + source=source, + document_id=document_id, + status="created", + ) + + async def _get_existing_materialization( + self, + db: AsyncSession, + *, + user_id: str, + namespace: str, + demo_source_id: str, + ) -> DemoMaterialization | None: + result = await db.execute( + select(DemoMaterialization) + .where(DemoMaterialization.user_id == user_id) + .where(DemoMaterialization.namespace == namespace) + .where(DemoMaterialization.demo_source_id == demo_source_id) + .with_for_update() + .limit(1) + ) + return result.scalar_one_or_none() + + async def _is_active_document( + self, + db: AsyncSession, + *, + document_id: str, + ) -> bool: + result = await db.execute( + select(Document.document_id) + .where(Document.document_id == document_id) + .where(Document.status == "active") + .limit(1) + ) + return result.scalar_one_or_none() is not None + + +def _deduplicate_source_ids(demo_source_ids: list[str]) -> list[str]: + selected: list[str] = [] + seen: set[str] = set() + for demo_source_id in demo_source_ids: + normalized = str(demo_source_id).strip() + if not normalized or normalized in seen: + continue + selected.append(normalized) + seen.add(normalized) + return selected + + +async def _lock_materialization_scope( + db: AsyncSession, + *, + user_id: str, + namespace: str, + demo_source_id: str, +) -> None: + lock_id = _materialization_lock_id( + user_id=user_id, + namespace=namespace, + demo_source_id=demo_source_id, + ) + await db.execute(select(func.pg_advisory_xact_lock(lock_id))) + + +def _materialization_lock_id( + *, + user_id: str, + namespace: str, + demo_source_id: str, +) -> int: + lock_key = f"{user_id}\0{namespace}\0{demo_source_id}" + digest = blake2b(lock_key.encode("utf-8"), digest_size=8).digest() + return int.from_bytes(digest, byteorder="big", signed=True) + + +def _materialized_source_payload( + *, + source: DemoSourceDefinition, + document_id: str, + status: str, +) -> MaterializedDemoSource: + return MaterializedDemoSource( + demo_source_id=source.demo_source_id, + document_id=document_id, + status=status, + title=source.title, + mime_type=source.mime_type, + size_bytes=source.size_bytes, + chunk_count=source.chunk_count, + ) + + +def _upload_demo_result_bundle( + *, + job_id: str, + source_directory: Path, +) -> dict[str, int | str]: + with tempfile.TemporaryDirectory(prefix="knowhere-demo-result-") as temp_directory: + zip_base_path = Path(temp_directory) / job_id + zip_file_path = Path( + shutil.make_archive( + str(zip_base_path), + "zip", + root_dir=source_directory, + ) + ) + zip_size = zip_file_path.stat().st_size + bundle = get_result_storage().upload( + job_id=job_id, + result_dir=str(source_directory), + zip_file_path=str(zip_file_path), + ) + + return { + "zip_key": bundle.zip_key, + "zip_size": zip_size, + } + + +def _utc_now() -> datetime: + return datetime.now(timezone.utc).replace(tzinfo=None) diff --git a/apps/api/app/services/demo/source_projection.py b/apps/api/app/services/demo/source_projection.py new file mode 100644 index 000000000..46a5f8ffe --- /dev/null +++ b/apps/api/app/services/demo/source_projection.py @@ -0,0 +1,304 @@ +"""Projection logic for canonical Demo Source data.""" + +from __future__ import annotations + +from typing import Any, Protocol +from urllib.parse import quote + + +class _DemoCitationDefinition(Protocol): + @property + def section_path(self) -> str: + raise NotImplementedError + + @property + def description(self) -> str: + raise NotImplementedError + + @property + def content(self) -> str: + raise NotImplementedError + + +class _DemoExampleDefinition(Protocol): + @property + def id(self) -> str: + raise NotImplementedError + + @property + def question(self) -> str: + raise NotImplementedError + + @property + def answer(self) -> str: + raise NotImplementedError + + @property + def citations(self) -> tuple[_DemoCitationDefinition, ...]: + raise NotImplementedError + + +class _DemoSourceDefinition(Protocol): + @property + def demo_source_id(self) -> str: + raise NotImplementedError + + @property + def canonical_document_id(self) -> str: + raise NotImplementedError + + @property + def title(self) -> str: + raise NotImplementedError + + @property + def mime_type(self) -> str: + raise NotImplementedError + + @property + def size_bytes(self) -> int: + raise NotImplementedError + + @property + def chunk_count(self) -> int: + raise NotImplementedError + + @property + def examples(self) -> tuple[_DemoExampleDefinition, ...]: + raise NotImplementedError + + +class DemoSourceProjection: + def source_catalog_payload( + self, + *, + source: _DemoSourceDefinition, + chunks: tuple[dict[str, Any], ...], + ) -> dict[str, Any]: + return { + "demo_source_id": source.demo_source_id, + "canonical_document_id": source.canonical_document_id, + "title": source.title, + "mime_type": source.mime_type, + "size_bytes": source.size_bytes, + "status": "ready", + "chunk_count": source.chunk_count, + "original_file": { + "url": f"/api/v1/demo/sources/{source.demo_source_id}/original", + "mime_type": source.mime_type, + "size_bytes": source.size_bytes, + "can_download": False, + }, + "examples": [ + self._example_payload(source=source, example=example, chunks=chunks) + for example in source.examples + ], + } + + def chunk_payload( + self, + *, + source: _DemoSourceDefinition, + chunk: dict[str, Any], + sort_order: int, + ) -> dict[str, Any]: + metadata = _metadata(chunk) + file_path = _first_string( + metadata.get("file_path"), + metadata.get("filePath"), + chunk.get("file_path"), + chunk.get("path") if _is_media_chunk(chunk) else None, + ) + return { + "id": self.canonical_chunk_id(source=source, chunk=chunk), + "chunk_id": chunk["chunk_id"], + "chunk_type": _normalize_chunk_type(chunk.get("type")), + "content": str(chunk.get("content") or ""), + "section_path": str(chunk.get("path") or "") or None, + "source_chunk_path": str(chunk.get("path") or "") or None, + "file_path": file_path, + "sort_order": sort_order, + "metadata": metadata, + "asset_url": _asset_url(source=source, file_path=file_path), + "created_at": None, + } + + def publication_chunks( + self, + *, + source: _DemoSourceDefinition, + chunks: tuple[dict[str, Any], ...], + ) -> list[dict[str, Any]]: + return [ + _publication_chunk(source=source, chunk=chunk) + for chunk in chunks + ] + + def canonical_chunk_id( + self, + *, + source: _DemoSourceDefinition, + chunk: dict[str, Any], + ) -> str: + return f"{source.demo_source_id}:{chunk['chunk_id']}" + + def matches_chunk_id( + self, + *, + source: _DemoSourceDefinition, + chunk: dict[str, Any], + demo_chunk_id: str, + ) -> bool: + return demo_chunk_id in { + self.canonical_chunk_id(source=source, chunk=chunk), + chunk["chunk_id"], + } + + def _example_payload( + self, + *, + source: _DemoSourceDefinition, + example: _DemoExampleDefinition, + chunks: tuple[dict[str, Any], ...], + ) -> dict[str, Any]: + return { + "id": example.id, + "question": example.question, + "answer": example.answer, + "citations": [ + self._citation_payload(source=source, citation=citation, chunks=chunks) + for citation in example.citations + ], + } + + def _citation_payload( + self, + *, + source: _DemoSourceDefinition, + citation: _DemoCitationDefinition, + chunks: tuple[dict[str, Any], ...], + ) -> dict[str, Any]: + chunk = _resolve_citation_chunk(source=source, citation=citation, chunks=chunks) + return { + "demo_source_id": source.demo_source_id, + "canonical_document_id": source.canonical_document_id, + "canonical_chunk_id": self.canonical_chunk_id(source=source, chunk=chunk), + "chunk_id": chunk["chunk_id"], + "chunk_type": _normalize_chunk_type(chunk.get("type")), + "content": citation.content, + "description": citation.description, + "source": { + "document_id": source.canonical_document_id, + "source_file_name": source.title, + "section_path": citation.section_path, + }, + } + + +def _publication_chunk( + *, + source: _DemoSourceDefinition, + chunk: dict[str, Any], +) -> dict[str, Any]: + materialized_chunk = dict(chunk) + metadata = _metadata(materialized_chunk) + raw_path = _first_string(metadata.get("path"), materialized_chunk.get("path")) + publication_path = _publication_path(source=source, raw_path=raw_path) + file_path = _first_string( + metadata.get("file_path"), + metadata.get("filePath"), + materialized_chunk.get("file_path"), + materialized_chunk.get("path") if _is_media_chunk(materialized_chunk) else None, + ) + + metadata["path"] = publication_path + if file_path: + metadata["file_path"] = file_path + materialized_chunk["file_path"] = file_path + materialized_chunk["path"] = publication_path + materialized_chunk["metadata"] = metadata + return materialized_chunk + + +def _publication_path( + *, + source: _DemoSourceDefinition, + raw_path: str | None, +) -> str: + prefix = source.title + raw = str(raw_path or "").strip() + if not raw: + return prefix + + if "-->" in raw: + sections = [part.strip() for part in raw.split("-->")[1:] if part.strip()] + return "/".join([prefix, *sections]) if sections else prefix + + if raw.startswith("images/") or raw.startswith("tables/"): + return f"{prefix}/Assets/{raw}" + + parts = [part.strip() for part in raw.split("/") if part.strip()] + if parts and parts[0] == source.title: + return raw + if len(parts) >= 2 and parts[0] == "Default_Root": + section_parts = parts[2:] if parts[1] == source.title else parts[1:] + return "/".join([prefix, *section_parts]) if section_parts else prefix + return prefix + + +def _resolve_citation_chunk( + *, + source: _DemoSourceDefinition, + citation: _DemoCitationDefinition, + chunks: tuple[dict[str, Any], ...], +) -> dict[str, Any]: + normalized_content = _normalize_text(citation.content) + if normalized_content: + for chunk in chunks: + if normalized_content in _normalize_text(str(chunk.get("content") or "")): + return chunk + + for chunk in chunks: + if str(chunk.get("path") or "") == citation.section_path: + return chunk + + raise ValueError( + "Demo citation does not resolve to a canonical chunk: " + f"demo_source_id={source.demo_source_id}, section_path={citation.section_path}" + ) + + +def _metadata(chunk: dict[str, Any]) -> dict[str, Any]: + metadata = chunk.get("metadata") + return dict(metadata) if isinstance(metadata, dict) else {} + + +def _is_media_chunk(chunk: dict[str, Any]) -> bool: + return _normalize_chunk_type(chunk.get("type")) in {"image", "table"} + + +def _asset_url( + *, + source: _DemoSourceDefinition, + file_path: str | None, +) -> str | None: + if not file_path: + return None + encoded_path = quote(file_path, safe="/") + return f"/api/v1/demo/sources/{source.demo_source_id}/assets/{encoded_path}" + + +def _normalize_chunk_type(value: object) -> str: + raw = str(value or "").strip().split("\n", 1)[0].lower() + return raw if raw in {"text", "image", "table"} else "text" + + +def _first_string(*values: object) -> str | None: + for value in values: + if isinstance(value, str) and value.strip(): + return value.strip() + return None + + +def _normalize_text(value: str) -> str: + return " ".join(value.lower().split()) diff --git a/apps/api/app/services/demo_document_service.py b/apps/api/app/services/demo_document_service.py deleted file mode 100644 index ea50a9b1d..000000000 --- a/apps/api/app/services/demo_document_service.py +++ /dev/null @@ -1,814 +0,0 @@ -"""API-owned canonical demo document catalog and materialization.""" - -from __future__ import annotations - -import json -import math -import shutil -import tempfile -from dataclasses import dataclass -from datetime import datetime, timezone -from functools import lru_cache -from hashlib import blake2b -from pathlib import Path -from typing import Any -from urllib.parse import quote -from uuid import uuid4 - -from sqlalchemy import func, select -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.exceptions.domain_exceptions import ValidationException -from shared.models.database.demo_materialization import DemoMaterialization -from shared.models.database.document import Document -from shared.models.database.job import Job -from shared.models.database.job_result import JobResult -from shared.services.retrieval.cache_service import invalidate_retrieval_cache_namespaces -from shared.services.retrieval.publication_service import RetrievalPublicationService -from shared.services.storage.result_storage import get_result_storage - - -@dataclass(frozen=True) -class DemoCitationDefinition: - """Curated answer citation that resolves to a canonical demo chunk.""" - - section_path: str - description: str - content: str - - -@dataclass(frozen=True) -class DemoExampleDefinition: - """Curated user-facing demo question and answer.""" - - id: str - question: str - answer: str - citations: tuple[DemoCitationDefinition, ...] - - -@dataclass(frozen=True) -class DemoSourceDefinition: - """Canonical demo source metadata and local asset pointers.""" - - demo_source_id: str - canonical_document_id: str - title: str - mime_type: str - size_bytes: int - asset_directory: str - chunk_count: int - examples: tuple[DemoExampleDefinition, ...] - - -@dataclass(frozen=True) -class MaterializedDemoSource: - """User-owned copy of one canonical demo source.""" - - demo_source_id: str - document_id: str - status: str - title: str - mime_type: str - size_bytes: int - chunk_count: int - - -_DATA_ROOT = Path(__file__).resolve().parents[1] / "data" / "demo_documents" -_ASSET_DIRECTORY_NAMES = frozenset({"images", "tables"}) -_DEMO_SOURCE_DEFINITIONS: tuple[DemoSourceDefinition, ...] = ( - DemoSourceDefinition( - demo_source_id="demo-tsla-q4-2025", - canonical_document_id="demo-doc-tsla-q4-2025", - title="TSLA-Q4-2025-Update.pdf", - mime_type="application/pdf", - size_bytes=5_648_867, - asset_directory="tsla-q4-2025", - chunk_count=70, - examples=( - DemoExampleDefinition( - id="demo-tsla-q4-2025-xai", - question="What does the document say about Tesla's xAI investment?", - answer=( - "Tesla entered an agreement on January 16, 2026 to invest " - "approximately $2 billion in xAI Series E Preferred Stock.\n\n" - "The document also says Tesla and xAI entered a framework " - "agreement to evaluate AI collaboration, with the investment " - "expected to close in Q1 2026 subject to customary regulatory " - "conditions." - ), - citations=( - DemoCitationDefinition( - section_path=( - "Default_Root/TSLA-Q4-2025-Update.pdf-->OTHER UPDATES" - ), - description="xAI investment", - content=( - "On January 16, 2026, Tesla entered into an agreement " - "to invest approximately" - ), - ), - ), - ), - DemoExampleDefinition( - id="demo-tsla-q4-2025-energy-storage", - question="What does the document say about energy storage?", - answer=( - "Tesla achieved its highest quarterly energy storage " - "deployments, driven by record Megapack deployments.\n\n" - "Energy gross profit reached a record $1.1 billion, marking " - "the fifth consecutive record quarter.\n\n" - "Tesla also plans to begin Megapack 3 and Megablock " - "production at Megafactory Houston in 2026." - ), - citations=( - DemoCitationDefinition( - section_path=( - "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->" - "Energy generation and storage" - ), - description="Storage deployment growth", - content=( - "We achieved our highest quarterly energy storage " - "deployments, driven by record Megapack deployments." - ), - ), - ), - ), - DemoExampleDefinition( - id="demo-tsla-q4-2025-production-plans", - question="What production plans does Tesla mention for 2026?", - answer=( - "Tesla says Cybercab, Tesla Semi, and Megapack 3 are on " - "schedule for volume production starting in 2026.\n\n" - "The same product update also notes that first-generation " - "Optimus production lines are being installed before volume " - "production." - ), - citations=( - DemoCitationDefinition( - section_path=( - "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->" - "Product" - ), - description="2026 production plans", - content=( - "Cybercab, Tesla Semi and Megapack 3 are on schedule " - "for volume production starting in 2026." - ), - ), - ), - ), - ), - ), -) - - -class DemoDocumentService: - """Serves canonical demo data and copies it into user namespaces.""" - - def __init__( - self, - *, - publication_service: RetrievalPublicationService | None = None, - ) -> None: - self._publication_service = publication_service or RetrievalPublicationService() - - def get_catalog(self) -> dict[str, Any]: - """Return the cacheable canonical demo source catalog.""" - return { - "sources": [ - self._source_catalog_payload(source) - for source in _DEMO_SOURCE_DEFINITIONS - ], - } - - def list_chunks( - self, - *, - demo_source_id: str, - page: int, - page_size: int, - ) -> dict[str, Any] | None: - """Return paginated canonical demo chunks.""" - source = _get_source_definition(demo_source_id) - if source is None: - return None - - chunks = _load_source_chunks(source) - start = (page - 1) * page_size - page_chunks = chunks[start : start + page_size] - return { - "demo_source_id": source.demo_source_id, - "canonical_document_id": source.canonical_document_id, - "title": source.title, - "mime_type": source.mime_type, - "chunks": [ - _chunk_payload(source=source, chunk=chunk) - for chunk in page_chunks - ], - "pagination": { - "page": page, - "page_size": page_size, - "total": len(chunks), - "total_pages": math.ceil(len(chunks) / page_size) if chunks else 0, - }, - } - - def get_chunk( - self, - *, - demo_source_id: str, - demo_chunk_id: str, - ) -> dict[str, Any] | None: - """Return one canonical demo chunk by canonical row id or parser chunk id.""" - source = _get_source_definition(demo_source_id) - if source is None: - return None - - for chunk in _load_source_chunks(source): - if demo_chunk_id in {_canonical_chunk_id(source, chunk), chunk["chunk_id"]}: - return { - "demo_source_id": source.demo_source_id, - "canonical_document_id": source.canonical_document_id, - "chunk": _chunk_payload(source=source, chunk=chunk), - } - - return None - - def get_original_file_path(self, *, demo_source_id: str) -> Path | None: - """Return the canonical original file path for a demo source.""" - source = _get_source_definition(demo_source_id) - if source is None: - return None - - file_path = _source_directory(source) / "original.pdf" - return file_path if file_path.is_file() else None - - def get_asset_file_path( - self, - *, - demo_source_id: str, - asset_path: str, - ) -> Path | None: - """Return a canonical parsed media/table asset path.""" - source = _get_source_definition(demo_source_id) - if source is None: - return None - - source_directory = _source_directory(source).resolve() - normalized_asset_path = _normalize_asset_path(asset_path) - if normalized_asset_path is None: - return None - - candidate = (source_directory / normalized_asset_path).resolve() - if not candidate.is_relative_to(source_directory): - return None - - return candidate if candidate.is_file() else None - - async def materialize_sources( - self, - db: AsyncSession, - *, - user_id: str, - namespace: str, - demo_source_ids: list[str], - ) -> list[MaterializedDemoSource]: - """Copy selected canonical demo sources into a user namespace.""" - selected_demo_source_ids = _deduplicate_source_ids(demo_source_ids) - if not selected_demo_source_ids: - raise ValidationException( - user_message="At least one demo source must be selected.", - violations=[ - { - "field": "demo_source_ids", - "description": "Select one or more demo source IDs.", - } - ], - ) - - selected_sources = [ - _require_source_definition(demo_source_id) - for demo_source_id in selected_demo_source_ids - ] - results: list[MaterializedDemoSource] = [] - for source in selected_sources: - result = await self._materialize_source( - db, - user_id=user_id, - namespace=namespace, - source=source, - ) - results.append(result) - - await db.commit() - await invalidate_retrieval_cache_namespaces( - user_id=user_id, - namespaces=[namespace], - ) - return results - - async def _materialize_source( - self, - db: AsyncSession, - *, - user_id: str, - namespace: str, - source: DemoSourceDefinition, - ) -> MaterializedDemoSource: - await _lock_materialization_scope( - db, - user_id=user_id, - namespace=namespace, - demo_source_id=source.demo_source_id, - ) - existing = await self._get_existing_materialization( - db, - user_id=user_id, - namespace=namespace, - demo_source_id=source.demo_source_id, - ) - if existing is not None and await self._is_active_document( - db, - document_id=existing.document_id, - ): - return _materialized_source_payload( - source=source, - document_id=existing.document_id, - status="existing", - ) - - document_id = f"doc_{uuid4().hex[:12]}" - job_id = f"job_demo_{uuid4().hex[:12]}" - job_result_id = str(uuid4()) - timestamp = _utc_now() - result_bundle = _upload_demo_result_bundle(job_id=job_id, source=source) - - db.add( - Job( - job_id=job_id, - user_id=user_id, - job_type="demo_materialization", - status="done", - source_type="demo", - webhook_enabled=False, - job_metadata={ - "document_id": document_id, - "namespace": namespace, - "source_type": "demo", - "source_file_name": source.title, - "demo_source_id": source.demo_source_id, - }, - version=0, - created_at=timestamp, - updated_at=timestamp, - credits_charged=0, - billing_status="skipped", - ) - ) - db.add( - JobResult( - id=job_result_id, - job_id=job_id, - delivery_mode="inline", - document_metadata={ - "source_file_name": source.title, - "demo_source_id": source.demo_source_id, - }, - inline_payload={"source": "canonical_demo"}, - result_s3_key=result_bundle["zip_key"], - result_size=result_bundle["zip_size"], - created_at=timestamp, - updated_at=timestamp, - ) - ) - await db.flush() - chunks = _publication_chunks(source) - await db.run_sync( - lambda sync_db: self._publication_service.publish_document_state( - sync_db, - job_id=job_id, - job_result_id=job_result_id, - chunks=[dict(chunk) for chunk in chunks], - ) - ) - await db.run_sync( - lambda sync_db: self._publication_service.publish_document_graph( - sync_db, - job_id=job_id, - job_result_id=job_result_id, - ) - ) - await db.flush() - - if existing is None: - db.add( - DemoMaterialization( - user_id=user_id, - namespace=namespace, - demo_source_id=source.demo_source_id, - document_id=document_id, - created_at=timestamp, - updated_at=timestamp, - ) - ) - else: - existing.document_id = document_id - existing.updated_at = timestamp - await db.flush() - return _materialized_source_payload( - source=source, - document_id=document_id, - status="created", - ) - - async def _get_existing_materialization( - self, - db: AsyncSession, - *, - user_id: str, - namespace: str, - demo_source_id: str, - ) -> DemoMaterialization | None: - result = await db.execute( - select(DemoMaterialization) - .where(DemoMaterialization.user_id == user_id) - .where(DemoMaterialization.namespace == namespace) - .where(DemoMaterialization.demo_source_id == demo_source_id) - .with_for_update() - .limit(1) - ) - return result.scalar_one_or_none() - - async def _is_active_document( - self, - db: AsyncSession, - *, - document_id: str, - ) -> bool: - result = await db.execute( - select(Document.document_id) - .where(Document.document_id == document_id) - .where(Document.status == "active") - .limit(1) - ) - return result.scalar_one_or_none() is not None - - def _source_catalog_payload(self, source: DemoSourceDefinition) -> dict[str, Any]: - return { - "demo_source_id": source.demo_source_id, - "canonical_document_id": source.canonical_document_id, - "title": source.title, - "mime_type": source.mime_type, - "size_bytes": source.size_bytes, - "status": "ready", - "chunk_count": source.chunk_count, - "original_file": { - "url": f"/api/v1/demo/sources/{source.demo_source_id}/original", - "mime_type": source.mime_type, - "size_bytes": source.size_bytes, - "can_download": False, - }, - "examples": [ - self._example_payload(source=source, example=example) - for example in source.examples - ], - } - - def _example_payload( - self, - *, - source: DemoSourceDefinition, - example: DemoExampleDefinition, - ) -> dict[str, Any]: - return { - "id": example.id, - "question": example.question, - "answer": example.answer, - "citations": [ - _citation_payload(source=source, citation=citation) - for citation in example.citations - ], - } - - -def _deduplicate_source_ids(demo_source_ids: list[str]) -> list[str]: - selected: list[str] = [] - seen: set[str] = set() - for demo_source_id in demo_source_ids: - normalized = str(demo_source_id).strip() - if not normalized or normalized in seen: - continue - selected.append(normalized) - seen.add(normalized) - return selected - - -async def _lock_materialization_scope( - db: AsyncSession, - *, - user_id: str, - namespace: str, - demo_source_id: str, -) -> None: - lock_id = _materialization_lock_id( - user_id=user_id, - namespace=namespace, - demo_source_id=demo_source_id, - ) - await db.execute(select(func.pg_advisory_xact_lock(lock_id))) - - -def _materialization_lock_id( - *, - user_id: str, - namespace: str, - demo_source_id: str, -) -> int: - lock_key = f"{user_id}\0{namespace}\0{demo_source_id}" - digest = blake2b(lock_key.encode("utf-8"), digest_size=8).digest() - return int.from_bytes(digest, byteorder="big", signed=True) - - -def _materialized_source_payload( - *, - source: DemoSourceDefinition, - document_id: str, - status: str, -) -> MaterializedDemoSource: - return MaterializedDemoSource( - demo_source_id=source.demo_source_id, - document_id=document_id, - status=status, - title=source.title, - mime_type=source.mime_type, - size_bytes=source.size_bytes, - chunk_count=source.chunk_count, - ) - - -def _upload_demo_result_bundle( - *, - job_id: str, - source: DemoSourceDefinition, -) -> dict[str, int | str]: - """Upload canonical demo result files so copied media URLs resolve.""" - source_directory = _source_directory(source) - with tempfile.TemporaryDirectory(prefix="knowhere-demo-result-") as temp_directory: - zip_base_path = Path(temp_directory) / job_id - zip_file_path = Path( - shutil.make_archive( - str(zip_base_path), - "zip", - root_dir=source_directory, - ) - ) - zip_size = zip_file_path.stat().st_size - bundle = get_result_storage().upload( - job_id=job_id, - result_dir=str(source_directory), - zip_file_path=str(zip_file_path), - ) - - return { - "zip_key": bundle.zip_key, - "zip_size": zip_size, - } - - -def _publication_chunks(source: DemoSourceDefinition) -> list[dict[str, Any]]: - return [ - _publication_chunk(source=source, chunk=chunk) - for chunk in _load_source_chunks(source) - ] - - -def _publication_chunk( - *, - source: DemoSourceDefinition, - chunk: dict[str, Any], -) -> dict[str, Any]: - materialized_chunk = dict(chunk) - metadata = _metadata(materialized_chunk) - raw_path = _first_string(metadata.get("path"), materialized_chunk.get("path")) - publication_path = _publication_path(source=source, raw_path=raw_path) - file_path = _first_string( - metadata.get("file_path"), - metadata.get("filePath"), - materialized_chunk.get("file_path"), - materialized_chunk.get("path") if _is_media_chunk(materialized_chunk) else None, - ) - - metadata["path"] = publication_path - if file_path: - metadata["file_path"] = file_path - materialized_chunk["file_path"] = file_path - materialized_chunk["path"] = publication_path - materialized_chunk["metadata"] = metadata - return materialized_chunk - - -def _publication_path( - *, - source: DemoSourceDefinition, - raw_path: str | None, -) -> str: - prefix = f"Default_Root/{source.title}" - raw = str(raw_path or "").strip() - if not raw: - return prefix - - if "-->" in raw: - sections = [ - part.strip() - for part in raw.split("-->")[1:] - if part.strip() - ] - return "/".join([prefix, *sections]) if sections else prefix - - if raw.startswith("images/") or raw.startswith("tables/"): - return f"{prefix}/Assets/{raw}" - - parts = [part.strip() for part in raw.split("/") if part.strip()] - if len(parts) >= 2 and parts[0] == "Default_Root": - return raw - return prefix - - -def _normalize_asset_path(asset_path: str) -> Path | None: - normalized = str(asset_path or "").strip().replace("\\", "/").lstrip("/") - parts = [part for part in normalized.split("/") if part and part != "."] - if not parts or parts[0] not in _ASSET_DIRECTORY_NAMES: - return None - if any(part == ".." or part.startswith(".") for part in parts): - return None - return Path(*parts) - - -def _citation_payload( - *, - source: DemoSourceDefinition, - citation: DemoCitationDefinition, -) -> dict[str, Any]: - chunk = _resolve_citation_chunk(source=source, citation=citation) - return { - "demo_source_id": source.demo_source_id, - "canonical_document_id": source.canonical_document_id, - "canonical_chunk_id": _canonical_chunk_id(source, chunk), - "chunk_id": chunk["chunk_id"], - "chunk_type": _normalize_chunk_type(chunk.get("type")), - "content": citation.content, - "description": citation.description, - "source": { - "document_id": source.canonical_document_id, - "source_file_name": source.title, - "section_path": citation.section_path, - }, - } - - -def _resolve_citation_chunk( - *, - source: DemoSourceDefinition, - citation: DemoCitationDefinition, -) -> dict[str, Any]: - chunks = _load_source_chunks(source) - normalized_content = _normalize_text(citation.content) - if normalized_content: - for chunk in chunks: - if normalized_content in _normalize_text(str(chunk.get("content") or "")): - return chunk - - for chunk in chunks: - if str(chunk.get("path") or "") == citation.section_path: - return chunk - - raise ValueError( - "Demo citation does not resolve to a canonical chunk: " - f"demo_source_id={source.demo_source_id}, section_path={citation.section_path}" - ) - - -def _chunk_payload( - *, - source: DemoSourceDefinition, - chunk: dict[str, Any], -) -> dict[str, Any]: - metadata = _metadata(chunk) - file_path = _first_string( - metadata.get("file_path"), - metadata.get("filePath"), - chunk.get("file_path"), - chunk.get("path") if _is_media_chunk(chunk) else None, - ) - return { - "id": _canonical_chunk_id(source, chunk), - "chunk_id": chunk["chunk_id"], - "chunk_type": _normalize_chunk_type(chunk.get("type")), - "content": str(chunk.get("content") or ""), - "section_path": str(chunk.get("path") or "") or None, - "source_chunk_path": str(chunk.get("path") or "") or None, - "file_path": file_path, - "sort_order": _sort_order(source=source, chunk=chunk), - "metadata": metadata, - "asset_url": _asset_url(source=source, file_path=file_path), - "created_at": None, - } - - -def _sort_order( - *, - source: DemoSourceDefinition, - chunk: dict[str, Any], -) -> int: - try: - return _load_source_chunks(source).index(chunk) - except ValueError: - return 0 - - -def _metadata(chunk: dict[str, Any]) -> dict[str, Any]: - metadata = chunk.get("metadata") - return dict(metadata) if isinstance(metadata, dict) else {} - - -def _is_media_chunk(chunk: dict[str, Any]) -> bool: - return _normalize_chunk_type(chunk.get("type")) in {"image", "table"} - - -def _canonical_chunk_id( - source: DemoSourceDefinition, - chunk: dict[str, Any], -) -> str: - return f"{source.demo_source_id}:{chunk['chunk_id']}" - - -def _asset_url( - *, - source: DemoSourceDefinition, - file_path: str | None, -) -> str | None: - if not file_path: - return None - encoded_path = quote(file_path, safe="/") - return f"/api/v1/demo/sources/{source.demo_source_id}/assets/{encoded_path}" - - -def _normalize_chunk_type(value: object) -> str: - raw = str(value or "").strip().split("\n", 1)[0].lower() - return raw if raw in {"text", "image", "table"} else "text" - - -def _first_string(*values: object) -> str | None: - for value in values: - if isinstance(value, str) and value.strip(): - return value.strip() - return None - - -def _normalize_text(value: str) -> str: - return " ".join(value.lower().split()) - - -def _get_source_definition(demo_source_id: str) -> DemoSourceDefinition | None: - return next( - ( - source - for source in _DEMO_SOURCE_DEFINITIONS - if source.demo_source_id == demo_source_id - ), - None, - ) - - -def _require_source_definition(demo_source_id: str) -> DemoSourceDefinition: - source = _get_source_definition(demo_source_id) - if source is None: - raise KeyError(demo_source_id) - return source - - -def _source_directory(source: DemoSourceDefinition) -> Path: - return _DATA_ROOT / source.asset_directory - - -@lru_cache(maxsize=8) -def _load_source_chunks(source: DemoSourceDefinition) -> tuple[dict[str, Any], ...]: - chunks_path = _source_directory(source) / "chunks.json" - with chunks_path.open("r", encoding="utf-8") as file: - payload = json.load(file) - - chunks = payload.get("chunks") if isinstance(payload, dict) else None - if not isinstance(chunks, list): - return () - - return tuple( - dict(chunk) - for chunk in chunks - if isinstance(chunk, dict) and isinstance(chunk.get("chunk_id"), str) - ) - - -def _utc_now() -> datetime: - return datetime.now(timezone.utc).replace(tzinfo=None) diff --git a/apps/api/app/services/document_ingestion/__init__.py b/apps/api/app/services/document_ingestion/__init__.py new file mode 100644 index 000000000..0e67b8017 --- /dev/null +++ b/apps/api/app/services/document_ingestion/__init__.py @@ -0,0 +1,3 @@ +from app.services.document_ingestion.service import DocumentIngestionService + +__all__ = ["DocumentIngestionService"] diff --git a/apps/api/app/services/document_ingestion/confirmation_service.py b/apps/api/app/services/document_ingestion/confirmation_service.py new file mode 100644 index 000000000..4cc3a7b8e --- /dev/null +++ b/apps/api/app/services/document_ingestion/confirmation_service.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +from app.repositories.job_repository import JobRepository +from app.services.document_ingestion.handoff_service import ( + DocumentIngestionHandoffService, +) +from app.services.jobs import check_job_permission +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + JobOperationException, + NotFoundException, + PermissionDeniedException, + UnavailableException, + ValidationException, +) +from shared.core.state_machine.states import JobStatus +from shared.services.storage.file_upload_service import FileUploadService + + +class DocumentIngestionConfirmationService: + def __init__( + self, + *, + job_repository: JobRepository | None = None, + file_upload_service: FileUploadService | None = None, + handoff_service: DocumentIngestionHandoffService | None = None, + ) -> None: + self._job_repository = job_repository or JobRepository() + self._file_upload_service = file_upload_service or FileUploadService() + self._handoff_service = handoff_service or DocumentIngestionHandoffService() + + async def confirm_upload( + self, + db: AsyncSession, + *, + job_id: str, + user_id: str, + ) -> dict[str, str]: + try: + job = await self._job_repository.get_job_by_id(db, job_id) + check_job_permission(job, user_id, job_id) + assert job is not None + + logger.info(f"Confirm upload - Job {job_id} current status: {job.status}") + if job.status not in [JobStatus.PENDING.value, JobStatus.WAITING_FILE.value]: + logger.info(f"Job {job_id} already processed, status: {job.status}") + return {"message": "Job status already updated"} + + if not job.s3_key: + raise ValidationException( + user_message="Job is missing S3 key information", + violations=[ + { + "field": "s3_key", + "description": "S3 key not set for this job", + } + ], + ) + + file_info = await self._file_upload_service.verify_s3_file_exists(job.s3_key) + if not bool(file_info.get("exists")): + raise ValidationException( + user_message="S3 file does not exist, please upload the file first", + violations=[ + {"field": "file", "description": "File not found in S3"} + ], + ) + + await self._handoff_service.start_uploaded_file_workflow( + db=db, + job=job, + user_id=user_id, + trigger="manual_upload_completed", + ) + return {"message": "File upload confirmed; processing started"} + except NotFoundException: + raise + except PermissionDeniedException: + raise + except UnavailableException: + raise + except ValidationException: + raise + except Exception as exc: + logger.error(f"Failed to confirm upload: {exc}") + raise JobOperationException( + internal_message=f"Failed to confirm upload: {str(exc)}" + ) diff --git a/apps/api/app/services/document_ingestion/creation_service.py b/apps/api/app/services/document_ingestion/creation_service.py new file mode 100644 index 000000000..9ca6d2ed6 --- /dev/null +++ b/apps/api/app/services/document_ingestion/creation_service.py @@ -0,0 +1,323 @@ +from __future__ import annotations + +import os +import uuid +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import cast +from urllib.parse import urlparse + +from app.repositories.job_repository import JobRepository +from app.services.document_ingestion.scope_service import ( + is_active_document_job_unique_violation, + raise_document_ingestion_conflict, +) +from app.services.jobs.result_projection import to_job_status_value +from app.services.rate_limit.data_structures import CurrentUser +from loguru import logger +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + JobOperationException, + ValidationException, +) +from shared.core.state_machine.states import JobStatus +from shared.models.database.job import Job +from shared.models.schemas.job import JobCreate, JobResponse +from shared.models.schemas.job_metadata import JobMetadataHelper +from shared.services.redis import JobInfoRedisService, RedisServiceFactory +from shared.services.redis.job_metadata_service import JobMetadataService +from shared.services.storage.file_upload_service import FileUploadService +from shared.services.http.url_file_type import resolve_file_extension_async + +_DOCUMENT_INGESTION_JOB_TYPE = "document_ingestion" +_URL_UPLOAD_TASK_NAME = ( + "app.core.tasks.document_ingestion_tasks.upload_url_file_task" +) +JobMetadata = dict[str, object] +UploadHeaders = dict[str, str] + + +@dataclass(frozen=True) +class ResolvedDocumentIngestionScope: + job_metadata: JobMetadata + document_id: str + namespace: str + + +class DocumentIngestionCreationService: + def __init__( + self, + *, + job_repository: JobRepository | None = None, + file_upload_service: FileUploadService | None = None, + ) -> None: + self._job_repository = job_repository or JobRepository() + self._file_upload_service = file_upload_service or FileUploadService() + + async def create_job( + self, + db: AsyncSession, + *, + payload: JobCreate, + job_id: str, + current_user: CurrentUser, + scope: ResolvedDocumentIngestionScope, + ) -> JobResponse: + if payload.source_type == "file": + return await self._create_file_job( + db, + payload=payload, + job_id=job_id, + current_user=current_user, + scope=scope, + ) + return await self._create_url_job( + db, + payload=payload, + job_id=job_id, + current_user=current_user, + scope=scope, + ) + + async def _create_waiting_job( + self, + db: AsyncSession, + *, + job_id: str, + user_id: str, + source_type: str, + webhook_url: str | None, + job_metadata: JobMetadata, + s3_key: str, + document_id: str, + ) -> Job: + try: + job = await self._job_repository.create_job( + db=db, + job_id=job_id, + user_id=user_id, + job_type=_DOCUMENT_INGESTION_JOB_TYPE, + source_type=source_type, + file_path=None, + webhook_url=webhook_url, + metadata=job_metadata, + initial_state=JobStatus.WAITING_FILE.value, + s3_key=s3_key, + ) + except IntegrityError as exc: + if is_active_document_job_unique_violation(exc): + raise_document_ingestion_conflict(document_id=document_id) + raise + + if job is None: + raise JobOperationException( + internal_message="Failed to create job in database" + ) + return job + + async def _cache_job_creation_state( + self, + *, + job_id: str, + s3_key: str, + user_id: str, + webhook_enabled: bool, + source_type: str, + job_metadata: JobMetadata, + ) -> None: + redis_service = RedisServiceFactory.get_service() + metadata_service = JobMetadataService(redis_service) + await metadata_service.save_metadata(job_id, job_metadata) + + job_info_service = JobInfoRedisService(redis_service) + job_info: dict[str, object] = { + "job_id": job_id, + "s3_key": s3_key, + "user_id": user_id, + "webhook_enabled": webhook_enabled, + "job_type": _DOCUMENT_INGESTION_JOB_TYPE, + "source_type": source_type, + "created_at": datetime.now(timezone.utc).isoformat(), + } + await job_info_service.save_job_info(job_id, job_info) + + async def _create_file_job( + self, + db: AsyncSession, + *, + payload: JobCreate, + job_id: str, + current_user: CurrentUser, + scope: ResolvedDocumentIngestionScope, + ) -> JobResponse: + assert payload.file_name is not None + file_extension = os.path.splitext(payload.file_name)[1] + s3_key = f"uploads/{job_id}{file_extension}" + JobMetadataHelper.set_file_source( + scope.job_metadata, + source_file_name=payload.file_name, + ) + + job = await self._create_waiting_job( + db, + job_id=job_id, + user_id=current_user.user_id, + source_type="file", + webhook_url=payload.webhook.url if payload.webhook else None, + job_metadata=scope.job_metadata, + s3_key=s3_key, + document_id=scope.document_id, + ) + + upload_info = await self._file_upload_service.generate_upload_url( + job_id, + file_extension, + ) + upload_url = cast(str, upload_info["upload_url"]) + upload_headers = cast(UploadHeaders, upload_info["upload_headers"]) + expires_in = cast(int, upload_info["expires_in"]) + + await self._cache_job_creation_state( + job_id=job_id, + s3_key=s3_key, + user_id=current_user.user_id, + webhook_enabled=bool(payload.webhook and payload.webhook.url), + source_type="file", + job_metadata=scope.job_metadata, + ) + + logger.info(f"Job {job_id} upload_url returned to client: {upload_url}") + return _build_job_response( + job_id=job_id, + job=job, + source_type="file", + data_id=payload.data_id, + namespace=scope.namespace, + upload_url=upload_url, + upload_headers=upload_headers, + expires_in=expires_in, + ) + + async def _create_url_job( + self, + db: AsyncSession, + *, + payload: JobCreate, + job_id: str, + current_user: CurrentUser, + scope: ResolvedDocumentIngestionScope, + ) -> JobResponse: + assert payload.source_url is not None + file_extension = await resolve_file_extension_async(payload.source_url) + if not file_extension: + raise ValidationException( + user_message=( + "Unsupported URL file type. Supported formats: " + f"{_get_supported_formats()}" + ), + violations=[ + { + "field": "source_url", + "description": "URL file type not supported", + } + ], + ) + + source_file_name = _resolve_url_source_file_name( + source_url=payload.source_url, + file_extension=file_extension, + ) + s3_key = f"uploads/{job_id}{file_extension}" + JobMetadataHelper.set_url_source( + scope.job_metadata, + source_file_name=source_file_name, + source_url=payload.source_url, + ) + + job = await self._create_waiting_job( + db, + job_id=job_id, + user_id=current_user.user_id, + source_type="url", + webhook_url=payload.webhook.url if payload.webhook else None, + job_metadata=scope.job_metadata, + s3_key=s3_key, + document_id=scope.document_id, + ) + + await self._cache_job_creation_state( + job_id=job_id, + s3_key=s3_key, + user_id=current_user.user_id, + webhook_enabled=bool(payload.webhook and payload.webhook.url), + source_type="url", + job_metadata=scope.job_metadata, + ) + _schedule_url_upload( + job_id=job_id, + source_url=payload.source_url, + user_id=current_user.user_id, + ) + + return _build_job_response( + job_id=job_id, + job=job, + source_type="url", + data_id=payload.data_id, + namespace=scope.namespace, + ) + + +def _get_supported_formats() -> str: + from shared.core.config import settings + + return ", ".join(sorted(settings.get_supported_extensions())) + + +def _build_job_response( + *, + job_id: str, + job: Job, + source_type: str, + data_id: str | None, + namespace: str | None = None, + document_id: str | None = None, + upload_url: str | None = None, + upload_headers: UploadHeaders | None = None, + expires_in: int | None = None, +) -> JobResponse: + return JobResponse( + job_id=job_id, + status=to_job_status_value(job.status), + source_type=source_type, + data_id=data_id, + namespace=namespace, + document_id=document_id, + created_at=job.created_at, + upload_url=upload_url, + upload_headers=upload_headers, + expires_in=expires_in, + ) + + +def _resolve_url_source_file_name(*, source_url: str, file_extension: str) -> str: + parsed_url = urlparse(source_url) + url_basename = str(os.path.basename(parsed_url.path)) + if url_basename and os.path.splitext(url_basename)[1].lower() == file_extension: + return url_basename + if url_basename: + return f"{url_basename}{file_extension}" + return f"url_file_{uuid.uuid4().hex[:8]}{file_extension}" + + +def _schedule_url_upload(*, job_id: str, source_url: str, user_id: str) -> None: + from shared.core.celery_app import get_celery_app + + celery_app = get_celery_app() + upload_url_file_task = celery_app.signature(_URL_UPLOAD_TASK_NAME) + upload_url_file_task.apply_async( + args=[job_id, source_url, user_id], + kwargs={"job_type": _DOCUMENT_INGESTION_JOB_TYPE}, + ) diff --git a/apps/api/app/services/document_ingestion/handoff_service.py b/apps/api/app/services/document_ingestion/handoff_service.py new file mode 100644 index 000000000..c8de180bd --- /dev/null +++ b/apps/api/app/services/document_ingestion/handoff_service.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +from typing import Protocol + +from app.services.document_ingestion.worker_dispatcher import ( + DocumentIngestionWorkerDispatcher, +) +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ( + UnavailableException, + ValidationException, +) +from shared.core.state_machine.service import AsyncStateMachineService +from shared.core.state_machine.states import JobStatus + +_DOCUMENT_PARSE_JOB_TYPE = "document_ingestion" +_LEGACY_DOCUMENT_PARSE_JOB_TYPE = "kb_management" +_SUPPORTED_DOCUMENT_PARSE_JOB_TYPES = frozenset( + { + _DOCUMENT_PARSE_JOB_TYPE, + _LEGACY_DOCUMENT_PARSE_JOB_TYPE, + } +) + + +class _UploadedFileJob(Protocol): + job_id: str + job_type: str + + +class DocumentIngestionHandoffService: + """Advance uploaded Document Ingestion Jobs into worker parsing.""" + + def __init__( + self, + *, + state_machine: AsyncStateMachineService | None = None, + worker_dispatcher: DocumentIngestionWorkerDispatcher | None = None, + ) -> None: + self._state_machine = state_machine or AsyncStateMachineService() + self._worker_dispatcher = ( + worker_dispatcher or DocumentIngestionWorkerDispatcher() + ) + + async def start_uploaded_file_workflow( + self, + db: AsyncSession, + *, + job: _UploadedFileJob, + user_id: str, + trigger: str, + ) -> None: + if job.job_type not in _SUPPORTED_DOCUMENT_PARSE_JOB_TYPES: + raise ValidationException( + user_message="Unsupported job type", + violations=[ + { + "field": "job_type", + "description": ( + f"Job type '{job.job_type}' is not supported" + ), + } + ], + ) + + outcome = await self._state_machine.transition_outcome( + db, + job.job_id, + JobStatus.PENDING.value, + trigger, + None, + "system", + ) + if not outcome.succeeded: + logger.warning( + "Upload handoff transition rejected: " + f"job_id={job.job_id}, reason={outcome.reason}" + ) + raise UnavailableException( + internal_message=( + f"Could not advance uploaded job {job.job_id} to pending: " + f"{outcome.reason}" + ), + retry_after=settings.DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN, + user_message="Job state is still settling. Retrying shortly.", + ) + + await self._worker_dispatcher.start_uploaded_file_parse( + job_id=job.job_id, + user_id=user_id, + ) + + async def mark_upload_expired( + self, + db: AsyncSession, + *, + job: _UploadedFileJob, + ) -> None: + outcome = await self._state_machine.mark_failed_outcome( + db, + job.job_id, + "Upload expired: file was not uploaded within the allowed time window", + error_code="UPLOAD_EXPIRED", + ) + if not outcome.succeeded: + logger.warning( + "Upload expiry transition rejected: " + f"job_id={job.job_id}, reason={outcome.reason}" + ) diff --git a/apps/api/app/services/job_document_scope_service.py b/apps/api/app/services/document_ingestion/scope_service.py similarity index 85% rename from apps/api/app/services/job_document_scope_service.py rename to apps/api/app/services/document_ingestion/scope_service.py index 80625fe87..afc13ada0 100644 --- a/apps/api/app/services/job_document_scope_service.py +++ b/apps/api/app/services/document_ingestion/scope_service.py @@ -1,5 +1,5 @@ """ -Document-scope rules used by job creation/update flows. +Document-scope rules used by document-ingestion workflows. """ from __future__ import annotations @@ -17,6 +17,7 @@ ValidationException, ) from shared.models.database.job import Job +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace _ACTIVE_JOB_STATUSES = ("waiting-file", "pending", "running", "converting") @@ -74,8 +75,9 @@ async def resolve_effective_document_scope( requested_namespace: Optional[str], repository: DocumentRepository | None = None, ) -> tuple[str, str]: + effective_requested_namespace = normalize_retrieval_namespace(requested_namespace) if not document_id: - return f"doc_{uuid.uuid4().hex[:12]}", requested_namespace or "default" + return f"doc_{uuid.uuid4().hex[:12]}", effective_requested_namespace document = await (repository or DocumentRepository()).get_document( db, @@ -88,7 +90,8 @@ async def resolve_effective_document_scope( resource_id=document_id, internal_message=f"Document not found for update flow: {document_id}", ) - if requested_namespace and requested_namespace != document.namespace: + has_requested_namespace = bool(str(requested_namespace or "").strip()) + if has_requested_namespace and effective_requested_namespace != document.namespace: raise ValidationException( user_message="namespace must match the existing document namespace", violations=[ diff --git a/apps/api/app/services/document_ingestion/service.py b/apps/api/app/services/document_ingestion/service.py new file mode 100644 index 000000000..7313c07f8 --- /dev/null +++ b/apps/api/app/services/document_ingestion/service.py @@ -0,0 +1,263 @@ +from __future__ import annotations + +import os +import uuid +from typing import cast + +from app.services.document_ingestion.confirmation_service import ( + DocumentIngestionConfirmationService, +) +from app.services.document_ingestion.creation_service import ( + DocumentIngestionCreationService, + ResolvedDocumentIngestionScope, +) +from app.services.document_ingestion.scope_service import ( + find_active_job_for_document, + raise_document_ingestion_conflict, + resolve_effective_document_scope, +) +from app.services.rate_limit.data_structures import CurrentUser +from app.services.rate_limit.job_admission_service import JobAdmissionService +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ( + ConflictException, + JobOperationException, + NotFoundException, + PermissionDeniedException, + RateLimitException, + UnavailableException, + ValidationException, +) +from shared.core.exceptions.webhook_exceptions import WebhookConfigException +from shared.models.schemas.job import ConfirmUploadRequest, JobCreate, JobResponse +from shared.models.schemas.job_metadata import JobMetadataHelper +from shared.services.http.url_file_type import resolve_file_extension_async +from shared.services.http.url_security import validate_http_url_and_resolve_ip_async + +JobMetadata = dict[str, object] + + +class DocumentIngestionService: + def __init__( + self, + *, + creation_service: DocumentIngestionCreationService | None = None, + confirmation_service: DocumentIngestionConfirmationService | None = None, + job_admission_service: JobAdmissionService | None = None, + ) -> None: + self._creation_service = creation_service or DocumentIngestionCreationService() + self._confirmation_service = ( + confirmation_service or DocumentIngestionConfirmationService() + ) + self._job_admission_service = job_admission_service or JobAdmissionService() + + async def create_job( + self, + db: AsyncSession, + *, + payload: JobCreate, + current_user: CurrentUser, + ) -> JobResponse: + try: + job_id = f"job_{uuid.uuid4().hex[:12]}" + await self._validate_create_payload(payload) + scope = await self._resolve_scope( + db, + payload=payload, + current_user=current_user, + ) + + await self._job_admission_service.enforce_job_creation_capacity( + db=db, + current_user=current_user, + ) + + return await self._creation_service.create_job( + db, + payload=payload, + job_id=job_id, + current_user=current_user, + scope=scope, + ) + except NotFoundException: + raise + except ValidationException: + raise + except ConflictException: + raise + except WebhookConfigException: + raise + except (RateLimitException, UnavailableException): + raise + except JobOperationException: + raise + except Exception as exc: + logger.error(f"Failed to create job: {exc}") + raise JobOperationException( + internal_message=f"Job creation failed: {str(exc)}" + ) + + async def confirm_upload( + self, + db: AsyncSession, + *, + job_id: str, + request_payload: ConfirmUploadRequest | None, + user_id: str, + ) -> dict[str, str]: + del request_payload + + try: + return await self._confirmation_service.confirm_upload( + db=db, + job_id=job_id, + user_id=user_id, + ) + except NotFoundException: + raise + except PermissionDeniedException: + raise + except UnavailableException: + raise + except ValidationException: + raise + except Exception as exc: + logger.error(f"Failed to confirm upload: {exc}") + raise JobOperationException( + internal_message=f"Failed to confirm upload: {str(exc)}" + ) + + async def _validate_create_payload(self, payload: JobCreate) -> None: + if payload.source_type == "file" and not payload.file_name: + raise ValidationException( + user_message="file_name is required when source_type is 'file'", + violations=[ + { + "field": "file_name", + "description": "Required for file source type", + } + ], + ) + if payload.source_type == "url" and not payload.source_url: + raise ValidationException( + user_message="source_url is required when source_type is 'url'", + violations=[ + { + "field": "source_url", + "description": "Required for url source type", + } + ], + ) + + if payload.webhook and payload.webhook.url: + validation_result = await validate_http_url_and_resolve_ip_async( + payload.webhook.url, + ) + if not validation_result.is_valid: + raise WebhookConfigException( + user_message="Invalid webhook URL", + internal_message=( + "Webhook validation failed: " + f"{validation_result.error_message}" + ), + ) + + if ( + payload.source_type == "file" + and payload.file_name + and not _is_supported_file_name(payload.file_name) + ): + raise ValidationException( + user_message=( + "Unsupported file type. Supported formats: " + f"{_get_supported_formats()}" + ), + violations=[ + {"field": "file_name", "description": "File type not supported"} + ], + ) + + if payload.source_type == "url": + assert payload.source_url is not None + file_extension = await resolve_file_extension_async(payload.source_url) + if not file_extension: + raise ValidationException( + user_message=( + "Unsupported URL file type. Supported formats: " + f"{_get_supported_formats()}" + ), + violations=[ + { + "field": "source_url", + "description": "URL file type not supported", + } + ], + ) + + async def _resolve_scope( + self, + db: AsyncSession, + *, + payload: JobCreate, + current_user: CurrentUser, + ) -> ResolvedDocumentIngestionScope: + job_metadata = cast(JobMetadata, JobMetadataHelper.create_from_request(payload)) + requested_document_id = JobMetadataHelper.get_document_id(job_metadata) + if requested_document_id: + active_job = await find_active_job_for_document( + db, + user_id=current_user.user_id, + document_id=requested_document_id, + ) + if active_job is not None: + raise_document_ingestion_conflict( + document_id=requested_document_id, + active_job_id=active_job.job_id, + ) + + ( + effective_document_id, + effective_namespace, + ) = await resolve_effective_document_scope( + db, + user_id=current_user.user_id, + document_id=requested_document_id, + requested_namespace=cast(str | None, payload.namespace), + ) + + if not requested_document_id: + active_job = await find_active_job_for_document( + db, + user_id=current_user.user_id, + document_id=effective_document_id, + ) + if active_job is not None: + raise_document_ingestion_conflict( + document_id=effective_document_id, + active_job_id=active_job.job_id, + ) + + JobMetadataHelper.set_document_scope( + job_metadata, + document_id=effective_document_id, + namespace=effective_namespace, + ) + return ResolvedDocumentIngestionScope( + job_metadata=job_metadata, + document_id=effective_document_id, + namespace=effective_namespace, + ) + + +def _get_supported_formats() -> str: + return ", ".join(sorted(settings.get_supported_extensions())) + + +def _is_supported_file_name(file_name: str) -> bool: + if not file_name: + return False + file_extension = os.path.splitext(file_name)[1].lower() + return file_extension in settings.get_supported_extensions() diff --git a/apps/api/app/services/document_ingestion/worker_dispatcher.py b/apps/api/app/services/document_ingestion/worker_dispatcher.py new file mode 100644 index 000000000..e340c8647 --- /dev/null +++ b/apps/api/app/services/document_ingestion/worker_dispatcher.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from celery import signature +from celery.canvas import Signature +from loguru import logger + +from shared.core.celery_router import CeleryTaskRouter, task_router +from shared.core.exceptions.domain_exceptions import WorkerHandlingException + +_DOCUMENT_PARSE_JOB_TYPE = "document_ingestion" +_DOCUMENT_PARSE_TASK_NAME = "app.core.tasks.document_ingestion_tasks.parse_task" + + +class DocumentIngestionWorkerDispatcher: + """Dispatch uploaded Document Ingestion Jobs to worker parsing.""" + + def __init__( + self, + *, + celery_task_router: CeleryTaskRouter | None = None, + ) -> None: + self._task_router = celery_task_router or task_router + + async def start_uploaded_file_parse(self, *, job_id: str, user_id: str) -> str: + task_signature = self._build_uploaded_file_parse_signature( + job_id=job_id, + user_id=user_id, + ) + result = task_signature.apply_async() + if result is None or result.id is None: + raise WorkerHandlingException( + internal_message=( + "Failed to start Document Ingestion worker parse: " + "missing Celery task id" + ) + ) + + task_id = str(result.id) + signature_options = task_signature.options or {} + queue_name = signature_options.get("queue") + logger.info( + "Document Ingestion worker parse started: " + f"job_id={job_id}, task_id={task_id}, queue={queue_name}" + ) + return task_id + + def _build_uploaded_file_parse_signature( + self, + *, + job_id: str, + user_id: str, + ) -> Signature: + queue_name = self._task_router.get_queue_for_job( + _DOCUMENT_PARSE_JOB_TYPE, + user_id, + ) + task_kwargs: dict[str, str] = { + "user_id": user_id, + "job_type": _DOCUMENT_PARSE_JOB_TYPE, + } + task_signature = signature( + _DOCUMENT_PARSE_TASK_NAME, + args=[job_id], + kwargs=task_kwargs, + ) + if task_signature is None: + raise WorkerHandlingException( + internal_message=( + "Failed to build Document Ingestion worker parse: " + "missing Celery signature" + ) + ) + + configured_signature = task_signature.set(queue=queue_name) + if configured_signature is None: + raise WorkerHandlingException( + internal_message=( + "Failed to configure Document Ingestion worker parse queue" + ) + ) + + return configured_signature diff --git a/apps/api/app/services/documents/__init__.py b/apps/api/app/services/documents/__init__.py new file mode 100644 index 000000000..53fc5491c --- /dev/null +++ b/apps/api/app/services/documents/__init__.py @@ -0,0 +1 @@ +"""Document lifecycle workflow package.""" diff --git a/apps/api/app/services/document_service.py b/apps/api/app/services/documents/lifecycle_service.py similarity index 98% rename from apps/api/app/services/document_service.py rename to apps/api/app/services/documents/lifecycle_service.py index 13dddd3e8..9986f29ad 100644 --- a/apps/api/app/services/document_service.py +++ b/apps/api/app/services/documents/lifecycle_service.py @@ -1,6 +1,4 @@ -""" -Application service for document lifecycle routes. -""" +"""Application workflow for document lifecycle routes.""" from __future__ import annotations @@ -16,7 +14,7 @@ from shared.services.retrieval.cache_service import ( invalidate_retrieval_cache_namespaces, ) -from shared.services.retrieval.graph_service import DocumentGraphService, GraphScope +from shared.services.retrieval.graph.service import DocumentGraphService, GraphScope def document_payload(document) -> dict[str, Any]: diff --git a/apps/api/app/services/guest/__init__.py b/apps/api/app/services/guest/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/apps/api/app/services/guest/guest_registration_service.py b/apps/api/app/services/guest/guest_registration_service.py index 58069fac2..4f189807e 100644 --- a/apps/api/app/services/guest/guest_registration_service.py +++ b/apps/api/app/services/guest/guest_registration_service.py @@ -149,8 +149,8 @@ async def _create_api_key_without_commit( ) -> str: """Generate an API key record and flush (but do not commit). - This avoids the internal commit inside APIKeyService.create_api_key() - which would make the key durable before the device row is inserted. + This keeps the guest API key row in the same transaction as the device + row so guest registration stays atomic. """ from shared.models.database.api_key import APIKey diff --git a/apps/api/app/services/jobs/__init__.py b/apps/api/app/services/jobs/__init__.py new file mode 100644 index 000000000..fcdc2ccdc --- /dev/null +++ b/apps/api/app/services/jobs/__init__.py @@ -0,0 +1,11 @@ +from app.services.jobs.read_service import ( + check_job_permission, + get_job_result_for_user, + list_jobs_for_user, +) + +__all__ = [ + "check_job_permission", + "get_job_result_for_user", + "list_jobs_for_user", +] diff --git a/apps/api/app/services/jobs/job_read_model.py b/apps/api/app/services/jobs/job_read_model.py new file mode 100644 index 000000000..1d0a4ab39 --- /dev/null +++ b/apps/api/app/services/jobs/job_read_model.py @@ -0,0 +1,209 @@ +from __future__ import annotations + +import math +from datetime import datetime, timedelta, timezone +from typing import Optional + +from app.repositories.job_repository import JobRepository +from app.services.jobs.result_projection import ( + build_job_result_response, + to_job_status_value, +) +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + JobOperationException, + NotFoundException, + PermissionDeniedException, + ValidationException, +) +from shared.models.schemas.job import JobList, JobResultResponse +from shared.services.redis import RedisServiceFactory +from shared.utils.utc_now import utc_now_naive + + +class JobReadModel: + async def list_jobs_for_user( + self, + db: AsyncSession, + *, + user_id: str, + page: int, + page_size: int, + job_status: Optional[str], + job_type: Optional[str], + recent_days: Optional[int], + start_time: Optional[datetime], + end_time: Optional[datetime], + ) -> JobList: + return await list_jobs_for_user( + db, + user_id=user_id, + page=page, + page_size=page_size, + job_status=job_status, + job_type=job_type, + recent_days=recent_days, + start_time=start_time, + end_time=end_time, + ) + + async def get_job_result_for_user( + self, + db: AsyncSession, + *, + job_id: str, + user_id: str, + ) -> JobResultResponse: + return await get_job_result_for_user(db, job_id=job_id, user_id=user_id) + + +def check_job_permission(job, user_id: str, job_id: str) -> None: + if not job: + raise NotFoundException( + resource="Job", resource_id=job_id, internal_message="Job not found" + ) + + if str(job.user_id) != user_id: + raise PermissionDeniedException( + user_message="You don't have permission to access this job", + ) + + +def normalize_naive_utc_filter_datetime(dt: Optional[datetime]) -> Optional[datetime]: + if dt is None: + return None + if dt.tzinfo is None or dt.utcoffset() is None: + return dt + return dt.astimezone(timezone.utc).replace(tzinfo=None) + + +async def list_jobs_for_user( + db: AsyncSession, + *, + user_id: str, + page: int, + page_size: int, + job_status: Optional[str], + job_type: Optional[str], + recent_days: Optional[int], + start_time: Optional[datetime], + end_time: Optional[datetime], +) -> JobList: + try: + job_repo = JobRepository() + + if recent_days not in (None, 1, 7, 30): + raise ValidationException( + user_message="recent_days only supports 1, 7, or 30", + violations=[{"field": "recent_days", "description": "Invalid value"}], + ) + + created_after: Optional[datetime] = None + if recent_days: + created_after = utc_now_naive() - timedelta(days=recent_days) + + normalized_start_time = normalize_naive_utc_filter_datetime(start_time) + normalized_end_time = normalize_naive_utc_filter_datetime(end_time) + + if ( + normalized_start_time + and normalized_end_time + and normalized_start_time > normalized_end_time + ): + raise ValidationException( + user_message="start_time cannot be later than end_time", + violations=[ + {"field": "start_time", "description": "Must be before end_time"} + ], + ) + + if normalized_start_time: + created_after = normalized_start_time + created_before = normalized_end_time + + total_count = await job_repo.count_jobs_by_user( + db=db, + user_id=user_id, + created_after=created_after, + created_before=created_before, + job_type=job_type, + job_status=job_status, + ) + jobs = await job_repo.get_jobs_by_user( + db=db, + user_id=user_id, + limit=page_size, + offset=(page - 1) * page_size, + created_after=created_after, + created_before=created_before, + job_type=job_type, + job_status=job_status, + ) + + redis_service = RedisServiceFactory.get_service() + job_responses = [] + for job in jobs: + job_metadata = await job_repo.get_job_metadata( + db, job.job_id, redis_service + ) + job_responses.append( + await build_job_result_response( + job=job, + job_metadata=job_metadata, + progress=None, + ) + ) + + total_pages = math.ceil(total_count / page_size) if total_count > 0 else 0 + return JobList( + jobs=job_responses, + total=total_count, + page=page, + page_size=page_size, + total_pages=total_pages, + ) + + except ValidationException: + raise + except Exception as exc: + logger.error(f"Failed to list jobs: {exc}") + raise JobOperationException( + internal_message=f"Failed to get job list: {str(exc)}" + ) + + +async def get_job_result_for_user( + db: AsyncSession, + *, + job_id: str, + user_id: str, +) -> JobResultResponse: + try: + job_repo = JobRepository() + job = await job_repo.get_job_by_id(db, job_id) + check_job_permission(job, user_id, job_id) + assert job is not None + + progress = None + if to_job_status_value(job.status) == "running": + progress = {"total_pages": 10, "processed_pages": 5} + + redis_service = RedisServiceFactory.get_service() + job_metadata = await job_repo.get_job_metadata(db, job_id, redis_service) + return await build_job_result_response( + job=job, + job_metadata=job_metadata, + progress=progress, + ) + + except NotFoundException: + raise + except PermissionDeniedException: + raise + except Exception as exc: + logger.error(f"Failed to get job result: {exc}") + raise JobOperationException( + internal_message=f"Failed to get job result: {str(exc)}" + ) diff --git a/apps/api/app/services/jobs/read_service.py b/apps/api/app/services/jobs/read_service.py new file mode 100644 index 000000000..c055dd2e5 --- /dev/null +++ b/apps/api/app/services/jobs/read_service.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +from datetime import datetime +from typing import Optional + +from app.services.jobs.job_read_model import JobReadModel +from app.services.jobs.job_read_model import check_job_permission +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.schemas.job import JobList, JobResultResponse + +__all__ = [ + "check_job_permission", + "get_job_result_for_user", + "list_jobs_for_user", +] + + +async def list_jobs_for_user( + db: AsyncSession, + *, + user_id: str, + page: int, + page_size: int, + job_status: Optional[str], + job_type: Optional[str], + recent_days: Optional[int], + start_time: Optional[datetime], + end_time: Optional[datetime], +) -> JobList: + return await JobReadModel().list_jobs_for_user( + db, + user_id=user_id, + page=page, + page_size=page_size, + job_status=job_status, + job_type=job_type, + recent_days=recent_days, + start_time=start_time, + end_time=end_time, + ) + + +async def get_job_result_for_user( + db: AsyncSession, + *, + job_id: str, + user_id: str, +) -> JobResultResponse: + return await JobReadModel().get_job_result_for_user( + db, + job_id=job_id, + user_id=user_id, + ) diff --git a/apps/api/app/services/jobs/result_projection.py b/apps/api/app/services/jobs/result_projection.py new file mode 100644 index 000000000..0a00ce2d2 --- /dev/null +++ b/apps/api/app/services/jobs/result_projection.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +import os +from datetime import datetime, timezone +from typing import Any, Literal, Optional, cast +from urllib.parse import urlparse + +from shared.core.billing import MicroDollar +from shared.core.exceptions.domain_exceptions import JobOperationException +from shared.models.schemas.job import JobResultResponse, StandardErrorObject +from shared.models.schemas.job_metadata import JobMetadataHelper +from shared.services.jobs.result_delivery import JobResultDeliveryResolver +from shared.utils.error_details import normalize_error_details + +JobStatusValue = Literal[ + "pending", "waiting-file", "running", "converting", "done", "failed" +] + + +def build_error_response( + job: Any, job_metadata: Optional[dict] = None +) -> Optional[StandardErrorObject]: + if not job.error_message: + return None + + error_details = None + if job_metadata and isinstance(job_metadata, dict): + error_details = normalize_error_details(job_metadata.get("error_details")) + + return StandardErrorObject( + code=job.error_code or "UNKNOWN", + message=job.error_message, + request_id=job.job_id, + details=error_details, + ) + + +def resolve_public_document_id(job: Any) -> Optional[str]: + job_result = getattr(job, "job_result", None) + published_document_id = getattr(job_result, "document_id", None) + if isinstance(published_document_id, str) and published_document_id: + return published_document_id + + return None + + +def ensure_utc(dt: Optional[datetime]) -> Optional[datetime]: + if not dt: + return None + if dt.tzinfo: + return dt.astimezone(timezone.utc) + return dt.replace(tzinfo=timezone.utc) + + +def require_utc(dt: Optional[datetime], *, field_name: str) -> datetime: + normalized_dt = ensure_utc(dt) + if normalized_dt is None: + raise JobOperationException( + internal_message=f"Job is missing required datetime field: {field_name}" + ) + return normalized_dt + + +def to_job_status_value(status: str) -> JobStatusValue: + return cast(JobStatusValue, status) + + +def _resolve_original_request(job_metadata: Optional[dict[str, Any]]) -> dict[str, Any]: + return JobMetadataHelper.get_original_request(job_metadata) + + +def _resolve_source_file_name(original_request: dict[str, Any]) -> str | None: + source_url = original_request.get("source_url") + file_name = None + if source_url: + parsed_source = urlparse(str(source_url)) + file_name = os.path.basename(parsed_source.path) or None + if not file_name: + file_name = original_request.get("file_name") + return str(file_name) if file_name else None + + +def _resolve_file_extension(file_name: str | None) -> str | None: + if not file_name: + return None + extension = os.path.splitext(file_name)[1] + return extension[1:].upper() if extension else None + + +def _resolve_parsing_params( + job_metadata: Optional[dict[str, Any]], + original_request: dict[str, Any], +) -> dict[str, Any]: + parsing_params = original_request.get("parsing_params") or {} + if not parsing_params: + parsing_params = JobMetadataHelper.get_parsing_params_dict(job_metadata) + return parsing_params if isinstance(parsing_params, dict) else {} + + +def _resolve_duration_seconds(job: Any) -> float | None: + if job.updated_at and job.created_at: + return (job.updated_at - job.created_at).total_seconds() + return None + + +async def _resolve_result_delivery( + job: Any, +) -> tuple[dict[str, Any] | None, str | None, datetime]: + default_expires_at = require_utc( + job.created_at, + field_name="created_at", + ) + delivery = JobResultDeliveryResolver().resolve( + job.job_result, + default_expires_at=default_expires_at, + ) + return ( + delivery.result, + delivery.result_url, + delivery.result_url_expires_at or default_expires_at, + ) + + +async def build_job_result_response( + *, + job: Any, + job_metadata: Optional[dict[str, Any]], + progress: dict[str, Any] | None, +) -> JobResultResponse: + original_request = _resolve_original_request(job_metadata) + file_name = _resolve_source_file_name(original_request) + parsing_params = _resolve_parsing_params(job_metadata, original_request) + result, result_url, result_url_expires_at = await _resolve_result_delivery(job) + + return JobResultResponse( + job_id=job.job_id, + namespace=JobMetadataHelper.get_namespace(job_metadata), + document_id=resolve_public_document_id(job), + status=to_job_status_value(job.status), + source_type=job.source_type, + data_id=JobMetadataHelper.get_data_id(job_metadata), + created_at=require_utc(job.created_at, field_name="created_at"), + progress=progress, + error=build_error_response(job, job_metadata), + result=result, + result_url=result_url, + result_url_expires_at=require_utc( + result_url_expires_at, + field_name="result_url_expires_at", + ), + file_name=file_name, + file_extension=_resolve_file_extension(file_name), + model=parsing_params.get("model"), + ocr_enabled=parsing_params.get("ocr_enabled"), + duration_seconds=_resolve_duration_seconds(job), + credits_spent=( + MicroDollar(job.credits_charged).to_credit() + if hasattr(job, "credits_charged") + else 0 + ), + ) diff --git a/apps/api/app/services/knowledge/__init__.py b/apps/api/app/services/knowledge/__init__.py deleted file mode 100644 index f95391f1e..000000000 --- a/apps/api/app/services/knowledge/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -""" -Knowledge-base services. -""" - -from .kb_orchestrator import KBOrchestrator - -__all__ = ["KBOrchestrator"] diff --git a/apps/api/app/services/knowledge/kb_orchestrator.py b/apps/api/app/services/knowledge/kb_orchestrator.py deleted file mode 100644 index b71b9af20..000000000 --- a/apps/api/app/services/knowledge/kb_orchestrator.py +++ /dev/null @@ -1,187 +0,0 @@ -""" -Knowledge-base workflow orchestration. -""" - -from typing import Optional - -from loguru import logger - -from shared.core.celery_router import task_router - -# Tasks now run in the Worker service and are referenced by task name. -from shared.core.exceptions.domain_exceptions import ( - KnowhereException, - WorkerHandlingException, -) - - -class KBOrchestrator: - """Coordinate knowledge-base processing jobs.""" - - def __init__(self): - self.task_router = task_router - - async def start_workflow( - self, - db, - job_id: str, - source_type: str, - file_path: Optional[str] = None, - file_url: Optional[str] = None, - user_id: Optional[str] = None, - ) -> str: - """ - Start the knowledge-base workflow. - - Args: - db: Database session. - job_id: Job identifier. - source_type: Source type. - file_path: Uploaded file path when direct upload is used. - file_url: Source URL when URL ingestion is used. - user_id: User identifier. - - Returns: - str: Celery task identifier. - """ - try: - # When the source is a URL, recover file_url from job metadata if needed. - if source_type == "url" and not file_url: - from app.repositories.job_repository import JobRepository - - from shared.models.schemas.job_metadata import JobMetadataHelper - from shared.services.redis import RedisServiceFactory - - job_repo = JobRepository() - redis_service = RedisServiceFactory.get_service() - job_metadata = await job_repo.get_job_metadata( - db, job_id, redis_service - ) - file_url = JobMetadataHelper.get_field(job_metadata, "file_url") - - # Resolve the queue name for this job. - effective_user_id = user_id or "" - queue_name = self.task_router.get_queue_for_job( - "kb_management", effective_user_id - ) - task_kwargs = { - "user_id": effective_user_id, - "job_type": "kb_management", - } - - # Start the single worker task. Upload is already complete via S3. - # The task handles parsing, vectorization, ZIP generation, S3 upload, - # and result publication. Webhook and email delivery stay in the API. - from celery import signature - - task_signature = signature( - "app.core.tasks.kb_tasks.parse_task", - args=[job_id], - kwargs=task_kwargs, - ) - if task_signature is None: - raise WorkerHandlingException( - internal_message=( - "Failed to build knowledge-base workflow: missing Celery signature" - ) - ) - task_signature = task_signature.set(queue=queue_name) - if task_signature is None: - raise WorkerHandlingException( - internal_message=( - "Failed to configure knowledge-base workflow queue" - ) - ) - - # Enqueue the task. - result = task_signature.apply_async() - if result is None or result.id is None: - raise WorkerHandlingException( - internal_message=( - "Failed to start knowledge-base workflow: missing Celery task id" - ) - ) - - logger.info( - f"Knowledge-base workflow started: job_id={job_id}, task_id={result.id}, queue={queue_name}" - ) - - return result.id - - except KnowhereException: - raise - except Exception as e: - logger.error(f"Failed to start knowledge-base workflow: {e}") - raise WorkerHandlingException( - internal_message=f"Failed to start knowledge-base workflow: {str(e)}", - original_exception=e, - ) - - def create_workflow_chain( - self, job_id: str, user_id: str, queue_name: Optional[str] = None - ): - """ - Create the workflow signature for tests or manual execution. - - Args: - job_id: Job identifier. - user_id: User identifier. - queue_name: Optional explicit queue name. - - Returns: - signature: Celery task signature. - """ - if not queue_name: - queue_name = self.task_router.get_queue_for_job("kb_management", user_id) - task_kwargs = { - "user_id": user_id, - "job_type": "kb_management", - } - - from celery import signature - - # Return the single-task signature. Parsing, vectorization, ZIP generation, - # S3 upload, and publication happen in the worker. Webhook and email - # delivery remain in the API service. - task_signature = signature( - "app.core.tasks.kb_tasks.parse_task", - args=[job_id], - kwargs=task_kwargs, - ) - if task_signature is None: - raise WorkerHandlingException( - internal_message=( - "Failed to build knowledge-base workflow: missing Celery signature" - ) - ) - configured_signature = task_signature.set(queue=queue_name) - if configured_signature is None: - raise WorkerHandlingException( - internal_message="Failed to configure knowledge-base workflow queue" - ) - return configured_signature - - def cancel_workflow(self, workflow_id: str) -> bool: - """ - Cancel a workflow. - - Args: - workflow_id: Workflow identifier. - - Returns: - bool: Whether cancellation succeeded. - """ - try: - from shared.core.celery_app import get_celery_app - - celery_app = get_celery_app() - - result = celery_app.AsyncResult(workflow_id) - result.revoke(terminate=True) - - logger.info(f"Workflow cancelled: {workflow_id}") - return True - - except Exception as e: - logger.error(f"Failed to cancel workflow: {e}") - return False diff --git a/apps/api/app/services/rate_limit/data_structures.py b/apps/api/app/services/rate_limit/data_structures.py index 0467c79aa..af206917b 100644 --- a/apps/api/app/services/rate_limit/data_structures.py +++ b/apps/api/app/services/rate_limit/data_structures.py @@ -15,6 +15,15 @@ class CurrentUser: user_tier: str +@dataclass(frozen=True) +class RouteAdmissionContext: + """HTTP route facts needed by the Job Admission workflow.""" + + method: str + path: str + limit_identifier: str + + @dataclass(frozen=True) class TierLimits: """ diff --git a/apps/api/app/services/rate_limit/dependencies.py b/apps/api/app/services/rate_limit/dependencies.py deleted file mode 100644 index c9a697b74..000000000 --- a/apps/api/app/services/rate_limit/dependencies.py +++ /dev/null @@ -1,423 +0,0 @@ -""" -FastAPI dependencies for the rate-limit layer. - -Dependency chain (outermost -> innermost): - require_billing_limits -> with_current_user -> get_current_user_id - -> get_db - -``with_current_user`` resolves the user's billing tier through TierService and -enforces the matched system limit (Layer 0). - -``require_billing_limits`` enforces billing RPM (Layer 1) when billing is -enabled and yields control to the route handler. Concurrency (Layer 2) and -daily quota (Layer 3) are enforced just before insert in the create-job route -only when billing is enabled. -""" - -import math -from fnmatch import fnmatch -from typing import AsyncGenerator - -from app.core.dependencies import get_current_user_id -from app.services.rate_limit.config import ( - CONCURRENCY_RETRY_AFTER_SECONDS, - RateLimitConfig, -) -from app.services.rate_limit.data_structures import CurrentUser, TierLimits -from app.services.rate_limit.limiter import RateLimiter -from app.services.rate_limit.system_limit import find_system_rule -from app.services.rate_limit.tier_service import TierService -from fastapi import Depends, Request -from loguru import logger -from sqlalchemy import func, select -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.config import settings -from shared.core.database import get_db -from shared.core.exceptions.domain_exceptions import ( - PermissionDeniedException, - RateLimitException, - UnavailableException, -) -from shared.core.logging import log_context -from shared.core.state_machine.states import JobStatus -from shared.models.database.job import Job -from shared.models.database.user_balance import UserBalance - -_ACTIVE_JOB_STATES: tuple[str, ...] = ( - JobStatus.WAITING_FILE.value, - JobStatus.PENDING.value, - JobStatus.RUNNING.value, - JobStatus.CONVERTING.value, -) -_GUEST_API_KEY_ALLOWED_ROUTE_PATTERNS: tuple[str, ...] = ( - "/v1/jobs", - "/v1/jobs/*", - "/v1/billing/credits", - "/v1/retrieval/query", - "/v1/documents", - "/v1/documents/*", - "/mcp", -) -_GUEST_API_KEY_REQUIRED_PERMISSION: str = ( - "jobs_documents_retrieval_mcp_or_billing_credits" -) -_GUEST_API_KEY_SCOPE_MESSAGE: str = ( - "Guest API keys can only access job, document, retrieval, MCP query, " - "and billing credits APIs" -) - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _get_route_path(request: Request) -> str: - """Return the request path without the application's root_path prefix.""" - scope_path: str = request.scope.get("path", request.url.path) - root_path: str = request.scope.get("root_path", "") - if root_path and scope_path.startswith(root_path): - return scope_path[len(root_path) :] - return scope_path - - -def _get_route_limit_identifier(request: Request) -> str: - """Return a stable identifier for route-scoped system limits.""" - route = request.scope.get("route") - route_path = getattr(route, "path", None) - if isinstance(route_path, str) and route_path: - return route_path - - route_path_format = getattr(route, "path_format", None) - if isinstance(route_path_format, str) and route_path_format: - return route_path_format - - return _get_route_path(request) - - -def _normalize_route_path(route_path: str) -> str: - """Normalize guest route checks across slash-redirect variants.""" - normalized_path = route_path.rstrip("/") - return normalized_path or "/" - - -def _is_guest_api_key_route_allowed(route_path: str) -> bool: - """Return whether a guest API key may access the given route.""" - normalized_path = _normalize_route_path(route_path) - return any( - fnmatch(normalized_path, pattern) - for pattern in _GUEST_API_KEY_ALLOWED_ROUTE_PATTERNS - ) - - -def _enforce_guest_api_key_scope(request: Request, user_tier: str) -> None: - """Reject guest API keys outside the guest-allowed API surface.""" - if user_tier != "guest": - return - - route_path = _get_route_path(request) - if _is_guest_api_key_route_allowed(route_path): - return - - raise PermissionDeniedException( - user_message=_GUEST_API_KEY_SCOPE_MESSAGE, - required_permission=_GUEST_API_KEY_REQUIRED_PERMISSION, - ) - - -# --------------------------------------------------------------------------- -# with_current_user -- Layer 0 (matched system limit) -# --------------------------------------------------------------------------- - - -async def with_current_user( - request: Request, - user_id: str = Depends(get_current_user_id), -) -> AsyncGenerator[CurrentUser, None]: - """Resolve the current user tier and enforce the matched system limit. - - Steps: - 1. ``get_current_user_id`` already authenticated the user (401 - on failure). - 2. Resolve ``user_tier`` through ``TierService.get_tier(user_id)``. - 3. If ``RATE_LIMIT_ENABLED=false`` is set, return immediately. - 4. Check the matched system limit via the rate limiter (fail-open on - Redis error). - """ - user_tier = await TierService.get_tier(user_id) - _enforce_guest_api_key_scope(request, user_tier) - current_user = CurrentUser(user_id=user_id, user_tier=user_tier) - - with log_context(user_id=user_id): - # -- Global rate-limit switch -- - config = RateLimitConfig.get_instance() - if not config.is_enabled: - yield current_user - return - - # -- Layer 0: matched system limit -- - try: - route_path = _get_route_path(request) - rule = find_system_rule(request.method, route_path, config.system_rules) - limiter = RateLimiter(config) - await limiter.check_system_limit( - identifier=user_id, - limit=rule.limit, - matched_pattern=rule.api_pattern, - period=rule.period, - ) - except RateLimitException: - raise - except Exception as exc: - # Fail-open: log and let the request through. - logger.warning( - "rate_limit: Redis error during system limit check, " - "failing open for user_id={}, error={}", - user_id, - exc, - ) - - yield current_user - - -# --------------------------------------------------------------------------- -# require_billing_limits -- Layer 1 -# (billing RPM) -# --------------------------------------------------------------------------- - -_RETRY_AFTER_SECONDS: int = 15 - - -async def require_billing_limits( - request: Request, - current_user: CurrentUser = Depends(with_current_user), - _db: AsyncSession = Depends(get_db), -) -> AsyncGenerator[CurrentUser, None]: - """Enforce billing RPM (Layer 1) around the route handler. - - This is an async-generator (yield) dependency so that teardown logic - can run after the route handler completes. - - Layer enforced before yield: - 1. Billing RPM -- per-user requests-per-minute - - Layers enforced inside route just before insert: - 2. Non-terminal jobs concurrency -- max pending/running jobs - 3. Daily quota (free tier only) -- hard daily cap - - When ``BILLING_ENABLED=false``, this yields after identity and system - route limiting. Otherwise, Redis failures raise 503 because billing - enforcement must not be silently skipped. - """ - if not settings.BILLING_ENABLED: - yield current_user - return - - config = RateLimitConfig.get_instance() - if not config.is_enabled: - yield current_user - return - - tier_limits: TierLimits | None = config.tier_map.get(current_user.user_tier) - if tier_limits is None: - logger.error( - "rate_limit: no tier config for tier='{}', user_id={}", - current_user.user_tier, - current_user.user_id, - ) - raise UnavailableException( - internal_message=(f"Missing tier config for tier={current_user.user_tier}"), - retry_after=_RETRY_AFTER_SECONDS, - ) - - # -- Layer 1: billing RPM -- - limiter = RateLimiter(config) - try: - await limiter.check_billing_rpm(current_user.user_id, tier_limits.rpm_limit) - except RateLimitException: - raise - except Exception as exc: - raise UnavailableException( - internal_message=(f"Redis error in billing RPM check: {exc}"), - retry_after=_RETRY_AFTER_SECONDS, - ) - - yield current_user - - -async def require_route_system_limit(request: Request) -> None: - """Apply the matched system limit to the current route using a route key. - - Prefer the framework route template so paths with different parameters - share the same budget bucket. If no explicit rule matches, the default - system rule still protects the route with the wider fallback budget. - Fail closed when Redis or limiter state is unavailable. - """ - config = RateLimitConfig.get_instance() - if not config.is_enabled: - return - - route_path = _get_route_path(request) - route_identifier = _get_route_limit_identifier(request) - rule = find_system_rule(request.method, route_path, config.system_rules) - limiter = RateLimiter(config) - try: - await limiter.check_system_limit( - identifier=route_identifier, - limit=rule.limit, - matched_pattern=rule.api_pattern, - period=rule.period, - use_global_key=True, - ) - except RateLimitException: - raise - except Exception as exc: - raise UnavailableException( - internal_message=(f"Redis error in route system limit: {exc}"), - retry_after=_RETRY_AFTER_SECONDS, - limit=rule.limit, - period=rule.period, - ) - - -async def enforce_job_creation_capacity( - request: Request, - db: AsyncSession, - current_user: CurrentUser, -) -> None: - """Enforce Layers 2-3 immediately before job insert.""" - if not settings.BILLING_ENABLED: - return - - config = RateLimitConfig.get_instance() - if not config.is_enabled: - return - - tier_limits = config.tier_map.get(current_user.user_tier) - if tier_limits is None: - logger.error( - "rate_limit: no tier config for tier='{}', user_id={}", - current_user.user_tier, - current_user.user_id, - ) - raise UnavailableException( - internal_message=(f"Missing tier config for tier={current_user.user_tier}"), - retry_after=_RETRY_AFTER_SECONDS, - ) - - limiter = RateLimiter(config) - - # -- Layer 2: non-terminal jobs concurrency (DB-locked) -- - if tier_limits.max_concurrent_jobs != -1: - try: - await _acquire_user_concurrency_lock(db, current_user.user_id) - active_jobs = await _count_non_terminal_jobs(db, current_user.user_id) - if active_jobs >= tier_limits.max_concurrent_jobs: - retry_after_seconds = _compute_concurrency_retry_after_seconds( - base_retry_after_seconds=CONCURRENCY_RETRY_AFTER_SECONDS, - rpm_limit=tier_limits.rpm_limit, - ) - exc = RateLimitException( - retry_after=retry_after_seconds, - limit=tier_limits.max_concurrent_jobs, - period="concurrent", - user_message=( - f"Too many concurrent requests " - f"({active_jobs}/{tier_limits.max_concurrent_jobs} active). " - f"Please retry after {retry_after_seconds} seconds." - ), - internal_message=( - "Concurrency limit exceeded: " - f"user_id={current_user.user_id}, " - f"active_jobs={active_jobs}, " - f"limit={tier_limits.max_concurrent_jobs}, " - f"retry_after={retry_after_seconds}s" - ), - ) - exc.details.update( - { - "active_jobs": active_jobs, - "available_slots": max( - 0, tier_limits.max_concurrent_jobs - active_jobs - ), - } - ) - raise exc - except RateLimitException: - raise - except Exception as exc: - raise UnavailableException( - internal_message=(f"DB error in concurrency check: {exc}"), - retry_after=_RETRY_AFTER_SECONDS, - ) - - # -- Layer 3: daily quota -- - if tier_limits.daily_quota != -1: - try: - await limiter.check_daily_quota( - current_user.user_id, tier_limits.daily_quota - ) - except RateLimitException: - raise - except Exception as exc: - raise UnavailableException( - internal_message=(f"Redis error in daily quota check: {exc}"), - retry_after=_RETRY_AFTER_SECONDS, - ) - - -# --------------------------------------------------------------------------- -# Internal helpers -# --------------------------------------------------------------------------- - - -async def _acquire_user_concurrency_lock( - db: AsyncSession, - user_id: str, -) -> None: - """Acquire a per-user row lock to serialize concurrent job creation. - - Locks the UserBalance row instead of User to avoid contention with - unrelated operations (profile updates, etc.) that may also lock User. - """ - result = await db.execute( - select(UserBalance.user_id) - .where(UserBalance.user_id == user_id) - .with_for_update() - ) - if result.scalar_one_or_none() is None: - raise RateLimitException( - internal_message=f"UserBalance row not found for user_id={user_id}" - ) - - -def _compute_concurrency_retry_after_seconds( - base_retry_after_seconds: int, - rpm_limit: int, -) -> int: - """ - Compute Retry-After hint for concurrency rejections. - - Concurrency has no deterministic reset timestamp, so we provide a - conservative client hint: - - floor: configured base retry (currently 30s) - - if billing RPM is finite, also respect one request spacing - (ceil(60 / rpm_limit)) to reduce immediate repeated 429s - """ - if rpm_limit <= 0: - return base_retry_after_seconds - return max(base_retry_after_seconds, int(math.ceil(60 / rpm_limit))) - - -async def _count_non_terminal_jobs( - db: AsyncSession, - user_id: str, -) -> int: - """Count non-terminal jobs for a user in the current transaction.""" - result = await db.execute( - select(func.count(Job.job_id)) - .where(Job.user_id == user_id) - .where(Job.status.in_(_ACTIVE_JOB_STATES)) - ) - return int(result.scalar_one() or 0) diff --git a/apps/api/app/services/rate_limit/job_admission_capacity_service.py b/apps/api/app/services/rate_limit/job_admission_capacity_service.py new file mode 100644 index 000000000..fb29bb7ad --- /dev/null +++ b/apps/api/app/services/rate_limit/job_admission_capacity_service.py @@ -0,0 +1,206 @@ +from __future__ import annotations + +import math + +from app.services.rate_limit.config import ( + CONCURRENCY_RETRY_AFTER_SECONDS, + RateLimitConfig, +) +from app.services.rate_limit.data_structures import CurrentUser, TierLimits +from app.services.rate_limit.limiter import RateLimiter +from loguru import logger +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ( + RateLimitException, + UnavailableException, +) +from shared.core.state_machine.states import JobStatus +from shared.models.database.job import Job +from shared.models.database.user_balance import UserBalance + +_ACTIVE_JOB_STATES: tuple[str, ...] = ( + JobStatus.WAITING_FILE.value, + JobStatus.PENDING.value, + JobStatus.RUNNING.value, + JobStatus.CONVERTING.value, +) +_RETRY_AFTER_SECONDS: int = 15 + + +class JobAdmissionCapacityService: + async def enforce_billing_limits( + self, + *, + current_user: CurrentUser, + ) -> None: + if not settings.BILLING_ENABLED: + return + + config = RateLimitConfig.get_instance() + if not config.is_enabled: + return + + tier_limits = self._require_tier_limits( + config=config, + current_user=current_user, + ) + limiter = RateLimiter(config) + + try: + await limiter.check_billing_rpm( + current_user.user_id, + tier_limits.rpm_limit, + ) + except RateLimitException: + raise + except Exception as exc: + raise UnavailableException( + internal_message=(f"Redis error in billing RPM check: {exc}"), + retry_after=_RETRY_AFTER_SECONDS, + ) + + async def enforce_job_creation_capacity( + self, + *, + db: AsyncSession, + current_user: CurrentUser, + ) -> None: + if not settings.BILLING_ENABLED: + return + + config = RateLimitConfig.get_instance() + if not config.is_enabled: + return + + tier_limits = self._require_tier_limits( + config=config, + current_user=current_user, + ) + limiter = RateLimiter(config) + + if tier_limits.max_concurrent_jobs != -1: + try: + await self._acquire_user_concurrency_lock( + db=db, + user_id=current_user.user_id, + ) + active_jobs = await self._count_non_terminal_jobs( + db=db, + user_id=current_user.user_id, + ) + if active_jobs >= tier_limits.max_concurrent_jobs: + retry_after_seconds = self._compute_concurrency_retry_after_seconds( + base_retry_after_seconds=CONCURRENCY_RETRY_AFTER_SECONDS, + rpm_limit=tier_limits.rpm_limit, + ) + exc = RateLimitException( + retry_after=retry_after_seconds, + limit=tier_limits.max_concurrent_jobs, + period="concurrent", + user_message=( + f"Too many concurrent requests " + f"({active_jobs}/{tier_limits.max_concurrent_jobs} active). " + f"Please retry after {retry_after_seconds} seconds." + ), + internal_message=( + "Concurrency limit exceeded: " + f"user_id={current_user.user_id}, " + f"active_jobs={active_jobs}, " + f"limit={tier_limits.max_concurrent_jobs}, " + f"retry_after={retry_after_seconds}s" + ), + ) + exc.details.update( + { + "active_jobs": active_jobs, + "available_slots": max( + 0, + tier_limits.max_concurrent_jobs - active_jobs, + ), + } + ) + raise exc + except RateLimitException: + raise + except Exception as exc: + raise UnavailableException( + internal_message=(f"DB error in concurrency check: {exc}"), + retry_after=_RETRY_AFTER_SECONDS, + ) + + if tier_limits.daily_quota != -1: + try: + await limiter.check_daily_quota( + current_user.user_id, + tier_limits.daily_quota, + ) + except RateLimitException: + raise + except Exception as exc: + raise UnavailableException( + internal_message=(f"Redis error in daily quota check: {exc}"), + retry_after=_RETRY_AFTER_SECONDS, + ) + + def _require_tier_limits( + self, + *, + config: RateLimitConfig, + current_user: CurrentUser, + ) -> TierLimits: + tier_limits = config.tier_map.get(current_user.user_tier) + if tier_limits is None: + logger.error( + "rate_limit: no tier config for tier='{}', user_id={}", + current_user.user_tier, + current_user.user_id, + ) + raise UnavailableException( + internal_message=( + f"Missing tier config for tier={current_user.user_tier}" + ), + retry_after=_RETRY_AFTER_SECONDS, + ) + return tier_limits + + async def _acquire_user_concurrency_lock( + self, + *, + db: AsyncSession, + user_id: str, + ) -> None: + result = await db.execute( + select(UserBalance.user_id) + .where(UserBalance.user_id == user_id) + .with_for_update() + ) + if result.scalar_one_or_none() is None: + raise RateLimitException( + internal_message=f"UserBalance row not found for user_id={user_id}" + ) + + async def _count_non_terminal_jobs( + self, + *, + db: AsyncSession, + user_id: str, + ) -> int: + result = await db.execute( + select(func.count(Job.job_id)) + .where(Job.user_id == user_id) + .where(Job.status.in_(_ACTIVE_JOB_STATES)) + ) + return int(result.scalar_one() or 0) + + def _compute_concurrency_retry_after_seconds( + self, + *, + base_retry_after_seconds: int, + rpm_limit: int, + ) -> int: + if rpm_limit <= 0: + return base_retry_after_seconds + return max(base_retry_after_seconds, int(math.ceil(60 / rpm_limit))) diff --git a/apps/api/app/services/rate_limit/job_admission_route_policy_service.py b/apps/api/app/services/rate_limit/job_admission_route_policy_service.py new file mode 100644 index 000000000..ce89931f6 --- /dev/null +++ b/apps/api/app/services/rate_limit/job_admission_route_policy_service.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +from fnmatch import fnmatch + +from app.services.rate_limit.config import RateLimitConfig +from app.services.rate_limit.data_structures import RouteAdmissionContext +from app.services.rate_limit.limiter import RateLimiter +from app.services.rate_limit.system_limit import find_system_rule + +from shared.core.exceptions.domain_exceptions import ( + PermissionDeniedException, + RateLimitException, + UnavailableException, +) + +_GUEST_API_KEY_ALLOWED_ROUTE_PATTERNS: tuple[str, ...] = ( + "/v1/jobs", + "/v1/jobs/*", + "/v1/billing/credits", + "/v1/retrieval/query", + "/v1/documents", + "/v1/documents/*", + "/mcp", +) +_GUEST_API_KEY_REQUIRED_PERMISSION: str = ( + "jobs_documents_retrieval_mcp_or_billing_credits" +) +_GUEST_API_KEY_SCOPE_MESSAGE: str = ( + "Guest API keys can only access job, document, retrieval, MCP query, " + "and billing credits APIs" +) +_RETRY_AFTER_SECONDS: int = 15 + + +class JobAdmissionRoutePolicyService: + async def enforce_user_system_limit( + self, + *, + route_context: RouteAdmissionContext, + config: RateLimitConfig, + user_id: str, + ) -> None: + rule = find_system_rule( + route_context.method, + route_context.path, + config.system_rules, + ) + limiter = RateLimiter(config) + await limiter.check_system_limit( + identifier=user_id, + limit=rule.limit, + matched_pattern=rule.api_pattern, + period=rule.period, + ) + + async def enforce_route_system_limit( + self, + *, + route_context: RouteAdmissionContext, + ) -> None: + config = RateLimitConfig.get_instance() + if not config.is_enabled: + return + + route_identifier = route_context.limit_identifier + rule = find_system_rule( + route_context.method, + route_context.path, + config.system_rules, + ) + limiter = RateLimiter(config) + + try: + await limiter.check_system_limit( + identifier=route_identifier, + limit=rule.limit, + matched_pattern=rule.api_pattern, + period=rule.period, + use_global_key=True, + ) + except RateLimitException: + raise + except Exception as exc: + raise UnavailableException( + internal_message=(f"Redis error in route system limit: {exc}"), + retry_after=_RETRY_AFTER_SECONDS, + limit=rule.limit, + period=rule.period, + ) + + def enforce_guest_api_key_scope( + self, + *, + route_context: RouteAdmissionContext, + user_tier: str, + ) -> None: + if user_tier != "guest": + return + + if self._is_guest_api_key_route_allowed(route_context.path): + return + + raise PermissionDeniedException( + user_message=_GUEST_API_KEY_SCOPE_MESSAGE, + required_permission=_GUEST_API_KEY_REQUIRED_PERMISSION, + ) + + def _normalize_route_path(self, route_path: str) -> str: + normalized_path = route_path.rstrip("/") + return normalized_path or "/" + + def _is_guest_api_key_route_allowed(self, route_path: str) -> bool: + normalized_path = self._normalize_route_path(route_path) + return any( + fnmatch(normalized_path, pattern) + for pattern in _GUEST_API_KEY_ALLOWED_ROUTE_PATTERNS + ) diff --git a/apps/api/app/services/rate_limit/job_admission_service.py b/apps/api/app/services/rate_limit/job_admission_service.py new file mode 100644 index 000000000..130a880c3 --- /dev/null +++ b/apps/api/app/services/rate_limit/job_admission_service.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from app.services.rate_limit.config import RateLimitConfig +from app.services.rate_limit.data_structures import CurrentUser, RouteAdmissionContext +from app.services.rate_limit.job_admission_capacity_service import ( + JobAdmissionCapacityService, +) +from app.services.rate_limit.job_admission_route_policy_service import ( + JobAdmissionRoutePolicyService, +) +from app.services.rate_limit.tier_service import TierService +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import RateLimitException +from shared.core.logging import log_context + + +class JobAdmissionService: + def __init__( + self, + *, + route_policy_service: JobAdmissionRoutePolicyService | None = None, + capacity_service: JobAdmissionCapacityService | None = None, + ) -> None: + self._route_policy_service = ( + route_policy_service or JobAdmissionRoutePolicyService() + ) + self._capacity_service = capacity_service or JobAdmissionCapacityService() + + async def resolve_current_user( + self, + *, + route_context: RouteAdmissionContext, + user_id: str, + ) -> CurrentUser: + user_tier = await TierService.get_tier(user_id) + self._route_policy_service.enforce_guest_api_key_scope( + route_context=route_context, + user_tier=user_tier, + ) + current_user = CurrentUser(user_id=user_id, user_tier=user_tier) + + with log_context(user_id=user_id): + config = RateLimitConfig.get_instance() + if not config.is_enabled: + return current_user + + try: + await self._route_policy_service.enforce_user_system_limit( + route_context=route_context, + config=config, + user_id=user_id, + ) + except RateLimitException: + raise + except Exception as exc: + logger.warning( + "rate_limit: Redis error during system limit check, " + "failing open for user_id={}, error={}", + user_id, + exc, + ) + + return current_user + + async def enforce_billing_limits( + self, + *, + current_user: CurrentUser, + ) -> None: + await self._capacity_service.enforce_billing_limits(current_user=current_user) + + async def enforce_route_system_limit( + self, + *, + route_context: RouteAdmissionContext, + ) -> None: + await self._route_policy_service.enforce_route_system_limit( + route_context=route_context, + ) + + async def enforce_job_creation_capacity( + self, + *, + db: AsyncSession, + current_user: CurrentUser, + ) -> None: + await self._capacity_service.enforce_job_creation_capacity( + db=db, + current_user=current_user, + ) diff --git a/apps/api/app/services/redis/__init__.py b/apps/api/app/services/redis/__init__.py deleted file mode 100644 index a7c271e34..000000000 --- a/apps/api/app/services/redis/__init__.py +++ /dev/null @@ -1,22 +0,0 @@ -"""API-specific Redis service exports.""" - -# Re-export shared Redis services used by the API runtime. -from shared.services.redis import ( - JobInfoRedisService, - RedisService, - RedisServiceFactory, - UserRedisService, -) -from shared.services.redis.job_metadata_service import JobMetadataService -from shared.services.redis.rate_limit_service import RateLimitService -from shared.services.redis.task_redis_service import TaskRedisService - -__all__ = [ - "JobMetadataService", - "TaskRedisService", - "RateLimitService", - "RedisService", # Re-exported from shared. - "RedisServiceFactory", # Re-exported from shared. - "UserRedisService", # Re-exported from shared for both Worker and API. - "JobInfoRedisService", # Re-exported from shared for both Worker and API. -] diff --git a/apps/api/app/services/s3_events/__init__.py b/apps/api/app/services/s3_events/__init__.py new file mode 100644 index 000000000..b68b3a02b --- /dev/null +++ b/apps/api/app/services/s3_events/__init__.py @@ -0,0 +1 @@ +"""S3-compatible storage event services.""" diff --git a/apps/api/app/services/s3_events/event_handlers.py b/apps/api/app/services/s3_events/event_handlers.py new file mode 100644 index 000000000..b3408012d --- /dev/null +++ b/apps/api/app/services/s3_events/event_handlers.py @@ -0,0 +1,176 @@ +"""Storage event protocol handlers.""" +from __future__ import annotations + +import json +import os +from typing import Any + +from app.services.s3_events.nested_message_decoder import decode_nested_json_message +from app.services.s3_events.signature_verification import ( + verify_minio_signature, + verify_oss_signature, +) +from app.services.s3_events.subscription_service import confirm_sns_subscription +from app.services.s3_events.upload_event_service import process_upload_events +from loguru import logger + +from shared.models.schemas.oss_event import OSSEvent +from shared.models.schemas.s3_event import S3Event + + +async def handle_sns_event(body: bytes) -> dict[str, str] | None: + try: + sns_message = json.loads(body.decode("utf-8")) + message_type = sns_message.get("Type") + logger.info(f"SNS message type: {message_type}") + + if message_type == "SubscriptionConfirmation": + logger.info("Received an SNS subscription confirmation request") + subscribe_url = sns_message.get("SubscribeURL") + if subscribe_url: + logger.info(f"SNS subscription confirmation URL: {subscribe_url}") + return await confirm_sns_subscription(subscribe_url) + + logger.warning("SNS subscription confirmation did not include SubscribeURL") + return {"message": "SNS subscription confirmation failed"} + + if message_type == "Notification": + logger.info("Received an SNS notification") + logger.info(f"SNS message payload: {sns_message}") + await _handle_sns_notification(sns_message) + return None + + logger.warning(f"Unknown SNS message type: {message_type}") + return {"message": f"Unknown SNS message type: {message_type}"} + + except Exception as exc: + logger.error(f"Failed to handle SNS event: {exc}") + raise + + +async def handle_minio_event(body: bytes, auth_token: str) -> None: + try: + from shared.core.config import settings + + expected_token = getattr(settings, "S3_WEBHOOK_AUTH_TOKEN", "") + if not verify_minio_signature(auth_token, expected_token): + logger.warning("MinIO webhook authentication failed") + return + + s3_event_data = json.loads(body.decode("utf-8")) + await process_upload_events(S3Event(**s3_event_data)) + + except Exception as exc: + logger.error(f"Failed to handle MinIO event: {exc}") + + +async def handle_direct_s3_event(body: bytes) -> None: + try: + s3_event_data = json.loads(body.decode("utf-8")) + await process_upload_events(S3Event(**s3_event_data)) + + except Exception as exc: + logger.error(f"Failed to handle direct S3 event: {exc}") + + +def is_oss_event(headers: dict[str, str]) -> bool: + storage_type = os.getenv("S3_TYPE", "s3").lower() + if storage_type == "oss": + return True + + if "x-oss-pub-key-url" in headers: + return True + + if "x-mns-version" in headers or "x-mns-signing-cert-url" in headers: + return True + user_agent = headers.get("user-agent") or headers.get("User-Agent") + return bool(user_agent and "Aliyun Notification Service Agent" in user_agent) + + +async def handle_oss_event(body: bytes, headers: dict[str, str]) -> None: + try: + if not verify_oss_signature(body, headers): + logger.warning("OSS event signature verification failed") + return + + event_data = json.loads(body.decode("utf-8")) + logger.info(f"OSS event payload: {event_data}") + event_data = _unwrap_mns_message(event_data) + + if "events" in event_data: + oss_event = OSSEvent(**event_data) + elif "Records" in event_data: + oss_event = convert_s3_format_to_oss(event_data) + else: + logger.error(f"Unknown OSS event format: {event_data}") + return + + await process_upload_events(oss_event.to_s3_event()) + + except Exception as exc: + logger.error(f"Failed to handle OSS event: {exc}") + raise + + +def convert_s3_format_to_oss(event_data: dict[str, Any]) -> OSSEvent: + from shared.models.schemas.oss_event import OSSEventRecord + + records = event_data.get("Records", []) + oss_records = [ + OSSEventRecord( + eventName=record.get("eventName", "").replace("s3:", ""), + eventSource="acs:oss", + eventTime=record.get("eventTime", ""), + region=record.get("awsRegion", ""), + oss={ + "bucket": record.get("s3", {}).get("bucket", {}), + "object": record.get("s3", {}).get("object", {}), + }, + ) + for record in records + ] + + return OSSEvent(events=oss_records) + + +async def _handle_sns_notification(sns_message: dict[str, Any]) -> None: + try: + s3_event_data = json.loads(sns_message["Message"]) + logger.info(f"S3 event payload: {s3_event_data}") + + if ( + isinstance(s3_event_data, dict) + and s3_event_data.get("Event") == "s3:TestEvent" + ): + logger.info("Skip S3 test event") + return + + await process_upload_events(S3Event(**s3_event_data)) + except Exception as exc: + logger.error(f"Failed to parse the S3 event payload: {exc}") + logger.error(f"SNS payload: {sns_message}") + try: + await process_upload_events(S3Event(**sns_message)) + except Exception as fallback_exc: + logger.error( + "Fallback parsing of the SNS payload as an S3 event also failed: " + f"{fallback_exc}" + ) + raise + + +def _unwrap_mns_message(event_data: dict[str, Any]) -> dict[str, Any]: + if not isinstance(event_data, dict) or "Message" not in event_data: + return event_data + + inner = event_data.get("Message") + if isinstance(inner, dict): + return inner + if not isinstance(inner, str): + return event_data + + decoded = decode_nested_json_message(inner) + if decoded is not None: + logger.info(f"Decoded MNS Message payload: {decoded}") + return decoded + return event_data diff --git a/apps/api/app/services/s3_events/intake_outcome.py b/apps/api/app/services/s3_events/intake_outcome.py new file mode 100644 index 000000000..24a02189f --- /dev/null +++ b/apps/api/app/services/s3_events/intake_outcome.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +_SENSITIVE_HEADER_NAMES: frozenset[str] = frozenset( + { + "authorization", + "x-amz-sns-signature", + "x-amz-sns-signing-cert-url", + "x-minio-auth-token", + "x-oss-signature", + "x-oss-pub-key-url", + "x-mns-signature", + "x-mns-signing-cert-url", + } +) + + +@dataclass(frozen=True) +class StorageEventIntakeOutcome: + message: str + reason: str + + def to_response(self) -> dict[str, str]: + return {"message": self.message} + + +def build_storage_event_error_acknowledgement() -> StorageEventIntakeOutcome: + return StorageEventIntakeOutcome( + message="Event handling completed", + reason="acked_after_handler_error", + ) + + +def sanitize_storage_event_headers(headers: dict[str, str]) -> dict[str, str]: + sanitized: dict[str, str] = {} + for name, value in headers.items(): + if name.lower() in _SENSITIVE_HEADER_NAMES: + sanitized[name] = "" + else: + sanitized[name] = value + return sanitized diff --git a/apps/api/app/services/s3_events/nested_message_decoder.py b/apps/api/app/services/s3_events/nested_message_decoder.py new file mode 100644 index 000000000..eb3e29c0f --- /dev/null +++ b/apps/api/app/services/s3_events/nested_message_decoder.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import base64 +import binascii +import json +from collections.abc import Callable +from typing import Any + +from loguru import logger + + +def decode_nested_json_message(inner: str) -> dict[str, Any] | None: + for decoder in (_decode_base64_json, _decode_plain_json): + decoded = _try_decode(inner, decoder) + if decoded is not None: + return decoded + return None + + +def _try_decode( + inner: str, + decoder: Callable[[str], object], +) -> dict[str, Any] | None: + try: + decoded = decoder(inner) + except (binascii.Error, json.JSONDecodeError, UnicodeDecodeError, ValueError) as exc: + logger.debug("Nested storage event message decode failed: {}", exc) + return None + return decoded if isinstance(decoded, dict) else None + + +def _decode_base64_json(inner: str) -> object: + decoded_bytes = base64.b64decode(inner, validate=True) + decoded_str = decoded_bytes.decode("utf-8") + return json.loads(decoded_str) + + +def _decode_plain_json(inner: str) -> object: + return json.loads(inner) diff --git a/apps/api/app/services/s3_events/service.py b/apps/api/app/services/s3_events/service.py new file mode 100644 index 000000000..9382c67bb --- /dev/null +++ b/apps/api/app/services/s3_events/service.py @@ -0,0 +1,57 @@ +"""Application service for S3-compatible storage event webhooks.""" +from __future__ import annotations + +from loguru import logger + +from app.services.s3_events.event_handlers import ( + handle_direct_s3_event, + handle_minio_event, + handle_oss_event, + handle_sns_event, + is_oss_event, +) +from app.services.s3_events.intake_outcome import ( + build_storage_event_error_acknowledgement, +) + + +async def handle_s3_event_post( + *, + body: bytes, + headers: dict[str, str], + sns_message_type: str | None, + minio_auth_token: str | None, +) -> dict[str, str]: + if sns_message_type: + result = await handle_sns_event(body) + if result: + return result + elif is_oss_event(headers): + await handle_oss_event(body, headers) + elif minio_auth_token: + await handle_minio_event(body, minio_auth_token) + else: + await handle_direct_s3_event(body) + + return {"message": "Event handled successfully"} + + +async def safely_handle_s3_event_post( + *, + body: bytes, + headers: dict[str, str], + sns_message_type: str | None, + minio_auth_token: str | None, +) -> dict[str, str]: + try: + return await handle_s3_event_post( + body=body, + headers=headers, + sns_message_type=sns_message_type, + minio_auth_token=minio_auth_token, + ) + except Exception as exc: + outcome = build_storage_event_error_acknowledgement() + logger.error(f"Failed to handle S3 event: {exc}") + logger.warning(f"S3 event intake outcome: reason={outcome.reason}") + return outcome.to_response() diff --git a/apps/api/app/services/s3_events/signature_verification.py b/apps/api/app/services/s3_events/signature_verification.py new file mode 100644 index 000000000..86d918b23 --- /dev/null +++ b/apps/api/app/services/s3_events/signature_verification.py @@ -0,0 +1,37 @@ +"""Signature and token checks for storage event callbacks.""" +from __future__ import annotations + +from loguru import logger + + +def verify_sns_signature(request_body: bytes, signature: str, message: str) -> bool: + del request_body, signature, message + return True + + +def verify_minio_signature(auth_token: str, expected_token: str) -> bool: + if not expected_token: + return True + return auth_token == expected_token + + +def verify_oss_signature(request_body: bytes, headers: dict[str, str]) -> bool: + del request_body, headers + try: + from shared.core.config import settings + + if not getattr(settings, "OSS_EVENT_VERIFY_SIGNATURE", True): + return True + + callback_key = getattr(settings, "OSS_EVENT_CALLBACK_KEY", "") + if not callback_key: + logger.warning( + "OSS_EVENT_CALLBACK_KEY is not configured; skipping signature verification" + ) + return True + + # TODO: Implement OSS callback signature verification. + return True + except Exception as exc: + logger.error(f"OSS signature verification failed: {exc}") + return False diff --git a/apps/api/app/services/s3_events/subscription_service.py b/apps/api/app/services/s3_events/subscription_service.py new file mode 100644 index 000000000..c64ab78e5 --- /dev/null +++ b/apps/api/app/services/s3_events/subscription_service.py @@ -0,0 +1,47 @@ +"""SNS subscription confirmation handling.""" +from __future__ import annotations + +from loguru import logger + +from shared.services.http.pinned_outbound import send_pinned_outbound_request +from shared.services.http.url_security import validate_http_url_and_resolve_ip_async + + +SNS_SUBSCRIPTION_TIMEOUT_SECONDS = 10 + + +async def confirm_sns_subscription(subscribe_url: str) -> dict[str, str]: + validation = await validate_http_url_and_resolve_ip_async(subscribe_url) + if not validation.is_valid: + logger.warning( + f"SNS subscription confirmation URL failed validation: {validation.error_message}" + ) + return {"message": "SNS subscription confirmation failed"} + + if not validation.validated_ip: + logger.warning("SNS subscription confirmation URL validation returned no IP") + return {"message": "SNS subscription confirmation failed"} + + try: + response = await send_pinned_outbound_request( + method="GET", + url=subscribe_url, + pinned_ip=validation.validated_ip, + timeout_seconds=SNS_SUBSCRIPTION_TIMEOUT_SECONDS, + ) + if response.status == 200: + logger.info("SNS subscription confirmed successfully") + return {"message": "SNS subscription confirmed"} + + if 300 <= response.status < 400: + logger.warning( + f"SNS subscription confirmation redirect blocked, status={response.status}" + ) + else: + logger.error( + f"SNS subscription confirmation failed, status={response.status}" + ) + return {"message": "SNS subscription confirmation failed"} + except Exception as exc: + logger.error(f"Failed to reach the SNS confirmation URL: {exc}") + return {"message": "SNS subscription confirmation failed"} diff --git a/apps/api/app/services/s3_events/upload_event_service.py b/apps/api/app/services/s3_events/upload_event_service.py new file mode 100644 index 000000000..8b41ca310 --- /dev/null +++ b/apps/api/app/services/s3_events/upload_event_service.py @@ -0,0 +1,73 @@ +"""Process storage upload-complete events into job workflow handoffs.""" +from __future__ import annotations + +import os + +from app.repositories.job_repository import JobRepository +from app.services.document_ingestion.handoff_service import ( + DocumentIngestionHandoffService, +) +from loguru import logger + +from shared.core.database import get_db_context +from shared.models.schemas.s3_event import S3Event + + +def extract_job_id_from_s3_key(s3_key: str) -> str | None: + if not s3_key.startswith("uploads/"): + return None + + filename = s3_key[8:] + return os.path.splitext(filename)[0] + + +async def process_upload_events(s3_event: S3Event) -> None: + try: + upload_events = s3_event.get_upload_events() + job_repo = JobRepository() + handoff_service = DocumentIngestionHandoffService() + + for event in upload_events: + s3_key = event.object_key or event.s3.get("object", {}).get("key") + if not s3_key: + continue + + job_id = extract_job_id_from_s3_key(s3_key) + if not job_id: + logger.warning(f"Could not extract job_id from S3 key: {s3_key}") + continue + + logger.info(f"Processing S3 upload event: {s3_key} -> job_id={job_id}") + + async with get_db_context() as db: + job = await job_repo.get_job_by_id(db, job_id) + if not job: + logger.warning(f"No job found for upload event: {job_id}") + continue + + if job.status != "waiting-file": + logger.info( + f"Job {job_id} is not in waiting-file status: {job.status}" + ) + continue + + from shared.core.config import settings + from shared.core.state_machine.states import is_job_expired + + if is_job_expired(job.updated_at, settings.JOB_WAITING_EXPIRE_SECONDS): + logger.warning(f"Job {job_id} upload expired, marking failed") + await handoff_service.mark_upload_expired(db, job=job) + continue + + await handoff_service.start_uploaded_file_workflow( + db=db, + job=job, + user_id=str(job.user_id), + trigger="s3_upload_completed", + ) + + logger.info(f"Triggered processing for job {job_id}") + + except Exception as exc: + logger.error(f"Failed to process upload events: {exc}") + raise diff --git a/apps/api/app/services/state_machine/__init__.py b/apps/api/app/services/state_machine/__init__.py deleted file mode 100644 index 2fc629999..000000000 --- a/apps/api/app/services/state_machine/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -""" -API-side compatibility layer for the shared state machine. - -The canonical implementation now lives under ``shared.core.state_machine``. -This package keeps stable imports for API callers while avoiding duplicate -state machine logic inside ``apps/api``. -""" - -from shared.core.state_machine.service import AsyncStateMachineService -from shared.core.state_machine.states import JobStatus - -from .manager import JobStateMachine - -StateMachineService = AsyncStateMachineService - -__all__ = [ - "AsyncStateMachineService", - "JobStateMachine", - "JobStatus", - "StateMachineService", -] diff --git a/apps/api/app/services/state_machine/manager.py b/apps/api/app/services/state_machine/manager.py deleted file mode 100644 index f23804115..000000000 --- a/apps/api/app/services/state_machine/manager.py +++ /dev/null @@ -1,128 +0,0 @@ -""" -API-side state machine facade. - -Core transition logic lives in ``shared.core.state_machine.service``. This -module keeps the ``JobStateMachine`` entry point that API code already uses. -""" - -from typing import Any, Dict, Optional - -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.state_machine.service import AsyncStateMachineService -from shared.services.redis import RedisServiceFactory - - -class JobStateMachine: - """Compatibility facade over ``AsyncStateMachineService``.""" - - def __init__(self, redis_service: Optional[Any] = None) -> None: - self.redis = redis_service or RedisServiceFactory.get_service() - self.state_machine = AsyncStateMachineService(self.redis) - - async def transition( - self, - db: AsyncSession, - job_id: str, - to_state: str, - transition_reason: str = "normal_transition", - operator_id: Optional[str] = None, - operator_type: str = "system", - metadata: Optional[Dict[str, Any]] = None, - auto_commit: bool = True, - ) -> bool: - """Execute a CAS-protected state transition.""" - try: - return await self.state_machine.transition( - db=db, - job_id=job_id, - to_state=to_state, - transition_reason=transition_reason, - operator_id=operator_id, - operator_type=operator_type, - metadata=metadata, - auto_commit=auto_commit, - ) - except Exception as err: - logger.error(f"Job {job_id} transition failed: {err}") - return False - - async def mark_failed( - self, - db: AsyncSession, - job_id: str, - error_message: str, - error_code: str = "UNKNOWN", - error_details: Optional[Dict[str, Any]] = None, - operator_id: Optional[str] = None, - metadata: Optional[Dict[str, Any]] = None, - auto_commit: bool = True, - ) -> bool: - """Mark a job as failed.""" - try: - return await self.state_machine.mark_failed( - db=db, - job_id=job_id, - error_message=error_message, - error_code=error_code, - error_details=error_details, - operator_id=operator_id, - metadata=metadata, - auto_commit=auto_commit, - ) - except Exception as err: - logger.error(f"Failed to mark Job {job_id} as failed: {err}") - return False - - async def mark_completed( - self, - db: AsyncSession, - job_id: str, - result_metadata: Optional[Dict[str, Any]] = None, - operator_id: Optional[str] = None, - auto_commit: bool = True, - ) -> bool: - """Mark a job as completed.""" - try: - return await self.state_machine.mark_completed( - db=db, - job_id=job_id, - result_metadata=result_metadata, - operator_id=operator_id, - auto_commit=auto_commit, - ) - except Exception as err: - logger.error(f"Failed to mark Job {job_id} as completed: {err}") - return False - - async def handle_retry( - self, - db: AsyncSession, - job_id: str, - retry_metadata: Optional[Dict[str, Any]] = None, - operator_id: Optional[str] = None, - ) -> bool: - """Retry a job through the shared state machine.""" - try: - return await self.state_machine.handle_retry( - db=db, - job_id=job_id, - retry_metadata=retry_metadata, - operator_id=operator_id, - ) - except Exception as err: - logger.error(f"Failed to retry Job {job_id}: {err}") - return False - - async def get_current_state( - self, - db: AsyncSession, - job_id: str, - ) -> Optional[str]: - """Read the current state through the shared service.""" - try: - return await self.state_machine.get_current_state(db=db, job_id=job_id) - except Exception as err: - logger.error(f"Failed to read Job {job_id} state: {err}") - return None diff --git a/apps/api/app/services/webhook/__init__.py b/apps/api/app/services/webhook/__init__.py new file mode 100644 index 000000000..c2528a3fc --- /dev/null +++ b/apps/api/app/services/webhook/__init__.py @@ -0,0 +1 @@ +"""Webhook delivery workflows.""" diff --git a/apps/api/app/services/webhook/qstash_callback_service.py b/apps/api/app/services/webhook/qstash_callback_service.py new file mode 100644 index 000000000..352e8fcd8 --- /dev/null +++ b/apps/api/app/services/webhook/qstash_callback_service.py @@ -0,0 +1,255 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Any, Literal +from uuid import NAMESPACE_URL, uuid5 + +from loguru import logger +from sqlalchemy import select + +from shared.core.config import app_config +from shared.core.database_sync import get_sync_db_context +from shared.models.database.webhook import WebhookEvent, WebhookEventStatus +from shared.models.database.webhook_log import WebhookLog + + +QStashCallbackOutcomeKind = Literal["processed", "missing_event_id", "event_not_found"] + + +@dataclass(frozen=True) +class QStashCallbackOutcome: + kind: QStashCallbackOutcomeKind + event_id: str | None = None + + @classmethod + def processed(cls, *, event_id: str) -> QStashCallbackOutcome: + return cls(kind="processed", event_id=event_id) + + @classmethod + def missing_event_id(cls) -> QStashCallbackOutcome: + return cls(kind="missing_event_id") + + @classmethod + def event_not_found(cls, *, event_id: str) -> QStashCallbackOutcome: + return cls(kind="event_not_found", event_id=event_id) + + +def get_qstash_verification_url(callback_path: str, request_url: str) -> str: + callback_base_url = app_config.QSTASH_CALLBACK_BASE_URL + if callback_base_url: + return f"{callback_base_url.rstrip('/')}{callback_path}" + return request_url + + +def verify_qstash_signature(raw_body: bytes, signature: str, url: str) -> bool: + current_key = app_config.QSTASH_CURRENT_SIGNING_KEY + next_key = app_config.QSTASH_NEXT_SIGNING_KEY + + if not current_key or not next_key: + logger.error("QStash signing keys not configured — rejecting callback") + return False + + try: + from qstash import Receiver + + receiver = Receiver( + current_signing_key=current_key, + next_signing_key=next_key, + ) + receiver.verify( + body=raw_body.decode("utf-8"), + signature=signature, + url=url, + ) + return True + except Exception as exc: + logger.warning( + "QStash signature verification failed: error_type={error_type}, url={url}", + error_type=type(exc).__name__, + url=url, + ) + return False + + +def handle_qstash_success_callback(raw_body: bytes) -> QStashCallbackOutcome: + data = extract_callback_data(raw_body) + event_id = find_event_id(data) + + if not event_id: + logger.warning("QStash callback: missing event_id, cannot correlate") + return QStashCallbackOutcome.missing_event_id() + + retried = data.get("retried", 0) + logger.info( + f"QStash callback: event_id={event_id}, status={data.get('status')}, " + f"retried={retried}, qstash_message_id={data.get('sourceMessageId')}" + ) + + return process_qstash_callback( + data, + event_id, + get_callback_event_status(data), + "callback", + ) + + +def handle_qstash_failure_callback(raw_body: bytes) -> QStashCallbackOutcome: + data = extract_callback_data(raw_body) + event_id = find_event_id(data) + + if not event_id: + logger.warning("QStash failure callback: missing event_id, cannot correlate") + return QStashCallbackOutcome.missing_event_id() + + retried = data.get("retried", 0) + max_retries = data.get("maxRetries", 0) + logger.warning( + f"QStash failure: event_id={event_id}, status={data.get('status')}, " + f"retried={retried}/{max_retries}, qstash_message_id={data.get('sourceMessageId')}" + ) + + return process_qstash_callback( + data, + event_id, + WebhookEventStatus.FAILED, + "failure", + ) + + +def extract_callback_data(body: bytes) -> dict[str, Any]: + try: + return json.loads(body) + except (json.JSONDecodeError, ValueError): + return {"raw": body.decode("utf-8", errors="replace")} + + +def find_event_id(data: dict[str, Any]) -> str | None: + source_header = data.get("sourceHeader", {}) or {} + event_id = normalize_header_value( + source_header.get("X-Knowhere-Event-Id") + or source_header.get("x-knowhere-event-id") + ) + if not event_id: + for key, value in source_header.items(): + if key.lower() == "x-knowhere-event-id": + event_id = normalize_header_value(value) + break + return event_id + + +def normalize_header_value(value: Any) -> str | None: + if isinstance(value, list): + if not value: + return None + first_value = value[0] + return first_value if isinstance(first_value, str) else str(first_value) + + if isinstance(value, str): + return value + + if value is None: + return None + + return str(value) + + +def build_callback_log_idempotency_key( + qstash_message_id: str | None, + event_id: str, +) -> str: + if qstash_message_id: + return str(uuid5(NAMESPACE_URL, qstash_message_id)) + return event_id + + +def get_response_status_code(value: Any) -> int | None: + if value is None: + return None + + try: + return int(value) + except (TypeError, ValueError): + return None + + +def is_success_response_status(status_code: int | None) -> bool: + return status_code is not None and 200 <= status_code < 300 + + +def get_callback_event_status(data: dict[str, Any]) -> str: + response_status = get_response_status_code(data.get("status")) + if is_success_response_status(response_status): + return WebhookEventStatus.DELIVERED + return WebhookEventStatus.DELIVERING + + +def resolve_event_status(current_status: str, callback_status: str) -> str: + if current_status in ( + WebhookEventStatus.DELIVERED, + WebhookEventStatus.FAILED, + WebhookEventStatus.CANCELED, + ): + return current_status + return callback_status + + +def process_qstash_callback( + data: dict[str, Any], + event_id: str, + callback_status: str, + log_label: str, +) -> QStashCallbackOutcome: + response_status_code = get_response_status_code(data.get("status")) + response_body = data.get("body", "") + qstash_message_id = data.get("sourceMessageId") + retried = data.get("retried", 0) + is_failed_delivery_attempt = ( + callback_status == WebhookEventStatus.FAILED + or ( + callback_status == WebhookEventStatus.DELIVERING + and not is_success_response_status(response_status_code) + ) + ) + error_message = None + if is_failed_delivery_attempt: + error_message = data.get("error") or response_body + + with get_sync_db_context() as db: + event = db.execute( + select(WebhookEvent).where(WebhookEvent.id == event_id) + ).scalar_one_or_none() + + if not event: + logger.warning(f"QStash {log_label}: event {event_id} not found in DB") + return QStashCallbackOutcome.event_not_found(event_id=event_id) + + now = datetime.now(timezone.utc).replace(tzinfo=None) + event_status = resolve_event_status(event.status, callback_status) + attempt_number = retried + 1 + event.status = event_status + event.attempts = max(event.attempts, attempt_number) + event.updated_at = now + + log = WebhookLog( + job_id=event.job_id, + event_id=event.id, + webhook_url=event.target_url, + attempt_number=attempt_number, + request_payload=event.payload, + signature="", + idempotency_key=build_callback_log_idempotency_key( + qstash_message_id, + event.id, + ), + response_status_code=response_status_code, + response_body=response_body[:4096] if response_body else None, + error_message=str(error_message)[:4096] if error_message else None, + duration_ms=0, + qstash_message_id=qstash_message_id, + ) + db.add(log) + db.commit() + + return QStashCallbackOutcome.processed(event_id=event_id) diff --git a/apps/api/app/services/webhook_service.py b/apps/api/app/services/webhook_service.py deleted file mode 100644 index 25b1f1b21..000000000 --- a/apps/api/app/services/webhook_service.py +++ /dev/null @@ -1,182 +0,0 @@ -""" -Consolidated Webhook Service - -Single responsibility: create WebhookEvents and publish them for async delivery. -""" - -from datetime import datetime, timezone -from typing import Any, Dict, Optional - -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.exceptions.webhook_exceptions import WebhookConfigException -from shared.core.response import build_standard_error_response -from shared.models.database.webhook import WebhookEvent, WebhookEventStatus - - -class WebhookService: - """ - Webhook Service - Transactional Outbox Pattern. - - Creates WebhookEvent records and publishes them via QStash for async delivery. - """ - - async def create_job_completion_event( - self, db: AsyncSession, job_id: str, webhook_url: str - ) -> WebhookEvent: - """ - Create webhook event for job completion. - - Args: - db: Database session (in active transaction) - job_id: Job ID - webhook_url: Webhook URL - - Returns: - Created WebhookEvent - """ - # Build minimal payload - dispatcher will enrich at delivery time - payload: Dict[str, Any] = { - "event": "job.completed", - "job_id": job_id, - "status": "completed", - "completed_at": datetime.now(timezone.utc).isoformat(), - # NOTE: result_url and result added by dispatcher at delivery - } - - # Create event - event = await self._create_event( - db=db, job_id=job_id, target_url=webhook_url, payload=payload - ) - - logger.info( - f"Job completion webhook event created: event_id={event.id}, job_id={job_id}" - ) - return event - - async def create_job_failure_event( - self, - db: AsyncSession, - job_id: str, - error_message: str, - error_type: Optional[str] = None, - error_code: str = "UNKNOWN", - error_details: Optional[Dict[str, Any]] = None, - webhook_url: Optional[str] = None, - ) -> WebhookEvent: - """ - Create webhook event for job failure. - - Args: - db: Database session (in active transaction) - job_id: Job ID - error_message: Error message - error_type: Error type (optional) - error_code: Error code - error_details: Structured error details (optional) - webhook_url: Webhook URL - - Returns: - Created WebhookEvent - """ - # Build payload with standardized error format - payload: Dict[str, Any] = { - "event": "job.failed", - "job_id": job_id, - "status": "failed", - "failed_at": datetime.now(timezone.utc).isoformat(), - "error": build_standard_error_response( - code=error_code, - message=error_message, - request_id=job_id, - details=error_details, - ), - } - - if webhook_url is None: - raise WebhookConfigException( - internal_message="Missing webhook target_url", - user_message="Webhook URL is required.", - ) - - # Create event - event = await self._create_event( - db=db, job_id=job_id, target_url=webhook_url, payload=payload - ) - - logger.info( - f"Job failure webhook event created: event_id={event.id}, job_id={job_id}" - ) - return event - - async def publish_to_queue(self, event_id: str) -> bool: - """Publish a webhook event for async delivery via QStash after commit.""" - return await self._publish_via_qstash(event_id) - - async def _publish_via_qstash(self, event_id: str) -> bool: - """Publish via QStash for managed delivery and retry.""" - try: - from shared.services.webhook.qstash_publisher import ( - get_qstash_webhook_publisher, - ) - - publisher = get_qstash_webhook_publisher() - message_id = publisher.publish_event(event_id) - if message_id: - logger.info( - f"Webhook published via QStash: event_id={event_id}, message_id={message_id}" - ) - return True - logger.warning( - f"QStash publish returned no message_id: event_id={event_id}" - ) - return False - except Exception as exc: - logger.error(f"QStash publish failed: event_id={event_id}, error={exc}") - return False - - async def _create_event( - self, db: AsyncSession, job_id: str, target_url: str, payload: Dict[str, Any] - ) -> WebhookEvent: - """ - Create a WebhookEvent record. - """ - # Validate - if not target_url: - raise WebhookConfigException( - internal_message="Missing webhook target_url", - user_message="Webhook URL is required.", - ) - - if not target_url.startswith(("http://", "https://")): - raise WebhookConfigException( - internal_message=f"Invalid webhook scheme: {target_url}", - user_message="Webhook URL must start with http:// or https://.", - ) - - # Create event - event = WebhookEvent( - job_id=job_id, - target_url=target_url, - payload=payload, - status=WebhookEventStatus.PENDING, - attempts=0, - ) - - db.add(event) - await db.flush() # Get ID, but don't commit (caller controls transaction) - - return event - - -# Singleton -_webhook_service: Optional[WebhookService] = None - - -def get_webhook_service() -> WebhookService: - """Get singleton WebhookService instance.""" - global _webhook_service - if _webhook_service is None: - _webhook_service = WebhookService() - return _webhook_service diff --git a/apps/api/main.py b/apps/api/main.py index ed2c2ae6d..db75099bf 100644 --- a/apps/api/main.py +++ b/apps/api/main.py @@ -1,6 +1,5 @@ import os from pathlib import Path -import httpx import uvicorn from fastapi import FastAPI from starlette.routing import Route @@ -18,7 +17,6 @@ from contextlib import asynccontextmanager from app.api.api_router import api_router from app.core.middleware import setup_cors, LoggingMiddleware -from app.core.image_cli import ImageCli from app.core.exception_handlers import setup_exception_handlers from app.mcp import create_retrieval_mcp_server from app.services.rate_limit.rule_loader import load_rules @@ -59,8 +57,6 @@ async def lifespan(app: FastAPI): await redis_pool_manager.init_pool() logger.info("Redis connection pool created.") - ImageCli.http_client = httpx.AsyncClient(timeout=30.0, follow_redirects=True) - # Initialize rate limiter rules from DB. # Changes now require a pod restart to take effect. from app.services.rate_limit.config import RateLimitConfig @@ -75,7 +71,7 @@ async def lifespan(app: FastAPI): mcp_server = getattr(app.state, "retrieval_mcp_server", None) mcp_session_manager = getattr(mcp_server, "session_manager", None) - logger.info("knowledge library API service started!") + logger.info("Document API service started!") if mcp_session_manager is not None: async with mcp_session_manager.run(): yield @@ -83,7 +79,7 @@ async def lifespan(app: FastAPI): yield try: - from shared.services.retrieval.app_service import ( + from shared.services.retrieval.stats.recorder import ( drain_retrieval_hit_stats_updates, ) @@ -92,13 +88,13 @@ async def lifespan(app: FastAPI): logger.error(f"retrieval hit stats drain failed: {e}") try: - from shared.utils.http_clients import close_async_client + from shared.services.http.client_pool import close_async_client await close_async_client() except Exception as e: logger.error(f"async HTTP client close failed: {e}") - logger.info("knowledge library API service stopped!") + logger.info("Document API service stopped!") await safe_dispose_engine(engine) logger.info("database engine connection pool disposed.") logger.info("service stopped.") @@ -139,7 +135,7 @@ def create_app() -> FastAPI: @app.get("/", tags=["Root"]) async def read_root(): - return {"message": f"Welcome to {app.title} - Knowledge Base API Service!"} + return {"message": f"Welcome to {app.title} - Document API Service!"} @app.api_route("/health", methods=["GET", "HEAD"], tags=["Health"]) async def health_check(): @@ -173,11 +169,10 @@ async def health_check(): return app -# Worker settings removed as DsTasks.py was deleted app = create_app() if __name__ == "__main__": - logger.info("Knowledge Base API service starting...") + logger.info("Document API service starting...") port = 5005 reload = False # Enable hot reload host = "0.0.0.0" diff --git a/apps/api/scripts/add_credits.py b/apps/api/scripts/add_credits.py index f468aa281..675f5ce26 100644 --- a/apps/api/scripts/add_credits.py +++ b/apps/api/scripts/add_credits.py @@ -33,7 +33,6 @@ def _bootstrap_python_path() -> None: from shared.models.database.tier_limit import TierLimit from shared.models.database.user import User from shared.models.database.user_balance import UserBalance -from shared.models.database.webhook import WebhookEvent # noqa: F401 from shared.services.billing import CreditsService from shared.utils.utc_now import utc_now_naive diff --git a/apps/api/scripts/validate_demo_documents.py b/apps/api/scripts/validate_demo_documents.py new file mode 100644 index 000000000..479e46556 --- /dev/null +++ b/apps/api/scripts/validate_demo_documents.py @@ -0,0 +1,185 @@ +from __future__ import annotations +# ruff: noqa: E402 + +import argparse +import json +import os +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +def _bootstrap_python_path() -> None: + current_dir = Path(__file__).resolve().parent + api_root = current_dir.parent + repo_root = api_root.parent.parent + shared_python_path = repo_root / "packages" / "shared-python" + + for path in (api_root, shared_python_path): + path_value = os.fspath(path) + if path_value not in sys.path: + sys.path.insert(0, path_value) + + +_bootstrap_python_path() + +from app.services.demo.source_catalog import DemoSourceCatalog, DemoSourceDefinition +from shared.services.storage.zip_result_schema import ZipResultSchemaBuilder + + +@dataclass(frozen=True) +class DemoValidationIssue: + source_id: str + message: str + + +def validate_demo_documents(*, write: bool = False) -> list[DemoValidationIssue]: + catalog = DemoSourceCatalog() + issues: list[DemoValidationIssue] = [] + for source in catalog.list_sources(): + issues.extend(_validate_source(catalog=catalog, source=source, write=write)) + return issues + + +def _validate_source( + *, + catalog: DemoSourceCatalog, + source: DemoSourceDefinition, + write: bool, +) -> list[DemoValidationIssue]: + issues: list[DemoValidationIssue] = [] + source_directory = catalog.source_directory(source) + chunks = _load_chunks(source_directory) + if len(chunks) != source.chunk_count: + issues.append( + DemoValidationIssue( + source_id=source.demo_source_id, + message=( + f"chunk_count mismatch: catalog={source.chunk_count}, " + f"chunks.json={len(chunks)}" + ), + ) + ) + + original_path = source_directory / "original.pdf" + if not original_path.is_file(): + issues.append( + DemoValidationIssue( + source_id=source.demo_source_id, + message="original.pdf is missing", + ) + ) + elif original_path.stat().st_size != source.size_bytes: + issues.append( + DemoValidationIssue( + source_id=source.demo_source_id, + message=( + f"size_bytes mismatch: catalog={source.size_bytes}, " + f"original.pdf={original_path.stat().st_size}" + ), + ) + ) + + _validate_catalog_projection(catalog=catalog, source=source, issues=issues) + _validate_doc_nav( + source=source, + source_directory=source_directory, + chunks=chunks, + write=write, + issues=issues, + ) + return issues + + +def _load_chunks(source_directory: Path) -> list[dict[str, Any]]: + chunks_path = source_directory / "chunks.json" + payload = json.loads(chunks_path.read_text(encoding="utf-8")) + raw_chunks = payload.get("chunks") if isinstance(payload, dict) else None + if not isinstance(raw_chunks, list): + return [] + return [chunk for chunk in raw_chunks if isinstance(chunk, dict)] + + +def _validate_catalog_projection( + *, + catalog: DemoSourceCatalog, + source: DemoSourceDefinition, + issues: list[DemoValidationIssue], +) -> None: + try: + catalog.get_catalog() + except Exception as exc: + issues.append( + DemoValidationIssue( + source_id=source.demo_source_id, + message=f"catalog projection failed: {exc}", + ) + ) + + +def _validate_doc_nav( + *, + source: DemoSourceDefinition, + source_directory: Path, + chunks: list[dict[str, Any]], + write: bool, + issues: list[DemoValidationIssue], +) -> None: + doc_nav_path = source_directory / "doc_nav.json" + expected_doc_nav = ZipResultSchemaBuilder().build_doc_nav(chunks, source.title) + expected_text = _serialize_json(expected_doc_nav) + if write: + doc_nav_path.write_text(expected_text, encoding="utf-8") + return + + if not doc_nav_path.is_file(): + issues.append( + DemoValidationIssue( + source_id=source.demo_source_id, + message="doc_nav.json is missing; run with --write", + ) + ) + return + + current_text = doc_nav_path.read_text(encoding="utf-8") + if current_text != expected_text: + issues.append( + DemoValidationIssue( + source_id=source.demo_source_id, + message="doc_nav.json is stale; run with --write", + ) + ) + + +def _serialize_json(payload: dict[str, Any]) -> str: + return json.dumps(payload, ensure_ascii=False, indent=2) + "\n" + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Validate and optionally regenerate canonical demo documents.", + ) + parser.add_argument( + "--write", + action="store_true", + help="Regenerate doc_nav.json files from chunks.json before validating.", + ) + return parser + + +def main() -> int: + args = _build_parser().parse_args() + issues = validate_demo_documents(write=bool(args.write)) + if issues: + for issue in issues: + print(f"[FAIL] {issue.source_id}: {issue.message}") + return 1 + + action = "regenerated and validated" if args.write else "validated" + print(f"Demo documents {action}.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/api/tests/contract/test_agentic_answer_policy_contract.py b/apps/api/tests/contract/test_agentic_answer_policy_contract.py deleted file mode 100644 index e0c0cf1d2..000000000 --- a/apps/api/tests/contract/test_agentic_answer_policy_contract.py +++ /dev/null @@ -1,25 +0,0 @@ -from __future__ import annotations - -import pytest - -from shared.services.retrieval.agentic.policy import attempt_answer -from shared.services.retrieval.agentic.types import AgentRunConfig, AgentState - - -async def _malformed_json_wrapper(_prompt: str) -> str: - return '{"status": "DONE", "answer": "truncated"' - - -@pytest.mark.asyncio -async def test_attempt_answer_should_not_expose_malformed_json_wrapper() -> None: - status, answer, reason = await attempt_answer( - _malformed_json_wrapper, - query="What changed?", - evidence_text="┈ evidence", - state=AgentState(), - config=AgentRunConfig(), - ) - - assert status == "NOT_FOUND" - assert answer == "" - assert reason == "attempt_answer returned malformed JSON" diff --git a/apps/api/tests/contract/test_billing_contract.py b/apps/api/tests/contract/test_billing_contract.py index 75986d40e..6f5eb2be6 100644 --- a/apps/api/tests/contract/test_billing_contract.py +++ b/apps/api/tests/contract/test_billing_contract.py @@ -432,8 +432,8 @@ async def test_should_return_a_checkout_url_when_buying_a_credit_package( ], monkeypatch: MonkeyPatch, ) -> None: - class FakeStripeService: - async def create_checkout_session_for_credits_package( + class FakeStripePurchaseService: + async def create_credits_package_checkout_session( self, db, user_id: str, @@ -452,8 +452,14 @@ async def create_checkout_session_for_credits_package( return "https://checkout.stripe.test/session/contract-package" async with developer_api_client_factory() as api_client: - billing_module = importlib.import_module("app.api.v1.routes.billing") - monkeypatch.setattr(billing_module, "StripeService", FakeStripeService) + billing_service_module = importlib.import_module( + "app.services.billing.billing_command_workflow" + ) + monkeypatch.setattr( + billing_service_module, + "StripePurchaseService", + FakeStripePurchaseService, + ) response = await api_client.post( "/api/v1/billing/buy-credits-package", json={"price_id": "price_contract_package", "quantity": 2}, @@ -473,7 +479,7 @@ async def test_should_return_a_payment_intent_payload_when_buying_credits( ], monkeypatch: MonkeyPatch, ) -> None: - class FakeStripeService: + class FakeStripePurchaseService: async def create_payment_intent( self, user_id: str, @@ -491,8 +497,14 @@ async def create_payment_intent( } async with developer_api_client_factory() as api_client: - billing_module = importlib.import_module("app.api.v1.routes.billing") - monkeypatch.setattr(billing_module, "StripeService", FakeStripeService) + billing_service_module = importlib.import_module( + "app.services.billing.billing_command_workflow" + ) + monkeypatch.setattr( + billing_service_module, + "StripePurchaseService", + FakeStripePurchaseService, + ) response = await api_client.post( "/api/v1/billing/buy-credits", json={"credits_amount": 350}, @@ -503,3 +515,50 @@ async def create_payment_intent( "client_secret": "pi_contract_secret", "payment_intent_id": "pi_contract_id", } + + +@pytest.mark.asyncio +async def test_should_delegate_the_webhook_endpoint_to_the_stripe_webhook_service( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + class FakeStripeWebhookService: + async def handle_webhook( + self, + db, + *, + payload: bytes, + sig_header: str, + ) -> dict[str, object]: + del db + assert payload == b'{"type":"payment_intent.succeeded"}' + assert sig_header == "sig_contract" + return { + "status": "success", + "event_type": "payment_intent.succeeded", + "user_id": "local-dev-user", + } + + async with developer_api_client_factory() as api_client: + billing_service_module = importlib.import_module( + "app.services.billing.billing_command_workflow" + ) + monkeypatch.setattr( + billing_service_module, + "StripeWebhookService", + FakeStripeWebhookService, + ) + response = await api_client.post( + "/api/v1/billing/webhook", + content=b'{"type":"payment_intent.succeeded"}', + headers={"stripe-signature": "sig_contract"}, + ) + + assert response.status_code == 200 + assert response.json() == { + "status": "success", + "event_type": "payment_intent.succeeded", + "user_id": "local-dev-user", + } diff --git a/apps/api/tests/contract/test_chunk_document_path_contract.py b/apps/api/tests/contract/test_chunk_document_path_contract.py new file mode 100644 index 000000000..d10a99c46 --- /dev/null +++ b/apps/api/tests/contract/test_chunk_document_path_contract.py @@ -0,0 +1,247 @@ +from typing import cast + +import pytest + +from shared.services.retrieval.search.lexical_text import section_path_from_chunk_path +from shared.services.storage.zip_result_schema import ZipResultSchemaBuilder + + +@pytest.mark.parametrize( + "chunk_path", + [ + "acme.com/report.pdf/Intro/Subsection", + "team/foo.fragment/Intro/Subsection", + "team/foo.atlas/Intro/Subsection", + "team/photo.png/Intro/Subsection", + "images/report.pdf/Intro/Subsection", + "tables/report.pdf/Intro/Subsection", + "client.pdf/report.pdf/Intro/Subsection", + ], +) +def test_should_read_legacy_namespace_chunk_paths_as_document_sections( + chunk_path: str, +) -> None: + source_file_name = chunk_path.split("/")[1] + + doc_nav = ZipResultSchemaBuilder().build_doc_nav( + [ + { + "chunk_id": "chunk_legacy_path", + "type": "text", + "content": "legacy path content", + "path": chunk_path, + "metadata": {"summary": "legacy path summary"}, + } + ], + source_file_name, + ) + + sections = cast(list[dict[str, object]], doc_nav["sections"]) + children = cast(list[dict[str, object]], sections[0]["children"]) + + assert ( + section_path_from_chunk_path( + chunk_path, + source_file_name=source_file_name, + ) + == "Intro / Subsection" + ) + assert sections[0]["title"] == "Intro" + assert sections[0]["path"] == "/".join(chunk_path.split("/")[:3]) + assert children[0]["title"] == "Subsection" + assert children[0]["path"] == chunk_path + + +@pytest.mark.parametrize( + "chunk_path", + [ + "images/photo.png", + "tables/table.html", + ], +) +def test_should_keep_media_resource_paths_at_root(chunk_path: str) -> None: + assert section_path_from_chunk_path(chunk_path) == "Root" + + +def test_should_not_treat_dotted_section_titles_as_legacy_document_files() -> None: + chunk_path = "report.pdf/1. Introduction/Details" + + doc_nav = ZipResultSchemaBuilder().build_doc_nav( + [ + { + "chunk_id": "chunk_numbered_section", + "type": "text", + "content": "numbered section content", + "path": chunk_path, + "metadata": {"summary": "numbered section summary"}, + } + ], + "report.pdf", + ) + + sections = cast(list[dict[str, object]], doc_nav["sections"]) + children = cast(list[dict[str, object]], sections[0]["children"]) + + assert section_path_from_chunk_path(chunk_path) == "1. Introduction / Details" + assert sections[0]["title"] == "1. Introduction" + assert sections[0]["path"] == "report.pdf/1. Introduction" + assert children[0]["title"] == "Details" + assert children[0]["path"] == chunk_path + + +def test_should_read_arrow_delimited_document_paths_as_section_paths() -> None: + chunk_path = "report.pdf-->Intro-->Subsection" + + doc_nav = ZipResultSchemaBuilder().build_doc_nav( + [ + { + "chunk_id": "chunk_arrow_delimited_path", + "type": "text", + "content": "arrow delimited content", + "path": chunk_path, + "metadata": {"summary": "arrow delimited summary"}, + } + ], + "report.pdf", + ) + + sections = cast(list[dict[str, object]], doc_nav["sections"]) + children = cast(list[dict[str, object]], sections[0]["children"]) + + assert ( + section_path_from_chunk_path( + chunk_path, + source_file_name="report.pdf", + ) + == "Intro / Subsection" + ) + assert sections[0]["title"] == "Intro" + assert sections[0]["path"] == "report.pdf/Intro" + assert children[0]["title"] == "Subsection" + assert children[0]["path"] == "report.pdf/Intro/Subsection" + + +def test_should_preserve_literal_arrow_text_in_slash_paths() -> None: + chunk_path = "report.pdf/Inputs --> Outputs/Details" + + doc_nav = ZipResultSchemaBuilder().build_doc_nav( + [ + { + "chunk_id": "chunk_literal_arrow_heading", + "type": "text", + "content": "literal arrow content", + "path": chunk_path, + "metadata": {"summary": "literal arrow summary"}, + } + ], + "report.pdf", + ) + + sections = cast(list[dict[str, object]], doc_nav["sections"]) + children = cast(list[dict[str, object]], sections[0]["children"]) + + assert ( + section_path_from_chunk_path( + chunk_path, + source_file_name="report.pdf", + ) + == "Inputs --> Outputs / Details" + ) + assert sections[0]["title"] == "Inputs --> Outputs" + assert sections[0]["path"] == "report.pdf/Inputs --> Outputs" + assert children[0]["title"] == "Details" + assert children[0]["path"] == chunk_path + + +def test_should_preserve_single_level_literal_arrow_headings() -> None: + chunk_path = "report.pdf/Inputs --> Outputs" + + doc_nav = ZipResultSchemaBuilder().build_doc_nav( + [ + { + "chunk_id": "chunk_single_arrow_heading", + "type": "text", + "content": "single arrow heading content", + "path": chunk_path, + "metadata": {"summary": "single arrow heading summary"}, + } + ], + "report.pdf", + ) + + sections = cast(list[dict[str, object]], doc_nav["sections"]) + + assert ( + section_path_from_chunk_path( + chunk_path, + source_file_name="report.pdf", + ) + == "Inputs --> Outputs" + ) + assert sections[0]["title"] == "Inputs --> Outputs" + assert sections[0]["path"] == chunk_path + assert sections[0]["children"] == [] + + +def test_should_preserve_filename_like_arrow_text_below_document_root() -> None: + chunk_path = "report.pdf/input.csv-->output/Details" + + doc_nav = ZipResultSchemaBuilder().build_doc_nav( + [ + { + "chunk_id": "chunk_filename_arrow_heading", + "type": "text", + "content": "filename arrow heading content", + "path": chunk_path, + "metadata": {"summary": "filename arrow heading summary"}, + } + ], + "report.pdf", + ) + + sections = cast(list[dict[str, object]], doc_nav["sections"]) + children = cast(list[dict[str, object]], sections[0]["children"]) + + assert ( + section_path_from_chunk_path( + chunk_path, + source_file_name="report.pdf", + ) + == "input.csv-->output / Details" + ) + assert sections[0]["title"] == "input.csv-->output" + assert sections[0]["path"] == "report.pdf/input.csv-->output" + assert children[0]["title"] == "Details" + assert children[0]["path"] == chunk_path + + +def test_should_preserve_filename_like_section_titles_for_new_paths() -> None: + chunk_path = "report.pdf/appendix.pdf/Details" + + doc_nav = ZipResultSchemaBuilder().build_doc_nav( + [ + { + "chunk_id": "chunk_filename_section", + "type": "text", + "content": "filename-like section content", + "path": chunk_path, + "metadata": {"summary": "filename-like section summary"}, + } + ], + "report.pdf", + ) + + sections = cast(list[dict[str, object]], doc_nav["sections"]) + children = cast(list[dict[str, object]], sections[0]["children"]) + + assert ( + section_path_from_chunk_path( + chunk_path, + source_file_name="report.pdf", + ) + == "appendix.pdf / Details" + ) + assert sections[0]["title"] == "appendix.pdf" + assert sections[0]["path"] == "report.pdf/appendix.pdf" + assert children[0]["title"] == "Details" + assert children[0]["path"] == chunk_path diff --git a/apps/api/tests/contract/test_demo_documents_contract.py b/apps/api/tests/contract/test_demo_documents_contract.py index ea21e7de6..db850c4c9 100644 --- a/apps/api/tests/contract/test_demo_documents_contract.py +++ b/apps/api/tests/contract/test_demo_documents_contract.py @@ -136,13 +136,16 @@ async def test_should_materialize_demo_source_without_parse_or_credit_charge( monkeypatch: MonkeyPatch, ) -> None: fake_result_storage = FakeResultStorage() - monkeypatch.setattr( - "shared.services.storage.result_storage.get_result_storage", - lambda: fake_result_storage, - ) monkeypatch.setenv("RETRIEVAL_AGENTIC_ENABLED", "false") async with developer_api_client_factory() as api_client: + import app.services.demo.source_materializer as source_materializer_module + + monkeypatch.setattr( + source_materializer_module, + "get_result_storage", + lambda: fake_result_storage, + ) empty_cached_response = await api_client.post( "/api/v1/retrieval/query", json={ @@ -269,19 +272,82 @@ async def test_should_materialize_demo_source_without_parse_or_credit_charge( @pytest.mark.asyncio -async def test_should_serialize_concurrent_first_demo_materialization( +async def test_should_materialize_each_normalized_demo_source_once_per_request( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], monkeypatch: MonkeyPatch, ) -> None: fake_result_storage = FakeResultStorage() - monkeypatch.setattr( - "shared.services.storage.result_storage.get_result_storage", - lambda: fake_result_storage, + + async with developer_api_client_factory() as api_client: + import app.services.demo.source_materializer as source_materializer_module + + monkeypatch.setattr( + source_materializer_module, + "get_result_storage", + lambda: fake_result_storage, + ) + response = await api_client.post( + "/api/v1/demo/materializations", + json={ + "namespace": "contract-demo-deduplicate", + "demo_source_ids": [ + DEMO_SOURCE_ID, + f" {DEMO_SOURCE_ID} ", + DEMO_SOURCE_ID, + "", + ], + }, + ) + + assert response.status_code == 200 + + body = cast(dict[str, Any], response.json()) + sources = cast(list[dict[str, Any]], body["sources"]) + assert [source["demo_source_id"] for source in sources] == [DEMO_SOURCE_ID] + + materialization_rows = await ContractDatabase.fetch_all( + """ + SELECT demo_source_id, document_id + FROM demo_materializations + WHERE user_id = 'local-dev-user' + AND namespace = 'contract-demo-deduplicate' + """, + ) + job_rows = await ContractDatabase.fetch_all( + """ + SELECT job_id + FROM jobs + WHERE user_id = 'local-dev-user' + AND job_metadata ->> 'namespace' = 'contract-demo-deduplicate' + AND job_metadata ->> 'demo_source_id' = :demo_source_id + """, + {"demo_source_id": DEMO_SOURCE_ID}, ) + assert len(materialization_rows) == 1 + assert len(job_rows) == 1 + assert len(fake_result_storage.raw_files_by_job_id) == 1 + + +@pytest.mark.asyncio +async def test_should_serialize_concurrent_first_demo_materialization( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + fake_result_storage = FakeResultStorage() + async with developer_api_client_factory() as api_client: + import app.services.demo.source_materializer as source_materializer_module + + monkeypatch.setattr( + source_materializer_module, + "get_result_storage", + lambda: fake_result_storage, + ) first_response, second_response = await asyncio.gather( api_client.post( "/api/v1/demo/materializations", @@ -369,12 +435,15 @@ async def test_should_reject_mixed_demo_materialization_selection_before_upload( monkeypatch: MonkeyPatch, ) -> None: fake_result_storage = FakeResultStorage() - monkeypatch.setattr( - "shared.services.storage.result_storage.get_result_storage", - lambda: fake_result_storage, - ) async with developer_api_client_factory() as api_client: + import app.services.demo.source_materializer as source_materializer_module + + monkeypatch.setattr( + source_materializer_module, + "get_result_storage", + lambda: fake_result_storage, + ) response = await api_client.post( "/api/v1/demo/materializations", json={ diff --git a/apps/api/tests/contract/test_documents_contract.py b/apps/api/tests/contract/test_documents_contract.py index bb1f8dcad..d0b401839 100644 --- a/apps/api/tests/contract/test_documents_contract.py +++ b/apps/api/tests/contract/test_documents_contract.py @@ -310,7 +310,7 @@ async def _insert_document_revision_with_chunks( ) VALUES ( :job_id, :user_id, - 'kb_management', + 'document_ingestion', 'done', 'url', FALSE, diff --git a/apps/api/tests/contract/test_job_creation_contract.py b/apps/api/tests/contract/test_job_creation_contract.py index b6cfa1375..572488351 100644 --- a/apps/api/tests/contract/test_job_creation_contract.py +++ b/apps/api/tests/contract/test_job_creation_contract.py @@ -163,7 +163,7 @@ async def _insert_active_job( { "job_id": job_id, "user_id": user_id, - "job_type": "kb_management", + "job_type": "document_ingestion", "status": status, "source_type": "file", "webhook_enabled": False, @@ -244,7 +244,7 @@ async def test_should_create_a_waiting_file_job_for_an_authenticated_developer( original_request = cast(dict[str, object], job_metadata["original_request"]) assert job_row["user_id"] == "local-dev-user" - assert job_row["job_type"] == "kb_management" + assert job_row["job_type"] == "document_ingestion" assert job_row["status"] == "waiting-file" assert job_row["source_type"] == "file" assert job_row["s3_key"] == f"uploads/{job_id}.pdf" @@ -272,7 +272,7 @@ async def test_should_create_a_waiting_file_job_for_an_authenticated_developer( assert cached_metadata["source_file_name"] == payload["file_name"] assert cached_job_info["job_id"] == job_id assert cached_job_info["user_id"] == "local-dev-user" - assert cached_job_info["job_type"] == "kb_management" + assert cached_job_info["job_type"] == "document_ingestion" assert cached_job_info["source_type"] == "file" assert cached_job_info["s3_key"] == f"uploads/{job_id}.pdf" assert cached_job_info["webhook_enabled"] is False @@ -429,10 +429,12 @@ async def test_should_reject_authenticated_user_id_missing_from_user_table( } async with api_client_factory() as api_client: - from app.core import dependencies as auth_dependencies + from app.services.auth.dashboard_jwt_authentication_service import ( + get_dashboard_jwt_authentication_service, + ) monkeypatch.setattr( - auth_dependencies, + get_dashboard_jwt_authentication_service(), "_get_verification_key", lambda _token: jwt_secret, ) @@ -564,6 +566,44 @@ async def test_should_return_not_found_when_creating_a_job_for_an_archived_docum assert await _count_jobs() == 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("namespace_value", [None, "", " "]) +async def test_should_inherit_existing_document_namespace_when_update_namespace_is_blank( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + namespace_value: str | None, +) -> None: + document_id = f"doc_contract_{uuid4().hex[:12]}" + existing_namespace = "contract.jobs.custom" + payload: dict[str, object] = { + "document_id": document_id, + "namespace": namespace_value, + "source_type": "file", + "file_name": "replacement-upload.pdf", + "data_id": "contract-document-update", + } + + async with developer_api_client_factory() as api_client: + await _insert_document( + document_id=document_id, + namespace=existing_namespace, + ) + response = await api_client.post("/api/v1/jobs", json=payload) + + assert response.status_code == 200 + + response_json: dict[str, object] = response.json() + job_id = cast(str, response_json["job_id"]) + job_metadata = cast(dict[str, object], (await _load_job_record(job_id))["job_metadata"]) + original_request = cast(dict[str, object], job_metadata["original_request"]) + + assert response_json["namespace"] == existing_namespace + assert job_metadata["document_id"] == document_id + assert job_metadata["namespace"] == existing_namespace + assert original_request["namespace"] == namespace_value + + @pytest.mark.asyncio async def test_should_create_a_waiting_file_job_for_a_url_source_and_enqueue_the_upload_worker( monkeypatch: MonkeyPatch, @@ -623,11 +663,20 @@ def __init__(self) -> None: def signature(self, task_name: str) -> _FakeCeleryTask: return _FakeCeleryTask(task_name) + def resolve_public_address( + host: str, + port: int | None, + *args: object, + **kwargs: object, + ) -> list[tuple[socket.AddressFamily, socket.SocketKind, int, str, tuple[str, int]]]: + return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 0))] + import shared.core.celery_app as celery_app_module - import shared.utils.http_clients as http_clients_module + import shared.services.http.client_pool as client_pool_module + monkeypatch.setattr(socket, "getaddrinfo", resolve_public_address) monkeypatch.setattr( - http_clients_module, + client_pool_module, "get_async_client", lambda: _FakeAsyncHttpClient(), ) @@ -664,7 +713,7 @@ def signature(self, task_name: str) -> _FakeCeleryTask: original_request = cast(dict[str, object], job_metadata["original_request"]) assert job_row["user_id"] == "local-dev-user" - assert job_row["job_type"] == "kb_management" + assert job_row["job_type"] == "document_ingestion" assert job_row["status"] == "waiting-file" assert job_row["source_type"] == "url" assert job_row["s3_key"] == f"uploads/{job_id}.pdf" @@ -694,15 +743,15 @@ def signature(self, task_name: str) -> _FakeCeleryTask: assert cached_metadata["source_url"] == payload["source_url"] assert cached_job_info["job_id"] == job_id assert cached_job_info["user_id"] == "local-dev-user" - assert cached_job_info["job_type"] == "kb_management" + assert cached_job_info["job_type"] == "document_ingestion" assert cached_job_info["source_type"] == "url" assert cached_job_info["s3_key"] == f"uploads/{job_id}.pdf" assert cached_job_info["webhook_enabled"] is False assert scheduled_tasks == [ { - "task_name": "app.core.tasks.kb_tasks.upload_url_file_task", + "task_name": "app.core.tasks.document_ingestion_tasks.upload_url_file_task", "args": [job_id, payload["source_url"], "local-dev-user"], - "kwargs": {"job_type": "kb_management"}, + "kwargs": {"job_type": "document_ingestion"}, } ] @@ -831,9 +880,9 @@ def resolve_private_address( assert job_metadata["source_url"] == source_url assert scheduled_tasks == [ { - "task_name": "app.core.tasks.kb_tasks.upload_url_file_task", + "task_name": "app.core.tasks.document_ingestion_tasks.upload_url_file_task", "args": [job_id, source_url, "local-dev-user"], - "kwargs": {"job_type": "kb_management"}, + "kwargs": {"job_type": "document_ingestion"}, } ] @@ -917,9 +966,9 @@ async def head( assert follow_redirects is False return _FakeHeadResponse() - import shared.utils.http_clients as http_clients_module + import shared.services.http.client_pool as client_pool_module monkeypatch.setattr( - http_clients_module, + client_pool_module, "get_async_client", lambda: _FakeAsyncHttpClient(), ) @@ -971,29 +1020,22 @@ async def _fake_verify_s3_file_exists( assert bucket is None return {"exists": True, "s3_key": s3_key} - async def _fake_start_workflow_for_job( - db: object, + async def _fake_start_uploaded_file_parse( + self: object, + *, job_id: str, - job_type: str, - source_type: str, user_id: str, - file_path: str | None = None, - file_url: str | None = None, - ) -> None: + ) -> str: started_workflows.append( { "job_id": job_id, - "job_type": job_type, - "source_type": source_type, "user_id": user_id, - "file_path": file_path, - "file_url": file_url, - "db_bound": db is not None, } ) + return "contract-task-id" async with developer_api_client_factory() as api_client: - import app.api.v1.routes.jobs as jobs_route_module + import app.services.document_ingestion.worker_dispatcher as dispatcher_module import shared.services.storage.file_upload_service as file_upload_service_module monkeypatch.setattr( @@ -1002,9 +1044,9 @@ async def _fake_start_workflow_for_job( _fake_verify_s3_file_exists, ) monkeypatch.setattr( - jobs_route_module, - "start_workflow_for_job", - _fake_start_workflow_for_job, + dispatcher_module.DocumentIngestionWorkerDispatcher, + "start_uploaded_file_parse", + _fake_start_uploaded_file_parse, ) create_response = await api_client.post("/api/v1/jobs", json=payload) @@ -1026,11 +1068,87 @@ async def _fake_start_workflow_for_job( assert started_workflows == [ { "job_id": job_id, - "job_type": "kb_management", - "source_type": "file", "user_id": "local-dev-user", - "file_path": None, - "file_url": None, - "db_bound": True, } ] + + +@pytest.mark.asyncio +async def test_should_preserve_retryable_confirm_upload_transition_rejection( + monkeypatch: MonkeyPatch, + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + payload: dict[str, str] = { + "namespace": "contract-jobs", + "source_type": "file", + "file_name": "confirm-upload-retry.pdf", + "data_id": "contract-job-confirm-upload-retry", + } + + async def _fake_verify_s3_file_exists( + self: object, + s3_key: str, + bucket: str | None = None, + ) -> dict[str, object]: + assert bucket is None + return {"exists": True, "s3_key": s3_key} + + async def _fake_transition_outcome( + self: object, + db: object, + job_id: str, + to_state: str, + transition_reason: str = "normal_transition", + operator_id: str | None = None, + operator_type: str = "system", + metadata: dict[str, object] | None = None, + auto_commit: bool = True, + ) -> object: + del self, db, transition_reason, operator_id, operator_type, metadata, auto_commit + + from shared.core.state_machine.transition_outcome import JobTransitionOutcome + + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=to_state, + reason="cas_conflict", + attempts=3, + ) + + async with developer_api_client_factory() as api_client: + import shared.core.state_machine.service as state_machine_module + import shared.services.storage.file_upload_service as file_upload_service_module + + monkeypatch.setattr( + file_upload_service_module.FileUploadService, + "verify_s3_file_exists", + _fake_verify_s3_file_exists, + ) + monkeypatch.setattr( + state_machine_module.AsyncStateMachineService, + "transition_outcome", + _fake_transition_outcome, + ) + + create_response = await api_client.post("/api/v1/jobs", json=payload) + assert create_response.status_code == 200 + + create_response_json: dict[str, object] = create_response.json() + job_id = cast(str, create_response_json["job_id"]) + + confirm_response = await api_client.post(f"/api/v1/jobs/{job_id}/confirm-upload") + + assert confirm_response.status_code == 503 + assert confirm_response.headers["retry-after"] == "120" + assert confirm_response.headers["x-request-id"] + + response_json: dict[str, object] = confirm_response.json() + error = cast(dict[str, object], response_json["error"]) + details = cast(dict[str, object], error["details"]) + + assert response_json["success"] is False + assert error["code"] == "UNAVAILABLE" + assert error["message"] == "Job state is still settling. Retrying shortly." + assert details["retry_after"] == 120 diff --git a/apps/api/tests/contract/test_qstash_callback_contract.py b/apps/api/tests/contract/test_qstash_callback_contract.py index 9b6734d2d..c9e0c8f40 100644 --- a/apps/api/tests/contract/test_qstash_callback_contract.py +++ b/apps/api/tests/contract/test_qstash_callback_contract.py @@ -48,8 +48,10 @@ async def test_should_return_unauthorized_for_an_invalid_qstash_callback_signatu monkeypatch: MonkeyPatch, ) -> None: async with api_client_factory() as api_client: - qstash_module = importlib.import_module("app.api.v1.routes.qstash_callbacks") - monkeypatch.setattr(qstash_module, "_verify_qstash_signature", lambda *args: False) + qstash_module = importlib.import_module( + "app.services.webhook.qstash_callback_service" + ) + monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: False) response = await api_client.post( "/api/v1/webhooks/qstash/callback", json={"status": 200}, @@ -70,8 +72,10 @@ async def test_should_mark_the_matching_event_delivered_and_persist_a_webhook_lo async with api_client_factory() as api_client: job_id, event_id = await _insert_qstash_event() - qstash_module = importlib.import_module("app.api.v1.routes.qstash_callbacks") - monkeypatch.setattr(qstash_module, "_verify_qstash_signature", lambda *args: True) + qstash_module = importlib.import_module( + "app.services.webhook.qstash_callback_service" + ) + monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: True) response = await api_client.post( "/api/v1/webhooks/qstash/callback", json={ @@ -133,8 +137,10 @@ async def test_should_keep_the_matching_event_delivering_for_retry_callback_with status="delivering", qstash_message_id="qstash-message-retry", ) - qstash_module = importlib.import_module("app.api.v1.routes.qstash_callbacks") - monkeypatch.setattr(qstash_module, "_verify_qstash_signature", lambda *args: True) + qstash_module = importlib.import_module( + "app.services.webhook.qstash_callback_service" + ) + monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: True) response = await api_client.post( "/api/v1/webhooks/qstash/callback", json={ @@ -197,8 +203,10 @@ async def test_should_not_downgrade_terminal_event_when_retry_callback_arrives_l attempts=4, qstash_message_id="qstash-message-late-retry", ) - qstash_module = importlib.import_module("app.api.v1.routes.qstash_callbacks") - monkeypatch.setattr(qstash_module, "_verify_qstash_signature", lambda *args: True) + qstash_module = importlib.import_module( + "app.services.webhook.qstash_callback_service" + ) + monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: True) response = await api_client.post( "/api/v1/webhooks/qstash/callback", json={ @@ -257,8 +265,10 @@ async def test_should_mark_the_matching_event_failed_and_persist_the_error_on_fa async with api_client_factory() as api_client: job_id, event_id = await _insert_qstash_event() - qstash_module = importlib.import_module("app.api.v1.routes.qstash_callbacks") - monkeypatch.setattr(qstash_module, "_verify_qstash_signature", lambda *args: True) + qstash_module = importlib.import_module( + "app.services.webhook.qstash_callback_service" + ) + monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: True) response = await api_client.post( "/api/v1/webhooks/qstash/failure", json={ @@ -318,8 +328,10 @@ async def test_should_return_ok_without_mutating_state_when_the_callback_has_no_ async with api_client_factory() as api_client: _, event_id = await _insert_qstash_event() - qstash_module = importlib.import_module("app.api.v1.routes.qstash_callbacks") - monkeypatch.setattr(qstash_module, "_verify_qstash_signature", lambda *args: True) + qstash_module = importlib.import_module( + "app.services.webhook.qstash_callback_service" + ) + monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: True) response = await api_client.post( "/api/v1/webhooks/qstash/callback", json={ @@ -344,6 +356,7 @@ async def test_should_return_ok_without_mutating_state_when_the_callback_has_no_ ) assert event_row is not None + assert log_count_row is not None assert event_row["status"] == "pending" assert event_row["attempts"] == 0 assert cast(int, log_count_row["count"]) == 0 diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index a95ed4d72..06cbc173b 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -1,12 +1,18 @@ from collections.abc import Callable from contextlib import AbstractAsyncContextManager +from datetime import datetime, timezone from typing import cast from uuid import uuid4 import pytest from httpx import AsyncClient +from pytest import MonkeyPatch +from sqlalchemy.ext.asyncio import AsyncSession from tests.support.contract_database import ContractDatabase +from shared.services.retrieval.agentic.core.types import AgenticResult +from shared.services.retrieval.workflow.run_request import WorkflowRunRequest +from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, WorkflowResult async def _seed_retrieval_document( @@ -16,12 +22,13 @@ async def _seed_retrieval_document( source_file_name: str, section_path: str, content: str, + chunk_id: str | None = None, ) -> dict[str, str]: document_id = f"doc_{uuid4().hex[:12]}" job_id = f"job_{uuid4().hex[:12]}" job_result_id = str(uuid4()) section_id = f"sec_{uuid4().hex[:12]}" - chunk_id = f"chunk_{uuid4().hex[:12]}" + resolved_chunk_id = chunk_id or f"chunk_{uuid4().hex[:12]}" await ContractDatabase.insert_job( job_id=job_id, @@ -67,7 +74,7 @@ async def _seed_retrieval_document( section_title=section_path.split("/")[-1], ) await ContractDatabase.insert_document_chunk( - chunk_id=chunk_id, + chunk_id=resolved_chunk_id, user_id=user_id, namespace=namespace, document_id=document_id, @@ -83,11 +90,165 @@ async def _seed_retrieval_document( "job_id": job_id, "job_result_id": job_result_id, "section_id": section_id, + "chunk_id": resolved_chunk_id, + "section_path": section_path, + } + + +async def _seed_retrieval_chunk_for_existing_document( + *, + user_id: str, + namespace: str, + document: dict[str, str], + section_path: str, + content: str, + chunk_id: str, +) -> dict[str, str]: + section_id = f"sec_{uuid4().hex[:12]}" + + await ContractDatabase.insert_document_section( + section_id=section_id, + user_id=user_id, + namespace=namespace, + document_id=document["document_id"], + job_result_id=document["job_result_id"], + section_path=section_path, + section_title=section_path.split("/")[-1], + ) + await ContractDatabase.insert_document_chunk( + chunk_id=chunk_id, + user_id=user_id, + namespace=namespace, + document_id=document["document_id"], + job_result_id=document["job_result_id"], + section_id=section_id, + chunk_type="text", + content=content, + section_path=section_path, + ) + + return { + "document_id": document["document_id"], + "job_id": document["job_id"], + "job_result_id": document["job_result_id"], + "section_id": section_id, "chunk_id": chunk_id, "section_path": section_path, } +def _result_source(result: dict[str, object]) -> dict[str, object]: + return cast(dict[str, object], result["source"]) + +@pytest.mark.asyncio +async def test_agentic_workflow_should_pass_full_request_policy_to_step_adapter( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + captured_requests: list[dict[str, object]] = [] + + async def fake_plan( + self: object, + *, + query: str, + corpus_total_docs: int = 0, + corpus_total_chunks: int = 0, + ) -> QueryPlan: + return QueryPlan( + original_query=query, + steps=[PlannedStep(id="request-policy", sub_query="policy marker")], + final_strategy="concat_final_parts", + reasoning_summary=( + f"request policy contract for {corpus_total_docs} docs " + f"and {corpus_total_chunks} chunks" + ), + ) + + async def fake_retrieval_run( + self: object, + db: object, + **kwargs: object, + ) -> AgenticResult: + del self, db + captured_requests.append(kwargs) + return AgenticResult( + evidence_text="policy evidence", + answer_text="policy answer", + referenced_chunks=[ + { + "chunk_id": policy_document["chunk_id"], + "document_id": policy_document["document_id"], + "chunk_type": "text", + "section_path": policy_document["section_path"], + "file_path": "", + "job_id": policy_document["job_id"], + } + ], + router_used="contract_fake_agent", + ) + + async with developer_api_client_factory() as api_client: + policy_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-agentic-request-policy", + source_file_name="policy.pdf", + section_path="Root/Policy", + content="policy marker content", + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-agentic-request-policy", + source_file_name="filler.pdf", + section_path="Root/Filler", + content="filler content", + ) + monkeypatch.setattr( + "shared.services.retrieval.workflow.planner.QueryPlanner.plan", + fake_plan, + ) + monkeypatch.setattr( + "shared.services.retrieval.workflow.step_runner.RetrievalAgent.run", + fake_retrieval_run, + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-agentic-request-policy", + "query": "policy marker", + "top_k": 1, + "data_type": 2, + "signal_paths": ["Root"], + "filter_mode": "keep", + "channels": ["content"], + "channel_weights": {"content": 2.0}, + "internal_recall_k": 23, + "threshold": 0.4, + "rerank": True, + "use_agentic": True, + }, + ) + + assert response.status_code == 200 + + assert len(captured_requests) == 1 + request = captured_requests[0] + assert request["user_id"] == "local-dev-user" + assert request["namespace"] == "contract-agentic-request-policy" + assert request["query"] == "policy marker" + assert request["top_k"] == 1 + assert request["exclude_document_ids"] == [] + assert request["exclude_sections"] == [] + assert request["data_type"] == 2 + assert request["signal_paths"] == ["Root"] + assert request["filter_mode"] == "keep" + assert request["channels"] == ["content"] + assert request["channel_weights"] == {"content": 2.0} + assert request["internal_recall_k"] == 23 + + @pytest.mark.asyncio async def test_should_return_seeded_retrieval_results_for_the_authenticated_user( developer_api_client_factory: Callable[ @@ -119,7 +280,7 @@ async def test_should_return_seeded_retrieval_results_for_the_authenticated_user assert response_json["namespace"] == "contract-retrieval" assert response_json["query"] == "alpha" - assert response_json["router_used"] == "small_kb_all" + assert response_json["router_used"] == "small_corpus_all" assert len(results) == 1 assert results[0]["chunk_type"] == "text" assert results[0]["content"] == "alpha contract retrieval content" @@ -161,7 +322,7 @@ async def test_should_default_the_namespace_to_default_when_it_is_omitted( assert response_json["namespace"] == "default" assert len(results) == 1 - assert results[0]["source"]["document_id"] == seeded_document["document_id"] + assert _result_source(results[0])["document_id"] == seeded_document["document_id"] @pytest.mark.asyncio @@ -187,6 +348,630 @@ async def test_should_return_empty_results_for_an_empty_query( } +@pytest.mark.asyncio +async def test_legacy_retrieval_should_rank_hot_chunk_before_cold_chunk_when_discovery_scores_tie( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + async with developer_api_client_factory() as api_client: + cold_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-hot-ranking", + source_file_name="cold.pdf", + section_path="ranking/cold", + content="same ranking marker cold", + ) + hot_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-hot-ranking", + source_file_name="hot.pdf", + section_path="ranking/hot", + content="same ranking marker hot", + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-hot-ranking", + source_file_name="filler.pdf", + section_path="ranking/filler", + content="same ranking marker filler", + ) + + now = datetime.now(timezone.utc).replace(tzinfo=None) + await ContractDatabase.execute( + """ + INSERT INTO retrieval_hit_stats ( + id, + user_id, + namespace, + hit_kind, + document_id, + chunk_id, + hit_count, + last_hit_at, + created_at, + updated_at + ) VALUES ( + :id, + :user_id, + :namespace, + 'chunk', + :document_id, + :chunk_id, + :hit_count, + :now, + :now, + :now + ) + """, + { + "id": f"rhs_{uuid4().hex[:12]}", + "user_id": "local-dev-user", + "namespace": "contract-hot-ranking", + "document_id": hot_document["document_id"], + "chunk_id": hot_document["chunk_id"], + "hit_count": 100, + "now": now, + }, + ) + + def to_channel_row(document: dict[str, str]) -> dict[str, object]: + return { + "document_id": document["document_id"], + "chunk_id": document["chunk_id"], + "section_id": document["section_id"], + "section_path": document["section_path"], + "source_file_name": "cold.pdf" + if document["document_id"] == cold_document["document_id"] + else "hot.pdf", + "chunk_type": "text", + "content": "same ranking marker", + "score": 1.0, + "file_path": None, + "chunk_metadata": {}, + "job_result_id": document["job_result_id"], + "job_id": document["job_id"], + "sort_order": 0, + } + + async def fake_content_channel(*_args: object, **_kwargs: object) -> list[dict[str, object]]: + return [ + to_channel_row(hot_document), + to_channel_row(cold_document), + ] + + async def fake_path_channel(*_args: object, **_kwargs: object) -> list[dict[str, object]]: + return [ + to_channel_row(cold_document), + to_channel_row(hot_document), + ] + + async def fake_graph_routing(*_args: object, **_kwargs: object) -> list[dict[str, object]]: + return [] + + monkeypatch.setattr( + "shared.services.retrieval.execution.legacy_route.path_channel", + fake_path_channel, + ) + monkeypatch.setattr( + "shared.services.retrieval.execution.legacy_route.content_channel", + fake_content_channel, + ) + monkeypatch.setattr( + "shared.services.retrieval.execution.legacy_route.list_graph_routed_chunks", + fake_graph_routing, + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-hot-ranking", + "query": "same ranking marker", + "top_k": 1, + "channels": ["path", "content"], + "channel_weights": {"path": 1.0, "content": 1.0}, + "use_agentic": False, + }, + ) + + assert response.status_code == 200 + + response_json = cast(dict[str, object], response.json()) + results = cast(list[dict[str, object]], response_json["results"]) + + assert len(results) == 1 + assert _result_source(results[0])["document_id"] == hot_document["document_id"] + + +@pytest.mark.asyncio +async def test_agentic_retrieval_should_reference_root_only_document_content( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + monkeypatch.setenv("LLM_MOCK_ENABLED", "true") + + async with developer_api_client_factory() as api_client: + rooted_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-root-retrieval", + source_file_name="root-only.pdf", + section_path="Root", + content="root only diluted earnings marker content", + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-root-retrieval", + source_file_name="filler.pdf", + section_path="filler/section", + content="unrelated filler content", + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-root-retrieval", + "query": "diluted earnings marker", + "top_k": 1, + "use_agentic": True, + }, + ) + + assert response.status_code == 200 + + response_json = cast(dict[str, object], response.json()) + referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) + results = cast(list[dict[str, object]], response_json["results"]) + + assert response_json["router_used"] == "workflow_single_step" + assert { + "chunk_id": rooted_document["chunk_id"], + "document_id": rooted_document["document_id"], + "chunk_type": "text", + "section_path": "root-only.pdf", + "file_path": None, + "job_id": rooted_document["job_id"], + } in referenced_chunks + assert results[0]["content"] == "root only diluted earnings marker content" + assert results[0]["source"] == { + "document_id": rooted_document["document_id"], + "source_file_name": "root-only.pdf", + "section_path": "Root", + } + + +@pytest.mark.asyncio +async def test_agentic_retrieval_should_reference_discovery_content_when_navigation_selects_nothing( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + monkeypatch.setenv("LLM_MOCK_ENABLED", "true") + + async with developer_api_client_factory() as api_client: + discovered_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-discovery-fallback", + source_file_name="discovery.pdf", + section_path="Findings", + content="discovery fallback EBITDA margin marker content", + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-discovery-fallback", + source_file_name="filler.pdf", + section_path="filler/section", + content="unrelated filler content", + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-discovery-fallback", + "query": "EBITDA margin marker", + "top_k": 1, + "use_agentic": True, + }, + ) + + assert response.status_code == 200 + + response_json = cast(dict[str, object], response.json()) + referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) + results = cast(list[dict[str, object]], response_json["results"]) + + assert response_json["router_used"] == "workflow_single_step" + assert { + "chunk_id": discovered_document["chunk_id"], + "document_id": discovered_document["document_id"], + "chunk_type": "text", + "section_path": discovered_document["section_path"], + "file_path": None, + "job_id": discovered_document["job_id"], + } in referenced_chunks + assert results[0]["content"] == "discovery fallback EBITDA margin marker content" + assert results[0]["source"] == { + "document_id": discovered_document["document_id"], + "source_file_name": "discovery.pdf", + "section_path": discovered_document["section_path"], + } + + +@pytest.mark.asyncio +async def test_agentic_retrieval_should_not_hydrate_references_outside_request_scope( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + class FakeWorkflowOrchestrator: + async def run_request( + self, + _db: AsyncSession, + *, + request: WorkflowRunRequest, + ) -> WorkflowResult: + return WorkflowResult( + namespace=request.namespace, + query=request.query, + router_used="workflow_single_step", + answer_text="foreign reference answer", + referenced_chunks=[ + { + "chunk_id": foreign_document["chunk_id"], + "document_id": foreign_document["document_id"], + "chunk_type": "text", + "section_path": foreign_document["section_path"], + "file_path": None, + "job_id": foreign_document["job_id"], + } + ], + ) + + async with developer_api_client_factory() as api_client: + request_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-visible-scope", + source_file_name="visible.pdf", + section_path="visible/section", + content="visible scoped content", + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-visible-scope", + source_file_name="visible-filler.pdf", + section_path="visible/filler", + content="visible scoped filler content", + ) + foreign_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-foreign-scope", + source_file_name="foreign.pdf", + section_path="foreign/section", + content="foreign scoped content should not leak", + ) + monkeypatch.setattr( + "shared.services.retrieval.workflow.orchestrator.WorkflowOrchestrator", + FakeWorkflowOrchestrator, + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-visible-scope", + "query": "visible", + "top_k": 1, + "use_agentic": True, + }, + ) + + assert response.status_code == 200 + + response_json = cast(dict[str, object], response.json()) + referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) + results = cast(list[dict[str, object]], response_json["results"]) + + assert request_document["document_id"] != foreign_document["document_id"] + assert referenced_chunks == [] + assert results == [] + + +@pytest.mark.asyncio +async def test_agentic_retrieval_should_drop_references_that_do_not_match_the_hydrated_section( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + class FakeWorkflowOrchestrator: + async def run_request( + self, + _db: AsyncSession, + *, + request: WorkflowRunRequest, + ) -> WorkflowResult: + return WorkflowResult( + namespace=request.namespace, + query=request.query, + router_used="workflow_single_step", + answer_text="mismatched section answer", + referenced_chunks=[ + { + "chunk_id": visible_chunk["chunk_id"], + "document_id": visible_chunk["document_id"], + "chunk_type": "text", + "section_path": "wrong/section", + "file_path": None, + "job_id": visible_chunk["job_id"], + } + ], + ) + + async with developer_api_client_factory() as api_client: + visible_chunk = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-reference-section-match", + source_file_name="visible.pdf", + section_path="right/section", + content="visible scoped content", + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-reference-section-match", + source_file_name="filler.pdf", + section_path="filler/section", + content="filler content", + ) + monkeypatch.setattr( + "shared.services.retrieval.workflow.orchestrator.WorkflowOrchestrator", + FakeWorkflowOrchestrator, + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-reference-section-match", + "query": "visible", + "top_k": 1, + "use_agentic": True, + }, + ) + + assert response.status_code == 200 + + response_json = cast(dict[str, object], response.json()) + referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) + results = cast(list[dict[str, object]], response_json["results"]) + + assert referenced_chunks == [] + assert results == [] + + +@pytest.mark.asyncio +async def test_agentic_workflow_should_preserve_references_with_the_same_chunk_id_across_documents( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + shared_chunk_id = f"chunk_{uuid4().hex[:12]}" + + async def fake_plan( + self: object, + *, + query: str, + corpus_total_docs: int = 0, + corpus_total_chunks: int = 0, + ) -> QueryPlan: + return QueryPlan( + original_query=query, + steps=[ + PlannedStep(id="first", sub_query="first shared reference"), + PlannedStep(id="second", sub_query="second shared reference"), + ], + final_strategy="concat_final_parts", + reasoning_summary=( + f"forced two-step contract plan for {corpus_total_docs} docs " + f"and {corpus_total_chunks} chunks" + ), + ) + + async def fake_retrieval_run( + self: object, + db: object, + **kwargs: object, + ) -> AgenticResult: + query = str(kwargs["query"]) + document = first_document if query == "first shared reference" else second_document + return AgenticResult( + evidence_text=f"evidence for {document['document_id']}", + answer_text=f"answer for {document['document_id']}", + referenced_chunks=[ + { + "chunk_id": shared_chunk_id, + "document_id": document["document_id"], + "chunk_type": "text", + "section_path": document["section_path"], + "file_path": "", + "job_id": document["job_id"], + } + ], + router_used="contract_fake_agent", + ) + + async with developer_api_client_factory() as api_client: + first_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-shared-chunk-id", + source_file_name="first.pdf", + section_path="first/section", + content="shared deterministic content", + chunk_id=shared_chunk_id, + ) + second_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-shared-chunk-id", + source_file_name="second.pdf", + section_path="second/section", + content="shared deterministic content", + chunk_id=shared_chunk_id, + ) + monkeypatch.setattr( + "shared.services.retrieval.workflow.planner.QueryPlanner.plan", + fake_plan, + ) + monkeypatch.setattr( + "shared.services.retrieval.workflow.step_runner.RetrievalAgent.run", + fake_retrieval_run, + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-shared-chunk-id", + "query": "show both shared references", + "top_k": 1, + "use_agentic": True, + }, + ) + + assert response.status_code == 200 + + response_json = cast(dict[str, object], response.json()) + referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) + results = cast(list[dict[str, object]], response_json["results"]) + + referenced_document_ids = { + cast(str, reference["document_id"]) for reference in referenced_chunks + } + result_document_ids = { + cast(str, _result_source(result)["document_id"]) for result in results + } + + assert referenced_document_ids == { + first_document["document_id"], + second_document["document_id"], + } + assert result_document_ids == { + first_document["document_id"], + second_document["document_id"], + } + + +@pytest.mark.asyncio +async def test_agentic_workflow_should_preserve_references_with_the_same_chunk_id_across_sections( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + shared_chunk_id = f"chunk_{uuid4().hex[:12]}" + + async def fake_plan( + self: object, + *, + query: str, + corpus_total_docs: int = 0, + corpus_total_chunks: int = 0, + ) -> QueryPlan: + return QueryPlan( + original_query=query, + steps=[ + PlannedStep(id="first", sub_query="first shared section"), + PlannedStep(id="second", sub_query="second shared section"), + ], + final_strategy="concat_final_parts", + reasoning_summary=( + f"forced section identity contract plan for {corpus_total_docs} docs " + f"and {corpus_total_chunks} chunks" + ), + ) + + async def fake_retrieval_run( + self: object, + db: object, + **kwargs: object, + ) -> AgenticResult: + query = str(kwargs["query"]) + chunk = first_chunk if query == "first shared section" else second_chunk + return AgenticResult( + evidence_text=f"evidence for {chunk['section_path']}", + answer_text=f"answer for {chunk['section_path']}", + referenced_chunks=[ + { + "chunk_id": shared_chunk_id, + "document_id": chunk["document_id"], + "chunk_type": "text", + "section_path": chunk["section_path"], + "file_path": "", + "job_id": chunk["job_id"], + } + ], + router_used="contract_fake_agent", + ) + + async with developer_api_client_factory() as api_client: + first_chunk = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-shared-section-chunk-id", + source_file_name="same-document.pdf", + section_path="first/section", + content="repeated deterministic content", + chunk_id=shared_chunk_id, + ) + second_chunk = await _seed_retrieval_chunk_for_existing_document( + user_id="local-dev-user", + namespace="contract-shared-section-chunk-id", + document=first_chunk, + section_path="second/section", + content="repeated deterministic content", + chunk_id=shared_chunk_id, + ) + monkeypatch.setattr( + "shared.services.retrieval.workflow.planner.QueryPlanner.plan", + fake_plan, + ) + monkeypatch.setattr( + "shared.services.retrieval.workflow.step_runner.RetrievalAgent.run", + fake_retrieval_run, + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-shared-section-chunk-id", + "query": "show both shared section references", + "top_k": 1, + "use_agentic": True, + }, + ) + + assert response.status_code == 200 + + response_json = cast(dict[str, object], response.json()) + referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) + results = cast(list[dict[str, object]], response_json["results"]) + + referenced_section_paths = { + cast(str, reference["section_path"]) for reference in referenced_chunks + } + result_section_paths = { + cast(str, _result_source(result)["section_path"]) for result in results + } + + assert referenced_section_paths == { + first_chunk["section_path"], + second_chunk["section_path"], + } + assert result_section_paths == { + first_chunk["section_path"], + second_chunk["section_path"], + } + + @pytest.mark.asyncio async def test_should_return_request_validation_failure_for_an_invalid_channel( developer_api_client_factory: Callable[ @@ -255,7 +1040,7 @@ async def test_should_exclude_matching_document_ids_from_the_response( results = cast(list[dict[str, object]], response_json["results"]) assert len(results) == 1 - assert results[0]["source"]["document_id"] == included_document["document_id"] + assert _result_source(results[0])["document_id"] == included_document["document_id"] @pytest.mark.asyncio @@ -300,5 +1085,5 @@ async def test_should_exclude_matching_sections_from_the_response( results = cast(list[dict[str, object]], response_json["results"]) assert len(results) == 1 - assert results[0]["source"]["document_id"] == included_document["document_id"] - assert results[0]["source"]["section_path"] == included_document["section_path"] + assert _result_source(results[0])["document_id"] == included_document["document_id"] + assert _result_source(results[0])["section_path"] == included_document["section_path"] diff --git a/apps/api/tests/contract/test_s3_event_contract.py b/apps/api/tests/contract/test_s3_event_contract.py index 949722136..b1348c0cc 100644 --- a/apps/api/tests/contract/test_s3_event_contract.py +++ b/apps/api/tests/contract/test_s3_event_contract.py @@ -31,7 +31,9 @@ def _build_s3_event_payload(job_id: str) -> dict[str, object]: } -async def _insert_waiting_file_job() -> tuple[str, str]: +async def _insert_waiting_file_job( + *, job_type: str = "document_ingestion" +) -> tuple[str, str]: user_id = f"contract-s3-user-{uuid4().hex[:12]}" job_id = f"job_{uuid4().hex[:12]}" @@ -40,7 +42,7 @@ async def _insert_waiting_file_job() -> tuple[str, str]: job_id=job_id, user_id=user_id, status="waiting-file", - job_type="kb_management", + job_type=job_type, source_type="file", s3_key=f"uploads/{job_id}.pdf", ) @@ -67,34 +69,89 @@ async def test_should_accept_a_direct_upload_complete_event_advance_the_waiting_ api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], monkeypatch: MonkeyPatch, ) -> None: - workflow_calls: list[dict[str, str | None]] = [] + workflow_calls: list[dict[str, str]] = [] user_id: str = "" job_id: str = "" - class FakeKBOrchestrator: - async def start_workflow( + class FakeDocumentIngestionWorkerDispatcher: + async def start_uploaded_file_parse( self, - db, + *, job_id: str, - source_type: str, - file_path: str | None, - file_url: str | None, user_id: str, - ) -> None: + ) -> str: workflow_calls.append( { "job_id": job_id, - "source_type": source_type, - "file_path": file_path, - "file_url": file_url, "user_id": user_id, } ) + return "contract-task-id" async with api_client_factory() as api_client: user_id, job_id = await _insert_waiting_file_job() - s3_events_module = importlib.import_module("app.api.v1.routes.s3_events") - monkeypatch.setattr(s3_events_module, "KBOrchestrator", FakeKBOrchestrator) + handoff_service = importlib.import_module( + "app.services.document_ingestion.handoff_service" + ) + monkeypatch.setattr( + handoff_service, + "DocumentIngestionWorkerDispatcher", + FakeDocumentIngestionWorkerDispatcher, + ) + response = await api_client.post( + "/api/v1/internal/s3-events", + json=_build_s3_event_payload(job_id), + ) + + assert response.status_code == 200 + assert response.json() == {"message": "Event handled successfully"} + + job_row = await ContractDatabase.fetch_job(job_id) + + assert job_row is not None + assert job_row["status"] == "pending" + assert workflow_calls == [ + { + "job_id": job_id, + "user_id": user_id, + } + ] + + +@pytest.mark.asyncio +async def test_should_accept_a_pre_rename_waiting_file_job_type_during_upload_handoff( + api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], + monkeypatch: MonkeyPatch, +) -> None: + workflow_calls: list[dict[str, str]] = [] + user_id: str = "" + job_id: str = "" + + class FakeDocumentIngestionWorkerDispatcher: + async def start_uploaded_file_parse( + self, + *, + job_id: str, + user_id: str, + ) -> str: + workflow_calls.append( + { + "job_id": job_id, + "user_id": user_id, + } + ) + return "contract-task-id" + + async with api_client_factory() as api_client: + user_id, job_id = await _insert_waiting_file_job(job_type="kb_management") + handoff_service = importlib.import_module( + "app.services.document_ingestion.handoff_service" + ) + monkeypatch.setattr( + handoff_service, + "DocumentIngestionWorkerDispatcher", + FakeDocumentIngestionWorkerDispatcher, + ) response = await api_client.post( "/api/v1/internal/s3-events", json=_build_s3_event_payload(job_id), @@ -110,9 +167,6 @@ async def start_workflow( assert workflow_calls == [ { "job_id": job_id, - "source_type": "file", - "file_path": None, - "file_url": None, "user_id": user_id, } ] @@ -125,22 +179,25 @@ async def test_should_accept_an_sns_wrapped_upload_complete_event_and_advance_a_ ) -> None: job_id: str = "" - class FakeKBOrchestrator: - async def start_workflow( + class FakeDocumentIngestionWorkerDispatcher: + async def start_uploaded_file_parse( self, - db, + *, job_id: str, - source_type: str, - file_path: str | None, - file_url: str | None, user_id: str, - ) -> None: - return None + ) -> str: + return "contract-task-id" async with api_client_factory() as api_client: _, job_id = await _insert_waiting_file_job() - s3_events_module = importlib.import_module("app.api.v1.routes.s3_events") - monkeypatch.setattr(s3_events_module, "KBOrchestrator", FakeKBOrchestrator) + handoff_service = importlib.import_module( + "app.services.document_ingestion.handoff_service" + ) + monkeypatch.setattr( + handoff_service, + "DocumentIngestionWorkerDispatcher", + FakeDocumentIngestionWorkerDispatcher, + ) response = await api_client.post( "/api/v1/internal/s3-events", content=json.dumps( @@ -198,7 +255,7 @@ def get(self, url: str, *args: object, **kwargs: object) -> object: async with api_client_factory() as api_client: monkeypatch.setattr(socket, "getaddrinfo", resolve_private_address) pinned_http_module = importlib.import_module( - "shared.utils.pinned_outbound_http" + "shared.services.http.pinned_outbound" ) monkeypatch.setattr( pinned_http_module.aiohttp, diff --git a/apps/api/tests/contract/test_webhook_contract.py b/apps/api/tests/contract/test_webhook_contract.py index 92f064d3c..890379af1 100644 --- a/apps/api/tests/contract/test_webhook_contract.py +++ b/apps/api/tests/contract/test_webhook_contract.py @@ -1,6 +1,7 @@ import importlib from collections.abc import Callable from contextlib import AbstractAsyncContextManager +from types import SimpleNamespace from typing import cast from uuid import uuid4 @@ -163,11 +164,15 @@ async def test_should_trigger_a_webhook_for_an_owned_terminal_job_with_a_matchin event_id: str = "" class FakeDispatcher: - async def _send_webhook(self, db, event, is_manual: bool = False): - assert is_manual is True + async def send_manual_webhook(self, db, event): assert event.id == event_id assert event.job_id == job_id - return True, 202, 118, None + return SimpleNamespace( + success=True, + status_code=202, + duration_ms=118, + error_message=None, + ) async with developer_api_client_factory() as api_client: job_id = await _insert_webhook_job(user_id="local-dev-user", status="done") diff --git a/apps/api/tests/migrations/test_schema_contract.py b/apps/api/tests/migrations/test_schema_contract.py index 7cbf27eb7..6d1a9db6a 100644 --- a/apps/api/tests/migrations/test_schema_contract.py +++ b/apps/api/tests/migrations/test_schema_contract.py @@ -88,7 +88,7 @@ def _insert_job( { "job_id": job_id, "user_id": user_id, - "job_type": "kb_management", + "job_type": "document_ingestion", "status": status, "source_type": "file", "webhook_enabled": False, diff --git a/apps/api/tests/support/contract_database.py b/apps/api/tests/support/contract_database.py index 664fff00a..2944cf14a 100644 --- a/apps/api/tests/support/contract_database.py +++ b/apps/api/tests/support/contract_database.py @@ -356,7 +356,7 @@ async def insert_job( *, job_id: str, user_id: str, - job_type: str = "kb_management", + job_type: str = "document_ingestion", status: str = "pending", source_type: str = "file", file_path: str | None = None, diff --git a/apps/docs/test_github_flow.md b/apps/docs/test_github_flow.md deleted file mode 100644 index 84fd82cd5..000000000 --- a/apps/docs/test_github_flow.md +++ /dev/null @@ -1,4 +0,0 @@ -# GitHub Flow Test - -This is a test file to verify the GitHub workflow simulation (Issue -> Branch -> PR -> Merge). -It is safe to ignore or delete this file later. diff --git a/apps/worker/app/core/tasks/__init__.py b/apps/worker/app/core/tasks/__init__.py index ba382c975..35cdd963b 100644 --- a/apps/worker/app/core/tasks/__init__.py +++ b/apps/worker/app/core/tasks/__init__.py @@ -1 +1 @@ -"""Celery task package for worker-side knowledge-base jobs.""" +"""Celery task package for worker-side Document Ingestion jobs.""" diff --git a/apps/worker/app/core/tasks/base_task.py b/apps/worker/app/core/tasks/base_task.py index 4a9f65543..c3474884f 100644 --- a/apps/worker/app/core/tasks/base_task.py +++ b/apps/worker/app/core/tasks/base_task.py @@ -1,5 +1,5 @@ """ -Base Celery task class for KB worker tasks. +Base Celery task class for worker-side Document Ingestion tasks. Provides centralized exception handling with direct DB writes for failure finalization. """ @@ -13,15 +13,15 @@ from shared.core.logging import LogEvent -class KBBaseTask(Task): - """Knowledge Base base task class - provides centralized exception handling.""" +class DocumentIngestionBaseTask(Task): + """Base task class for worker-side Document Ingestion error handling.""" def on_success(self, retval, task_id, args, kwargs): """Task success callback.""" logger.bind( event=LogEvent.WORKER_TASK_COMPLETE.value, task_id=task_id, - ).info("KB task completed successfully") + ).info("Document Ingestion task completed successfully") def on_failure(self, exc, task_id, args, kwargs, einfo): """Task failure callback — finalize failure directly to the database.""" @@ -44,7 +44,7 @@ def on_failure(self, exc, task_id, args, kwargs, einfo): # Finalize failure directly to the database. if job_id: try: - from shared.services.job_lifecycle_sync import ( + from shared.services.jobs.lifecycle.service import ( get_sync_job_lifecycle_service, ) @@ -84,4 +84,4 @@ def on_retry(self, exc, task_id, args, kwargs, einfo): task_id=task_id, job_id=job_id, retry_count=self.request.retries, - ).warning(f"KB task retrying: {exc}") + ).warning(f"Document Ingestion task retrying: {exc}") diff --git a/apps/worker/app/core/tasks/document_ingestion_tasks.py b/apps/worker/app/core/tasks/document_ingestion_tasks.py new file mode 100644 index 000000000..1683a05f0 --- /dev/null +++ b/apps/worker/app/core/tasks/document_ingestion_tasks.py @@ -0,0 +1,170 @@ +""" +Document Ingestion Celery Tasks + +Sync implementation for gevent worker pool. +All I/O operations use sync services that yield cooperatively under gevent. +""" + +from app.core.tasks.base_task import DocumentIngestionBaseTask +from app.services.document_ingestion.service import parse_uploaded_file_job +from app.services.workload.url_upload_service import upload_url_file +from loguru import logger + +from shared.core.celery_app import get_celery_app +from shared.core.config import settings +from shared.core.exceptions import RETRYABLE_EXCEPTIONS + +# Exception handling +from shared.core.exceptions.domain_exceptions import ( + WorkerHandlingException, +) +from shared.core.logging import LogEvent, log_context + +# Get Celery application +celery_app = get_celery_app() + + +_DOCUMENT_INGESTION_JOB_TYPE = "document_ingestion" +_LEGACY_DOCUMENT_INGESTION_JOB_TYPE = "kb_management" + + +@celery_app.task( + bind=True, + base=DocumentIngestionBaseTask, + name="app.core.tasks.document_ingestion_tasks.upload_url_file_task", + ignore_result=True, + autoretry_for=RETRYABLE_EXCEPTIONS, + retry_kwargs={ + "countdown": settings.DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN, + "max_retries": settings.DOCUMENT_INGESTION_TASK_MAX_RETRIES, + }, +) +def upload_url_file_task( + self, + job_id: str, + source_url: str, + user_id: str | None = None, + job_type: str | None = None, +): + """Download file from URL and upload to S3.""" + return _run_upload_url_file_task(self, job_id, source_url, user_id, job_type) + + +@celery_app.task( + bind=True, + base=DocumentIngestionBaseTask, + name="app.core.tasks.kb_tasks.upload_url_file_task", + ignore_result=True, + autoretry_for=RETRYABLE_EXCEPTIONS, + retry_kwargs={ + "countdown": settings.DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN, + "max_retries": settings.DOCUMENT_INGESTION_TASK_MAX_RETRIES, + }, +) +def legacy_upload_url_file_task( + self, + job_id: str, + source_url: str, + user_id: str | None = None, + job_type: str | None = _LEGACY_DOCUMENT_INGESTION_JOB_TYPE, +): + """Drain pre-rename URL-upload messages from legacy broker queues.""" + return _run_upload_url_file_task(self, job_id, source_url, user_id, job_type) + + +def _run_upload_url_file_task( + self, + job_id: str, + source_url: str, + user_id: str | None = None, + job_type: str | None = None, +): + with log_context(task_id=self.request.id): + if not job_id: + raise WorkerHandlingException( + user_message="An unexpected system error occurred", + internal_message="Worker task 'upload_url_file_task' called without job_id", + ) + + result = _upload_url_file(job_id, source_url, user_id, job_type) + + logger.bind(event=LogEvent.WORKER_TASK_COMPLETE.value).info( + "Task completed: upload_url_file_task" + ) + return result + + +def _upload_url_file( + job_id: str, source_url: str, user_id: str | None, job_type: str | None = None +): + """Sync URL file download and upload to S3.""" + return upload_url_file(job_id, source_url, user_id, job_type) + + +@celery_app.task( + bind=True, + base=DocumentIngestionBaseTask, + name="app.core.tasks.document_ingestion_tasks.parse_task", + ignore_result=True, + autoretry_for=RETRYABLE_EXCEPTIONS, + retry_kwargs={ + "countdown": settings.DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN, + "max_retries": settings.DOCUMENT_INGESTION_TASK_MAX_RETRIES, + }, +) +def parse_task( + self, + job_id: str, + user_id: str | None = None, + job_type: str = _DOCUMENT_INGESTION_JOB_TYPE, +): + """Parse and vectorize task (file already uploaded to S3).""" + return _run_parse_task(self, job_id, user_id, job_type) + + +@celery_app.task( + bind=True, + base=DocumentIngestionBaseTask, + name="app.core.tasks.kb_tasks.parse_task", + ignore_result=True, + autoretry_for=RETRYABLE_EXCEPTIONS, + retry_kwargs={ + "countdown": settings.DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN, + "max_retries": settings.DOCUMENT_INGESTION_TASK_MAX_RETRIES, + }, +) +def legacy_parse_task( + self, + job_id: str, + user_id: str | None = None, + job_type: str = _LEGACY_DOCUMENT_INGESTION_JOB_TYPE, +): + """Drain pre-rename parse messages from legacy broker queues.""" + return _run_parse_task(self, job_id, user_id, job_type) + + +def _run_parse_task( + self, + job_id: str, + user_id: str | None = None, + job_type: str = _DOCUMENT_INGESTION_JOB_TYPE, +): + del job_type + with log_context(task_id=self.request.id): + if not job_id: + raise WorkerHandlingException( + user_message="An unexpected system error occurred", + internal_message="Worker task 'parse_task' called without job_id", + ) + + result = _parse(job_id, user_id) + + logger.bind(event=LogEvent.WORKER_TASK_COMPLETE.value).info( + "Task completed: parse_task" + ) + return result + + +def _parse(job_id: str, user_id: str | None): + """Sync parse and vectorize (file already uploaded to S3).""" + return parse_uploaded_file_job(job_id, user_id) diff --git a/apps/worker/app/core/tasks/kb_tasks.py b/apps/worker/app/core/tasks/kb_tasks.py deleted file mode 100644 index 19acc3605..000000000 --- a/apps/worker/app/core/tasks/kb_tasks.py +++ /dev/null @@ -1,712 +0,0 @@ -""" -Knowledge Base Management Celery Tasks - -Sync implementation for gevent worker pool. -All I/O operations use sync services that yield cooperatively under gevent. -""" - -import os -from datetime import datetime, timezone - -import pandas as pd - -# Base task class -from app.core.tasks.base_task import KBBaseTask -from app.core.tasks.task_utils import ( - cleanup_task_workspace, - create_task_workspace, - download_s3_file_to_temp, -) -from app.services.common.job_start_service import mark_job_running -from app.services.document_parser.stage_profiler import stage_timer - -# Storage operations -from app.services.storage.sync_storage_service import ( - download_file_from_url, - generate_download_url, - upload_to_s3, - verify_s3_file_exists, -) -from app.services.workload.page_estimator import PageEstimator -from loguru import logger -from sqlalchemy import select - -from shared.core.celery_app import get_celery_app -from shared.core.config import settings -from shared.core.database_sync import get_sync_db_context -from shared.core.exceptions import RETRYABLE_EXCEPTIONS - -# Exception handling -from shared.core.exceptions.domain_exceptions import ( - InsufficientCreditsException, - NotFoundException, - StorageServiceException, - ValidationException, - WorkerHandlingException, -) -from shared.core.logging import LogEvent, log_context -from shared.models.database.job import Job - -# Domain services -from shared.models.schemas.job_metadata import JobMetadataHelper -from shared.services.billing.work_billing_service import WorkBillingService -from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks -from shared.services.job_lifecycle_sync import get_sync_job_lifecycle_service -from shared.services.redis.distributed_lock import RedisJobLock - -# Sync services for gevent worker -from shared.services.redis.redis_sync_service import ( - SyncJobInfoRedisService, - SyncJobMetadataService, - SyncRedisServiceFactory, -) -from shared.services.storage.result_storage import get_result_storage -from shared.services.storage.zip_result_service import ZipResultService -from app.services.connect_builder.summary_builder import ( - build_section_summary_lookup, - enrich_doc_nav_summaries, - ensure_doc_nav_json, - load_nav_top_summary, -) - -# Get Celery application -celery_app = get_celery_app() - - - -@celery_app.task( - bind=True, - base=KBBaseTask, - name="app.core.tasks.kb_tasks.upload_url_file_task", - ignore_result=True, - autoretry_for=RETRYABLE_EXCEPTIONS, - retry_kwargs={ - "countdown": settings.KB_TASK_RETRY_COUNTDOWN, - "max_retries": settings.KB_TASK_MAX_RETRIES, - }, -) -def upload_url_file_task( - self, - job_id: str, - source_url: str, - user_id: str | None = None, - job_type: str | None = None, -): - """Download file from URL and upload to S3.""" - with log_context(task_id=self.request.id): - if not job_id: - raise WorkerHandlingException( - user_message="An unexpected system error occurred", - internal_message="Worker task 'upload_url_file_task' called without job_id", - ) - - result = _upload_url_file(job_id, source_url, user_id, job_type) - - logger.bind(event=LogEvent.WORKER_TASK_COMPLETE.value).info( - "Task completed: upload_url_file_task" - ) - return result - - -def _upload_url_file( - job_id: str, source_url: str, user_id: str | None, job_type: str | None = None -): - """Sync URL file download and upload to S3.""" - lifecycle_service = get_sync_job_lifecycle_service() - - # Get job info from Redis - redis_service = SyncRedisServiceFactory.get_service() - job_info_service = SyncJobInfoRedisService(redis_service) - job_info = job_info_service.get_job_info(job_id) - - if not job_info: - metadata_service = SyncJobMetadataService(redis_service) - job_metadata = metadata_service.get_metadata(job_id) - if job_metadata: - s3_key = job_metadata.get("s3_key") - else: - raise NotFoundException( - resource="JobInfo", - resource_id=job_id, - internal_message="Job info not found in Redis or Metadata", - ) - else: - s3_key = job_info.get("s3_key") - - if not s3_key: - raise NotFoundException( - resource="JobInfo", - resource_id="s3_key", - internal_message=f"Missing s3_key in Redis job info for job_id={job_id}", - ) - - # Publish progress: validating file type - lifecycle_service.update_progress( - job_id, progress=3, message="Validating URL file type..." - ) - - # Step 1: Validate URL file type (path first, then Content-Type header) - from shared.utils.url_file_type import resolve_file_extension_sync - - file_extension = resolve_file_extension_sync(source_url) - - if not file_extension: - all_supported_extensions = settings.get_supported_extensions() - supported_formats = ", ".join(sorted(all_supported_extensions)) - raise ValidationException( - user_message="Unsupported file type", - violations=[ - { - "field": "file_extension", - "description": f"Must be one of: {supported_formats}", - } - ], - ) - - # Publish progress: downloading - lifecycle_service.update_progress( - job_id, progress=10, message="Downloading file from URL..." - ) - - # Step 2: Download file to temp directory - try: - temp_file_path = download_file_from_url(source_url) - except Exception as e: - raise ValidationException( - user_message="Failed to download file from URL", - violations=[ - { - "field": "source_url", - "description": "Could not download file from the provided URL", - } - ], - internal_message=f"Failed to download file from URL: {source_url}, error: {e}", - ) - - try: - # Publish progress: validating file size - lifecycle_service.update_progress( - job_id, progress=30, message="Validating file size..." - ) - - # Step 3: Validate file size - file_size = os.path.getsize(temp_file_path) - - if file_size > settings.MAX_FILE_SIZE: - limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024) - raise ValidationException( - user_message=f"File size exceeds limit (max {limit_mb}MB for {file_extension})", - violations=[ - { - "field": "file_size", - "description": f"Size {file_size} bytes exceeds limit of {settings.MAX_FILE_SIZE} bytes", - } - ], - ) - - # Publish progress: uploading to S3 - lifecycle_service.update_progress( - job_id, progress=50, message="Uploading file to S3..." - ) - - # Step 4: Upload to S3 - uploads_bucket = settings.S3_BUCKET_NAME - upload_to_s3(temp_file_path, s3_key, uploads_bucket) - logger.info(f"File uploaded to S3: {s3_key}") - - finally: - if os.path.exists(temp_file_path): - os.remove(temp_file_path) - logger.debug(f"Temp file cleaned up: {temp_file_path}") - - # Publish progress: verifying upload - lifecycle_service.update_progress( - job_id, progress=80, message="Verifying upload result..." - ) - - # Step 5: Verify S3 file exists - file_info = verify_s3_file_exists(s3_key) - if not file_info.get("exists"): - raise StorageServiceException( - user_message="We failed to verify your file upload", - internal_message=f"S3 file verification failed for {s3_key}", - ) - - # Publish progress: complete - lifecycle_service.update_progress( - job_id, - progress=100, - message="URL file upload complete, waiting for processing...", - ) - - logger.info( - f"URL file upload complete, waiting for S3 webhook: {job_id} -> {s3_key}" - ) - - return { - "status": "success", - "job_id": job_id, - "s3_key": s3_key, - "file_size": file_info.get("size"), - } - - -@celery_app.task( - bind=True, - base=KBBaseTask, - name="app.core.tasks.kb_tasks.parse_task", - ignore_result=True, - autoretry_for=RETRYABLE_EXCEPTIONS, - retry_kwargs={ - "countdown": settings.KB_TASK_RETRY_COUNTDOWN, - "max_retries": settings.KB_TASK_MAX_RETRIES, - }, -) -def parse_task( - self, job_id: str, user_id: str | None = None, job_type: str = "kb_management" -): - """Parse and vectorize task (file already uploaded to S3).""" - with log_context(task_id=self.request.id): - if not job_id: - raise WorkerHandlingException( - user_message="An unexpected system error occurred", - internal_message="Worker task 'parse_task' called without job_id", - ) - - result = _parse(job_id, user_id) - - logger.bind(event=LogEvent.WORKER_TASK_COMPLETE.value).info( - "Task completed: parse_task" - ) - return result - - -def _parse(job_id: str, user_id: str | None): - """Sync parse and vectorize (file already uploaded to S3).""" - logger.info(f"Parse started: job_id={job_id}, user_id={user_id}") - lifecycle_service = get_sync_job_lifecycle_service() - - # Get job info from Redis (sync) - redis_service = SyncRedisServiceFactory.get_service() - job_info_service = SyncJobInfoRedisService(redis_service) - job_info = job_info_service.get_job_info(job_id) - - if not job_info: - # Redis JobInfo has expired or been flushed — fall back to the DB, which is - # the durable source of truth for s3_key and user_id written at job creation. - logger.warning( - f"JobInfo not found in Redis for job_id={job_id}; falling back to database" - ) - with get_sync_db_context() as fallback_db: - job_row = fallback_db.execute( - select(Job).where(Job.job_id == job_id) - ).scalar_one_or_none() - - if not job_row or not job_row.s3_key: - raise NotFoundException( - resource="JobInfo", - resource_id=job_id, - internal_message="job info not found in Redis or database", - ) - - s3_key: str = job_row.s3_key - job_user_id: str | None = str(job_row.user_id) if job_row.user_id else user_id - logger.info( - f"Recovered JobInfo from database: job_id={job_id}, s3_key={s3_key}" - ) - else: - raw_s3_key = job_info.get("s3_key") - if not isinstance(raw_s3_key, str) or not raw_s3_key: - raise NotFoundException( - resource="JobInfo", - resource_id="s3_key", - internal_message="Missing s3_key in job_info", - ) - - s3_key = raw_s3_key - raw_job_user_id = job_info.get("user_id") - job_user_id = raw_job_user_id if isinstance(raw_job_user_id, str) else user_id - - # Verify S3 file exists (sync) - file_info = verify_s3_file_exists(s3_key) - if not file_info.get("exists"): - raise NotFoundException( - resource="S3File", - resource_id=s3_key, - internal_message=f"S3 file not found: {s3_key}", - ) - - logger.info(f"S3 file verified: {s3_key}") - - # Validate file size - file_size = file_info.get("size", 0) - file_extension = os.path.splitext(s3_key)[1].lower() - - if file_size > settings.MAX_FILE_SIZE: - limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024) - raise ValidationException( - user_message=f"File size exceeds limit (max {limit_mb}MB for {file_extension})", - violations=[ - { - "field": "file_size", - "description": f"Size {file_size} bytes exceeds limit of {settings.MAX_FILE_SIZE} bytes", - } - ], - ) - - # Get job_metadata from Redis - metadata_service = SyncJobMetadataService(redis_service) - job_metadata = metadata_service.get_metadata(job_id) - if not job_metadata: - raise NotFoundException( - resource="JobMetadata", - resource_id=job_id, - internal_message=f"Job metadata not found for job_id={job_id}", - ) - - should_process = mark_job_running(job_id, redis_service) - if not should_process: - logger.warning(f"Skipping parse_task for inactive job: job_id={job_id}") - return { - "status": "skipped", - "job_id": job_id, - "reason": "job_already_terminal", - } - - # Acquire distributed lock to prevent concurrent processing of the same - # job when the broker redelivers a task before the original worker acks. - # If another worker already holds the lock, UnavailableException is raised - # and Celery auto-retries after KB_TASK_RETRY_COUNTDOWN seconds. - with RedisJobLock(redis_service, job_id): - task_workspace_dir = create_task_workspace(job_id) - input_dir = os.path.join(task_workspace_dir, "input") - output_dir = os.path.join(task_workspace_dir, "output") - os.makedirs(input_dir, exist_ok=True) - os.makedirs(output_dir, exist_ok=True) - logger.info( - f"Task workspace ready: job_id={job_id}, workspace={task_workspace_dir}" - ) - - try: - # Publish progress: start parsing - lifecycle_service.update_progress( - job_id, progress=10, message="Parsing document..." - ) - - # Generate download URL and download file (sync) - file_url_response = generate_download_url(s3_key, settings.S3_BUCKET_NAME) - file_url = file_url_response["download_url"] - - filename = JobMetadataHelper.get_field(job_metadata, "source_file_name") - - # Download file to the task workspace - page_count = 1 - - # Derive file extension from s3_key (always has the correct extension) - # rather than filename, which may not have a real extension for URLs - # like arxiv.org/pdf/1706.03762 - file_ext = os.path.splitext(s3_key)[1].lower() if s3_key else "" - local_temp_path = download_s3_file_to_temp(file_url, file_ext, input_dir) - - logger.info( - f"File downloaded: job_id={job_id}, local_path={local_temp_path}" - ) - - from app.services.document_parser.internal_parse_name import ( - prepare_internal_parse_input, - ) - from app.services.document_parser.parse_service import ( - checkerboard_inject_parse, - ) - - prepared_parse_input = prepare_internal_parse_input( - local_temp_path, - filename, - fallback_ext=file_ext, - prefer_fallback_ext=True, - ) - internal_parse_name = prepared_parse_input.internal_filename - local_temp_path = prepared_parse_input.file_path - logger.info( - f"File prepared for parsing: job_id={job_id}, " - f"internal_filename={internal_parse_name}, local_path={local_temp_path}" - ) - - # Estimate workload - page_count = PageEstimator.estimate(local_temp_path) - logger.info( - f"Workload estimation: job_id={job_id}, page_count={page_count}" - ) - - processing_started_at = datetime.now(timezone.utc) - - if not job_user_id: - raise NotFoundException( - resource="JobInfo", - resource_id="user_id", - internal_message=f"Missing user_id in job info for job_id={job_id}", - ) - - billing_service = WorkBillingService() - billing_status = "skipped" - billing_amount_micro_dollars = 0 - billing_credits = 0.0 - with get_sync_db_context() as db: - job_result = db.execute( - select(Job).where(Job.job_id == job_id).with_for_update() - ) - job = job_result.scalar_one_or_none() - - if job and getattr(job, "billing_status", "") == "charged": - logger.info(f"Job already charged: {job_id}") - billing_status = "charged" - billing_amount_micro_dollars = int(job.credits_charged or 0) - billing_credits = billing_amount_micro_dollars / 1_000_000 - else: - try: - billing_result = billing_service.charge_for_pages( - session=db, - user_id=job_user_id, - page_count=page_count, - filename=filename, - ) - except InsufficientCreditsException: - logger.warning( - f"Billing failed: job_id={job_id}, user_id={job_user_id}" - ) - billing_amount = billing_service.estimate_page_charge( - page_count=page_count - ) - if job: - job.page_count = page_count - job.credits_charged = billing_amount.amount_micro_dollars - job.billing_status = "billing_failed" - db.commit() - - raise InsufficientCreditsException( - user_message=( - "Insufficient credits to process this document " - f"({page_count} pages required, cost: " - f"{billing_amount.credits})." - ), - required_credits=billing_amount.credits, - internal_message=( - f"job_id={job_id}, user_id={job_user_id}, " - f"page_count={page_count}" - ), - ) - - billing_status = billing_result.billing_status - billing_amount_micro_dollars = billing_result.amount_micro_dollars - billing_credits = billing_result.credits - if job: - job.page_count = page_count - job.credits_charged = billing_amount_micro_dollars - job.billing_status = billing_status - - # Store billing info in Redis - metadata_updates = { - "page_count": page_count, - "billing_status": billing_status, - "billing_amount_micro_dollars": billing_amount_micro_dollars, - "billing_credits": billing_credits, - "processing_started_at": processing_started_at.isoformat(), - } - metadata_service.update_metadata(job_id, metadata_updates) - job_metadata.update(metadata_updates) - - doc_type = JobMetadataHelper.get_parsing_param( - job_metadata, "doc_type", "auto" - ) - logger.info( - f"Start parse: job_id={job_id}, filename={filename}, " - f"internal_filename={internal_parse_name}, type={doc_type}" - ) - - with stage_timer( - "worker.parse.document", - job_id=job_id, - filename=filename, - doc_type=doc_type, - ): - add_dir, add_contents_df = checkerboard_inject_parse( - file_full_path=local_temp_path, - filename=filename, - output_dir=output_dir, - job_id=job_id, - internal_output_filename=internal_parse_name, - kb_dir=JobMetadataHelper.get_parsing_param( - job_metadata, "kb_dir", "Default_Root" - ), - doc_type=doc_type, - smart_title_parse=JobMetadataHelper.get_parsing_param( - job_metadata, "smart_title_parse", True - ), - summary_image=JobMetadataHelper.get_parsing_param( - job_metadata, "summary_image", True - ), - summary_table=JobMetadataHelper.get_parsing_param( - job_metadata, "summary_table", True - ), - summary_txt=JobMetadataHelper.get_parsing_param( - job_metadata, "summary_txt", True - ), - add_frag_desc=JobMetadataHelper.get_parsing_param( - job_metadata, "add_frag_desc", "" - ), - s3_key=s3_key, - ) - parsed_contents_df: pd.DataFrame | None = add_contents_df - - logger.info( - f"File parsing completed: job_id={job_id}, add_dir={add_dir}, chunks={len(parsed_contents_df) if parsed_contents_df is not None else 0}" - ) - - if parsed_contents_df is None: - raise WorkerHandlingException( - user_message="We could not extract content from your file", - internal_message="File parsing failed, no content returned from parser", - ) - - if parsed_contents_df.empty: - logger.warning( - f"No content returned from file parsing: job_id={job_id}, filename={filename}" - ) - - lifecycle_service.update_progress( - job_id, progress=30, message="Parse completed, preparing chunks..." - ) - - chunks = dataframe_to_chunks(parsed_contents_df) - - lifecycle_service.update_progress( - job_id, progress=70, message="Chunks ready, generating zip..." - ) - logger.info(f"Chunks prepared: job_id={job_id}, count={len(chunks)}") - - # Get source file name - source_file_name = JobMetadataHelper.get_field( - job_metadata, "source_file_name" - ) or JobMetadataHelper.get_field(job_metadata, "source_url") - if isinstance(source_file_name, str) and "/" in source_file_name: - source_file_name = os.path.basename(source_file_name) - - document_top_summary = "" - section_summaries: dict[str, str] = {} - if add_dir and source_file_name: - if add_contents_df is not None and "path" in add_contents_df.columns: - ensure_doc_nav_json( - str(add_dir), - chunks, - source_file_name=str(source_file_name), - ) - # Enrich non-leaf section summaries (bottom-up aggregation) - try: - kb_dir_for_enrich = os.path.dirname(str(add_dir)) - summary_use_llm = JobMetadataHelper.get_parsing_param( - job_metadata, "summary_use_llm", False - ) - enrich_doc_nav_summaries( - kb_dir_for_enrich, - source_file=str(source_file_name), - use_llm=summary_use_llm, - ) - section_summaries = build_section_summary_lookup(str(add_dir)) - except Exception as _e: - logger.warning(f"doc_nav enrichment failed (non-fatal): {_e}") - document_top_summary = load_nav_top_summary(str(add_dir), str(source_file_name)) - if document_top_summary: - for chunk in chunks: - metadata = chunk.get("metadata") - if not isinstance(metadata, dict): - metadata = {} - chunk["metadata"] = metadata - metadata["document_top_summary"] = document_top_summary - - data_id = JobMetadataHelper.get_field(job_metadata, "data_id") - - lifecycle_service.update_progress( - job_id, progress=80, message="Generating ZIP package..." - ) - processing_completed_at = datetime.now(timezone.utc) - processing_timing_updates = { - "processing_completed_at": processing_completed_at.isoformat(), - "processing_duration_ms": max( - 0, - int( - ( - processing_completed_at - processing_started_at - ).total_seconds() - * 1000 - ), - ), - } - metadata_service.update_metadata(job_id, processing_timing_updates) - job_metadata.update(processing_timing_updates) - - # Generate ZIP package - zip_service = ZipResultService() - zip_file_path, checksum, statistics, zip_size = ( - zip_service.generate_zip_package( - job_id=job_id, - chunks=chunks, - add_dir=str(add_dir) if add_dir else "", - source_file_name=source_file_name, - data_id=data_id, - job_metadata=job_metadata, - parsed_df=parsed_contents_df, - temp_dir=task_workspace_dir, - ) - ) - - checksum_value = ( - checksum.get("value", "") - if isinstance(checksum, dict) - else (checksum or "") - ) - - lifecycle_service.update_progress( - job_id, progress=90, message="Uploading results to S3..." - ) - - result_bundle = get_result_storage().upload( - job_id=job_id, - result_dir=str(add_dir) if add_dir else "", - zip_file_path=zip_file_path, - ) - result_s3_key = result_bundle.zip_key - - stored_count = 0 - - lifecycle_service.update_progress( - job_id, progress=100, message="Task complete!" - ) - - # Finalize job success directly to the database - lifecycle_service.finalize_job_success( - job_id=job_id, - chunks=chunks, - result_s3_key=result_s3_key, - checksum=checksum_value, - zip_size=zip_size, - stored_count=stored_count, - delivery_mode="url", - section_summaries=section_summaries, - ) - - logger.info( - f"Worker processing complete: job_id={job_id}, result_s3_key={result_s3_key}" - ) - - return { - "status": "success", - "job_id": job_id, - "add_dir": None, - "vectors_count": 0, - "contents_count": len(parsed_contents_df), - "stored_count": stored_count, - "delivery_mode": "url", - "result_s3_key": result_s3_key, - } - finally: - cleanup_task_workspace(task_workspace_dir) diff --git a/apps/worker/app/core/tasks/stale_job_sweeper.py b/apps/worker/app/core/tasks/stale_job_sweeper.py index 30142a614..235c5aabd 100644 --- a/apps/worker/app/core/tasks/stale_job_sweeper.py +++ b/apps/worker/app/core/tasks/stale_job_sweeper.py @@ -103,14 +103,14 @@ def expire_stale_jobs() -> dict: return {"status": "success", "expired": 0, "skipped": 0} for job in expired_jobs: - success = state_machine.mark_failed( + outcome = state_machine.mark_failed_outcome( db, job.job_id, error_message=JOB_EXPIRED_ERROR_MESSAGE, error_code=JOB_EXPIRED_ERROR_CODE, metadata={"sweeper": True, "stale_status": job.status}, ) - if success: + if outcome.succeeded: expired_count += 1 logger.info( f"Expired stale job {job.job_id} (was {job.status})" @@ -118,7 +118,7 @@ def expire_stale_jobs() -> dict: else: skipped_count += 1 logger.debug( - f"Job {job.job_id} already transitioned (CAS miss)" + f"Job {job.job_id} was not expired: reason={outcome.reason}" ) if expired_count > 0: diff --git a/apps/worker/app/core/worker_bootstrap.py b/apps/worker/app/core/worker_bootstrap.py index 191f2fa27..5a0827a22 100644 --- a/apps/worker/app/core/worker_bootstrap.py +++ b/apps/worker/app/core/worker_bootstrap.py @@ -15,7 +15,7 @@ def _register_task_modules() -> None: """Import task modules for Celery side-effect registration.""" - import app.core.tasks.kb_tasks # noqa: F401 + import app.core.tasks.document_ingestion_tasks # noqa: F401 import app.core.tasks.stale_job_sweeper # noqa: F401 import app.core.tasks.webhook_tasks # noqa: F401 @@ -74,7 +74,7 @@ def shutdown_worker(**kwargs) -> None: logger.warning(f"Worker heartbeat cleanup failed: {exc}") try: - from shared.utils.http_clients import close_sync_client + from shared.services.http.client_pool import close_sync_client close_sync_client() logger.info("Worker sync HTTP client closed") @@ -106,6 +106,18 @@ def run_worker() -> None: node_name = f"celery@{hostname}-{pid}" log_level = os.getenv("LOG_LEVEL", "INFO").lower() concurrency = settings.WORKER_CONCURRENCY + worker_queues = ",".join( + [ + "document_ingestion_high", + "document_ingestion_medium", + "document_ingestion_low", + "kb_high", + "kb_medium", + "kb_low", + "ai_high_priority", + "default", + ] + ) celery_args = [ "worker", @@ -114,7 +126,7 @@ def run_worker() -> None: f"--loglevel={log_level}", f"--hostname={node_name}", "-Q", - "kb_high,kb_medium,kb_low,ai_high_priority,default", + worker_queues, "--without-gossip", "--without-mingle", ] diff --git a/apps/worker/app/services/__init__.py b/apps/worker/app/services/__init__.py index 62d28a4ab..8f5b2f6ca 100644 --- a/apps/worker/app/services/__init__.py +++ b/apps/worker/app/services/__init__.py @@ -1,5 +1,3 @@ """Worker service package.""" -# KBOrchestrator was removed and now lives only in the API service. - __all__ = [] diff --git a/apps/worker/app/services/billing/__init__.py b/apps/worker/app/services/billing/__init__.py deleted file mode 100644 index 50a959c99..000000000 --- a/apps/worker/app/services/billing/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -""" -Billing services for the worker. -""" - -from shared.services.billing.credits_sync_service import SyncCreditsService - -__all__ = ["SyncCreditsService"] diff --git a/apps/worker/app/services/common/__init__.py b/apps/worker/app/services/common/__init__.py index ca66e606f..05787f01d 100644 --- a/apps/worker/app/services/common/__init__.py +++ b/apps/worker/app/services/common/__init__.py @@ -1,48 +1,17 @@ -"""Common worker services, including reusable knowledge-base helpers.""" +"""Common worker services.""" -from shared.utils.device_utils import check_internet -from shared.utils.file_utils import clean_file, path_handle -from shared.utils.gc_utils import gc_collect as _gc -from shared.utils.math_utils import min_max_normalize +from app.services.common.device_checks import check_internet +from app.services.common.file_utils import clean_file, path_handle +from app.services.common.resource_cleanup import gc_collect as _gc +from app.services.common.math_helpers import min_max_normalize from shared.utils.text_utils import ( count_cn_en, remove_duplicates_orderkept, tokenize2stw_remove, ) -from .kb_utils import ( - find_images, - find_matches_parsing, - flatten_dic2paths, - flatten_list, - gen_str_codes, - get_str_time, - html2txt, - merge_df, - process_dup_paths_df, - process_path_texts, - remove_spaces, - restore_graph_by_paths, - traverse_dict, -) - __all__ = [ - # From kb_utils "count_cn_en", - "find_images", - "find_matches_parsing", - "flatten_dic2paths", - "flatten_list", - "gen_str_codes", - "get_str_time", - "html2txt", - "merge_df", - "process_dup_paths_df", - "process_path_texts", - "remove_spaces", - "restore_graph_by_paths", - "traverse_dict", - # From shared-python "check_internet", "clean_file", "min_max_normalize", diff --git a/packages/shared-python/shared/utils/device_utils.py b/apps/worker/app/services/common/device_checks.py similarity index 100% rename from packages/shared-python/shared/utils/device_utils.py rename to apps/worker/app/services/common/device_checks.py diff --git a/packages/shared-python/shared/utils/CommonHelperSync.py b/apps/worker/app/services/common/file_loading.py similarity index 73% rename from packages/shared-python/shared/utils/CommonHelperSync.py rename to apps/worker/app/services/common/file_loading.py index 924501ddb..522516cea 100644 --- a/packages/shared-python/shared/utils/CommonHelperSync.py +++ b/apps/worker/app/services/common/file_loading.py @@ -1,22 +1,24 @@ -"""Sync helpers for gevent worker code paths. - -Keep API async helpers in `CommonHelper.py`; worker should import this module. -""" +"""Sync file-loading helpers for worker parsing paths.""" from pathlib import Path -from typing import Optional +from urllib.parse import ParseResult import httpx -def is_remote(path): +def is_remote(path: object) -> bool: """Return True if `path` is an HTTP(S) URL.""" if path is None or not isinstance(path, str): return False return path.startswith("http://") or path.startswith("https://") -def load_file_bytes(file_path, *, file_url: str = "", timeout: Optional[float] = None): +def load_file_bytes( + file_path: str | Path, + *, + file_url: str | ParseResult = "", + timeout: float | None = None, +) -> bytes: """Load bytes from local path or remote URL synchronously.""" if isinstance(file_path, str) and is_remote(file_path): url_to_use = file_path diff --git a/packages/shared-python/shared/utils/file_utils.py b/apps/worker/app/services/common/file_utils.py similarity index 100% rename from packages/shared-python/shared/utils/file_utils.py rename to apps/worker/app/services/common/file_utils.py diff --git a/apps/worker/app/services/common/kb_utils.py b/apps/worker/app/services/common/kb_utils.py deleted file mode 100755 index 43389a56f..000000000 --- a/apps/worker/app/services/common/kb_utils.py +++ /dev/null @@ -1,406 +0,0 @@ -import os -import re -import uuid -from datetime import datetime - -import pandas as pd -from bs4 import BeautifulSoup - -from shared.core.config import settings -from shared.utils.chunk_refs import extract_chunk_refs -from shared.utils.file_utils import path_handle -from shared.utils.text_utils import _CN_EN_NUM_RE - -SUMMARY_PATH_MARKERS: tuple[str, ...] = ("summary", "\u6458\u8981\u603b\u7ed3") - - -def gen_str_codes(input_string): - """Generate a UUID5 code from a string.""" - namespace = uuid.NAMESPACE_DNS - return str(uuid.uuid5(namespace, input_string)) - - -def get_str_time(): - """Get the current time as a string.""" - now = datetime.now() - return now.strftime("%Y-%m-%d %H:%M:%S") - - -def find_images(folder_path): - """Find image files inside a folder tree.""" - image_extensions = {".png", ".jpg", ".jpeg"} - image_files = [] - - for _, _, files in os.walk(folder_path): - files.sort() - for file in files: - if os.path.splitext(file)[1].lower() in image_extensions: - image_files.append(file) - return image_files - - -def find_matches_parsing(content, path): - """Parse table and image markers from content.""" - matches = extract_chunk_refs(content) - if len(matches) == 0: - match_type = "PTXT" - else: - match_type = "\n".join((["PTXT"] + matches)) - - split_char = settings.SPLIT_CHAR or ";" - if any( - f"{split_char}{summary_marker}" in path - for summary_marker in SUMMARY_PATH_MARKERS - ): - parent_path = path.split(split_char)[-2] - match_type = "SUMMARY_" + parent_path + "_SUMMARY" - return match_type - - -def flatten_list(nested_list): - """Flatten a nested list.""" - flat_list = [] - for item in nested_list: - if isinstance(item, list): - flat_list.extend(flatten_list(item)) - else: - flat_list.append(item) - return flat_list - - -def flatten_dic2paths(d, current_path=None, result=None): - """Flatten a nested dict into path strings.""" - if result is None: - result = [] - if current_path is None: - current_path = [] - - for key, value in d.items(): - if not isinstance(key, str): - continue - new_path = current_path + [key] - if isinstance(value, dict) and value: - flatten_dic2paths(value, new_path, result) - else: - split_char = settings.SPLIT_CHAR or ";" - result.append(split_char.join(new_path)) - return result - - -def merge_df(input_df): - """Merge DataFrame rows that share the same path.""" - dfs_by_path = list(input_df.groupby("path", sort=False)) - processed_dfs = [] - - for key, df in dfs_by_path: - content_to_merge = [] - types_to_merge = [] - total_length = 0 - - for i, row in df.iterrows(): - content_to_merge.append(row["content"]) - types_to_merge.extend(row["type"].split("\n")) - total_length += len(row["content"]) - - content_to_merge = "\n".join(content_to_merge) - temp_merge_df = pd.DataFrame( - [ - { - "content": content_to_merge, - "type": "\n".join(list(set(types_to_merge))), - "path": key, - "length": total_length, - "know_id": gen_str_codes(content_to_merge), - } - ] - ) - processed_dfs.append(temp_merge_df) - - final_df = pd.concat(processed_dfs, axis=0, ignore_index=True) - return final_df - - -def process_path_texts(path_, last=50): - """Normalize path text for downstream use.""" - temp_path = path_handle(path_, mode="sanitize") - if temp_path is None: - return "" - return "_".join(temp_path.split(os.sep))[:last] - - -def process_dup_paths_df(df): - """ - de-duplicate kbs dataframe for final output - - Args: - df: initial dataframe after all heading stacking - - Returns: - Dataframe without any duplicate paths - """ - if "path" not in df.columns: - return df - - split_char = settings.SPLIT_CHAR or "/" - - # Step 1: detect if there are any duplicated paths - dup_mask = df["path"].duplicated(keep=False) - if not dup_mask.any(): - return df - - # Step 2: record ids of duplicated paths as a mapping - path_occurrences = {} # path -> list of row indices - for idx, path in enumerate(df["path"]): - if path not in path_occurrences: - path_occurrences[path] = [] - path_occurrences[path].append(idx) - - # path_renames: row_index -> new_path (recording rows renamed) - # parent_rename_map: original_path -> {row_index: new_path} - path_renames = {} - parent_rename_map = {} - - for path, indices in path_occurrences.items(): - if len(indices) > 1: # only process duplicated paths - parent_rename_map[path] = {} - for occurrence, idx in enumerate(indices): - if occurrence == 0: - # keep the first appearance as it is - path_renames[idx] = path - else: - # add suffix to subsequent appearances - new_path = f"{path}_{occurrence + 1}" - path_renames[idx] = new_path - parent_rename_map[path][idx] = new_path - - # Step 3: process all rows, update paths - new_paths = [] - - for idx, row in df.iterrows(): - path = row["path"] - - # Check whether this row itself needs renaming. - new_path = path_renames.get(idx, path) - path_parts = new_path.split(split_char) - - # Check whether this row is under a renamed parent path. - for parent_path, rename_info in parent_rename_map.items(): - parent_parts = parent_path.split(split_char) - - # Check whether the current path starts with that parent path. - if ( - len(path_parts) > len(parent_parts) - and path_parts[: len(parent_parts)] == parent_parts - ): - # Find the nearest renamed parent path that appears earlier. - matching_parent_idx = None - for parent_idx in sorted(rename_info.keys(), reverse=True): - if parent_idx < idx: - matching_parent_idx = parent_idx - break - - if matching_parent_idx is not None: - renamed_parent = rename_info[matching_parent_idx] - renamed_parent_parts = renamed_parent.split(split_char) - new_path_parts = ( - renamed_parent_parts + path_parts[len(parent_parts) :] - ) - new_path = split_char.join(new_path_parts) - break - new_paths.append(new_path) - - df = df.copy() - df["path"] = new_paths - return df - - -def remove_spaces(text, handle_punctuation=False): - """Remove spaces between Chinese chars while keeping English word spacing.""" - if handle_punctuation: - punctuation = ( - r"""!"#$%&'()*+,-./:;<=>?@[\]^_`{|}~，。、【】《》？；：''""（）…—-！""" - ) - res_text = re.sub(f"[{re.escape(punctuation)}]", "", text) - else: - pattern = re.compile(r"([\u4e00-\u9fff])\s+|(?<=\s)([\u4e00-\u9fff])") - - def replacer(match): - return match.group(1) or match.group(2) - - res_text = pattern.sub(replacer, text) - - res_text = re.sub(r"\s+", " ", res_text) - return res_text.strip() - - -def traverse_dict(d, parent=None): - """Traverse a dictionary and generate description text.""" - dic_texts = [] - for key, value in d.items(): - if value: - child_keys = ", ".join(value.keys()) - text = f"'{key}' includes {child_keys}" - dic_texts.append(text) - dic_texts.extend(traverse_dict(value, key)) - return dic_texts - - -def restore_graph_by_paths(paths): - """Rebuild a graph structure from path strings.""" - root_dict = {} - split_char = settings.SPLIT_CHAR or ";" - for path in paths: - nodes = path.split(split_char) - current_dict = root_dict - for node in nodes: - if node not in current_dict: - current_dict[node] = {} - current_dict = current_dict[node] - dic_texts = traverse_dict(root_dict) - return root_dict, dic_texts - - -def html2txt(html_text): - """Convert HTML into plain text.""" - soup = BeautifulSoup(html_text, "html.parser") - text = soup.get_text() - return text - - -def normalize_md(s: str) -> str: - """Normalize markdown string for comparison - - Removes heading markers (###) and whitespace, converts to lowercase. - Used for TOC keyword matching. - """ - s = re.sub(r"^\s*#+\s*", "", s) - s = re.sub(r"\s+", "", s) - return s.lower() - - -# --------------------------------------------------------------------------- -# truncate_text (character-based) — KEPT for table-cell display callers in -# doc_parser.py and html_parser.py where a per-character limit is intentional. -# Do NOT use for heading / semantic text truncation; use truncate_text_by_tokens. -# --------------------------------------------------------------------------- -def truncate_text(text: str, start_limit: int, end_limit: int) -> str: - """Truncate text by raw character count, keeping start and end parts. - - Intended for short display values (table headers, file names, etc.) where - a fixed character budget is appropriate. For heading / semantic text where - English words must not be split mid-word, use ``truncate_text_by_tokens``. - - Args: - text: Text to truncate. - start_limit: Number of characters to keep from start. - end_limit: Number of characters to keep from end (0 = no tail). - - Returns: - Truncated text with '...' in the middle if it exceeds the limits. - """ - text = str(text) - total_limit = start_limit + end_limit - if len(text) <= total_limit: - return text - start_part = text[:start_limit] - end_part = text[-end_limit:] if end_limit > 0 else "" - return f"{start_part}...{end_part}" - - -# --------------------------------------------------------------------------- -# Language detection & language-aware token truncation -# --------------------------------------------------------------------------- - -_CN_CHAR_RE = re.compile(r"[\u4e00-\u9fff]") - -EN_START_LIMIT = 15 # token budget for English-dominant headings -CN_RATIO_THRESHOLD = 0.3 # if ≥30 % of tokens are Chinese chars → "Chinese" - - -def detect_primary_lang(text: str) -> str: - """Detect whether *text* is primarily Chinese or English/other. - - Uses the semantic tokens already defined by ``_CN_EN_NUM_RE`` - (Chinese chars, English word runs, number groups). If Chinese - characters account for at least ``CN_RATIO_THRESHOLD`` of all - tokens the text is classified as ``'zh'``; otherwise ``'en'``. - - Args: - text: Input text (heading or any short string). - - Returns: - ``'zh'`` for Chinese-dominant text, ``'en'`` otherwise. - """ - if not text: - return "en" - tokens = _CN_EN_NUM_RE.findall(text) - if not tokens: - return "en" - cn_count = sum(1 for t in tokens if _CN_CHAR_RE.fullmatch(t)) - return "zh" if (cn_count / len(tokens)) >= CN_RATIO_THRESHOLD else "en" - - -def count_cn_en(text: str) -> int: - """Count semantic Chinese/English/number tokens in a string.""" - return len(_CN_EN_NUM_RE.findall(str(text))) - - -def truncate_text_by_tokens( - text: str, - start_limit: int, - end_limit: int, - lang_aware: bool = True, -) -> str: - """Truncate text by semantic token count, preserving whole words. - - Uses the same token definition as ``count_cn_en``: - - - each Chinese character = 1 token - - each run of English letters = 1 token - - each number group = 1 token - - punctuation and whitespace are excluded from the count but - preserved in the output up to the split point. - - When *lang_aware* is ``True`` (default), the function auto-detects - whether the text is English-dominant and caps ``start_limit`` at - ``EN_START_LIMIT`` (15) in that case. Chinese-dominant text keeps - the caller-supplied ``start_limit`` (typically 30). This prevents - over-long English heading chunks while still allowing a generous - budget for dense Chinese text. - - Cut points are placed *after* the last character of the start - token and *before* the first character of the first tail token, - so no word is ever split in the middle. - - Args: - text: Text to truncate. - start_limit: Max tokens to keep from the start. When - *lang_aware* is True and the text is English-dominant, - this is silently capped at ``EN_START_LIMIT``. - end_limit: Max tokens to keep from the end (0 = no tail). - lang_aware: When True, auto-detect language and apply a tighter - budget for English text. Set to False to use raw limits. - - Returns: - Truncated text with ``'...'`` in the middle when the token - count exceeds ``start_limit + end_limit``. Returns the - original text unchanged when the count is within the budget. - """ - text = str(text) - matches = list(_CN_EN_NUM_RE.finditer(text)) - total = len(matches) - - if lang_aware and total > 0: - lang = detect_primary_lang(text) - if lang == "en": - start_limit = min(start_limit, EN_START_LIMIT) - - if total <= start_limit + end_limit: - return text - # Cut position: end of the start_limit-th token - cut_start = matches[start_limit - 1].end() if start_limit > 0 else 0 - # Tail position: start of the (total - end_limit)-th token - cut_end = matches[total - end_limit].start() if end_limit > 0 else len(text) - if cut_start >= cut_end: - return text - return text[:cut_start] + "..." + text[cut_end:] diff --git a/packages/shared-python/shared/utils/math_utils.py b/apps/worker/app/services/common/math_helpers.py similarity index 100% rename from packages/shared-python/shared/utils/math_utils.py rename to apps/worker/app/services/common/math_helpers.py diff --git a/packages/shared-python/shared/utils/gc_utils.py b/apps/worker/app/services/common/resource_cleanup.py similarity index 100% rename from packages/shared-python/shared/utils/gc_utils.py rename to apps/worker/app/services/common/resource_cleanup.py diff --git a/apps/worker/app/services/connect_builder/__init__.py b/apps/worker/app/services/connect_builder/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/apps/worker/app/services/connect_builder/builder.py b/apps/worker/app/services/connect_builder/builder.py deleted file mode 100644 index b788c7ff1..000000000 --- a/apps/worker/app/services/connect_builder/builder.py +++ /dev/null @@ -1,527 +0,0 @@ -""" -ConnectTo Builder — KB-level post-processor for inter-chunk relationships. - -This module discovers relationships between chunks across different files -within a knowledge base, populating the `connectto` column in the DataFrame. -""" - -import json -import re -from collections import defaultdict -from difflib import SequenceMatcher -from typing import Any, Dict, List, Optional, Tuple - -from loguru import logger - -from shared.utils.chunk_refs import CHUNK_REF_RE - -# ─── Relation Type Registry (extensible, not hard-coded) ────────────────────── - -RELATION_REGISTRY: Dict[str, Dict[str, Any]] = { - "related": { - "description": "Chunks share common concepts or topics", - "requires_llm": False, - }, - # TODO: LLM-classified relation types — uncomment when classify_relation() is implemented - # "contradicts": { - # "description": "Chunks describe opposing or contradictory facts", - # "requires_llm": True, - # }, - # "causal": { - # "description": "Chunks have a cause-effect relationship", - # "requires_llm": True, - # }, - # "extends": { - # "description": "One chunk extends or improves upon the other", - # "requires_llm": True, - # }, - # "supports": { - # "description": "One chunk provides evidence supporting the other", - # "requires_llm": True, - # }, - # "same_method": { - # "description": "Chunks discuss the same methodology or technique", - # "requires_llm": True, - # }, - # "same_data": { - # "description": "Chunks reference the same dataset", - # "requires_llm": True, - # }, -} - - -# ─── Default Configuration ──────────────────────────────────────────────────── - -DEFAULT_CONFIG: Dict[str, Any] = { - # Minimum number of shared keywords to consider a connection - "min_keyword_overlap": 3, - # Weight multiplier for keyword score (linear) - "keyword_score_weight": 1.0, - # Maximum connections per chunk (top-N by score) - "max_connections_per_chunk": 10, - # Minimum score threshold to create a connection - "min_score_threshold": 0.8, - # Only connect chunks from different files (skip intra-file) - "cross_file_only": True, - # Maximum character overlap ratio to allow (filter near-duplicates) - # Pairs with SequenceMatcher.ratio() >= this are considered duplicates - "max_content_overlap": 0.8, -} - - -# ─── Keyword Normalization ──────────────────────────────────────────────────── - -# TODO: Synonym dictionary for advanced normalization (e.g. "RL" ↔ "reinforcement learning") -_SYNONYM_MAP: Dict[str, str] = {} - - -def _normalize_keyword(keyword: str) -> str: - """ - Normalize a keyword for matching: - - lowercase - - strip whitespace - - collapse multiple spaces - - apply synonym mapping (TODO) - - Args: - keyword: Raw keyword string. - - Returns: - Normalized keyword string. - """ - kw = keyword.lower().strip() - kw = re.sub(r"\s+", " ", kw) - - # Apply synonym mapping if available - return _SYNONYM_MAP.get(kw, kw) - - -def _extract_file_key(path: str) -> str: - """ - Extract a file-level key from a chunk's path to determine - whether two chunks belong to the same file. - - Example paths: - "Default_Root/paper.pdf/Section 1/Subsection" → "Default_Root/paper.pdf" - "KB_DATA/reports/annual.docx/Table 1" → "KB_DATA/reports/annual.docx" - - Heuristic: take the path up to and including the first segment - that looks like a filename (has an extension). - """ - if not path: - return "" - - parts = path.replace("\\", "/").split("/") - file_parts = [] - for part in parts: - file_parts.append(part) - # Check if this segment looks like a file (has extension) - if "." in part and not part.startswith("."): - break - - return "/".join(file_parts) - - -# ─── Keyword Inverted Index ────────────────────────────────────────────────── - - -def _build_keyword_index( - chunks: List[Dict[str, Any]], -) -> Dict[str, List[Tuple[str, str]]]: - """ - Build an inverted index: normalized_keyword → [(chunk_id, file_key)]. - - Args: - chunks: List of chunk dicts, each having: - - "chunk_id": str - - "metadata" or "keywords": keyword source - - "path": str - - Returns: - Dict mapping normalized keyword → list of (chunk_id, file_key) tuples. - """ - index: Dict[str, List[Tuple[str, str]]] = defaultdict(list) - - for chunk in chunks: - chunk_id = chunk.get("chunk_id") or chunk.get("know_id", "") - path = chunk.get("path", "") - file_key = _extract_file_key(path) - - # Extract keywords from metadata or top-level - keywords = _get_keywords(chunk) - if not keywords: - continue - - for kw in keywords: - normalized = _normalize_keyword(kw) - if normalized: - index[normalized].append((str(chunk_id), file_key)) - - return dict(index) - - -def _get_keywords(chunk: Dict[str, Any]) -> List[str]: - """ - Extract keywords from a chunk, supporting multiple input formats: - - chunk["metadata"]["keywords"] (list) - - chunk["keywords"] (list or semicolon-separated string) - - chunk["metadata"]["tokens"] or chunk["tokens"] (fallback: jieba word chain) - """ - # Try metadata.keywords first - metadata = chunk.get("metadata", {}) - if isinstance(metadata, dict): - kws = metadata.get("keywords", []) - if isinstance(kws, list) and kws: - return kws - - # Try top-level keywords - kws = chunk.get("keywords", []) - if isinstance(kws, list) and kws: - return kws - if isinstance(kws, str) and kws.strip(): - # Parse semicolon or comma separated - if ";" in kws: - return [k.strip() for k in kws.split(";") if k.strip()] - elif "," in kws: - return [k.strip() for k in kws.split(",") if k.strip()] - return [kws.strip()] - - # ─── Fallback: tokens (jieba word chain) ────────────────────────────── - tokens = _parse_tokens_field( - metadata.get("tokens") if isinstance(metadata, dict) else None - ) - if not tokens: - tokens = _parse_tokens_field(chunk.get("tokens")) - return tokens - - -# Pre-compiled patterns for token noise filtering -_UUID_LIKE_RE = re.compile(r"^[0-9a-f]{4,}$", re.IGNORECASE) -_MARKER_PREFIXES = ("IMAGE_", "TABLE_", "PTXT", "image-", "table-") - - -def _parse_tokens_field(raw) -> List[str]: - """ - Parse the tokens field into a filtered keyword list. - - Accepts: - - List[str]: already parsed (from chunks.json after safe_parse_tokens) - - str with ';': semicolon-separated tokens (new format, matches keywords) - - str with '->': arrow-separated jieba word chain (legacy format) - - str with "['...']": legacy list-repr format - - Filters out noise: single-char tokens, UUIDs, IMAGE_/TABLE_ markers. - """ - if raw is None: - return [] - - # Already a list (from chunks.json) - if isinstance(raw, list): - words = raw - elif isinstance(raw, str): - raw = raw.strip() - if not raw: - return [] - # List-repr format: "['w1;w2;w3']" or "['w1->w2->w3']" - if raw.startswith("[") and raw.endswith("]"): - inner = raw[1:-1].strip() - if (inner.startswith("'") and inner.endswith("'")) or ( - inner.startswith('"') and inner.endswith('"') - ): - inner = inner[1:-1] - raw = inner - # Determine separator: semicolon (new) or arrow (legacy) - if ";" in raw: - words = [t.strip() for t in raw.split(";") if t.strip()] - elif "->" in raw: - words = [t.strip() for t in raw.split("->") if t.strip()] - else: - return [] - else: - return [] - - # Filter noise - filtered = [] - for w in words: - if len(w) <= 1: - continue - if any(w.startswith(p) for p in _MARKER_PREFIXES) or CHUNK_REF_RE.fullmatch(w): - continue - if _UUID_LIKE_RE.match(w): - continue - filtered.append(w) - return filtered - - -# ─── Scoring ───────────────────────────────────────────────────────────────── - - -def _compute_keyword_score( - shared_kws: set, - kws_a: set, - kws_b: set, - weight: float = 1.0, -) -> float: - """ - Compute keyword overlap score using character-length-weighted scoring. - - Longer tokens contribute more: a four-character term contributes twice the - weight of a two-character term. - Formula: score = weight * sum(len(kw) for shared) / min(sum(len) for A, sum(len) for B) - - Args: - shared_kws: Set of shared (normalized) keywords. - kws_a: Full keyword set of chunk A. - kws_b: Full keyword set of chunk B. - weight: Score multiplier. - - Returns: - Float score in [0, weight]. - """ - weighted_a = sum(len(k) for k in kws_a) - weighted_b = sum(len(k) for k in kws_b) - denominator = min(weighted_a, weighted_b) - if denominator == 0: - return 0.0 - weighted_shared = sum(len(k) for k in shared_kws) - return weight * weighted_shared / denominator - - -# ─── Main Entry Point ──────────────────────────────────────────────────────── - - -def build_connections( - chunks: List[Dict[str, Any]], - config: Optional[Dict[str, Any]] = None, -) -> Dict[str, List[Dict[str, Any]]]: - """ - Compute "related" connections between chunks based on keyword overlap. - - Args: - chunks: List of chunk dicts. Each must have: - - chunk_id (or know_id): str - - path: str - - metadata.keywords or keywords: List[str] - config: Optional overrides for DEFAULT_CONFIG. - - Returns: - Dict mapping chunk_id → list of connection dicts: - [{"target": "...", "relation": "related", "score": 0.82, "keywords": ["PPO", "RL"]}] - """ - cfg = {**DEFAULT_CONFIG, **(config or {})} - - min_overlap = cfg["min_keyword_overlap"] - kw_weight = cfg["keyword_score_weight"] - max_conns = cfg["max_connections_per_chunk"] - min_score = cfg["min_score_threshold"] - cross_only = cfg["cross_file_only"] - max_overlap = cfg.get("max_content_overlap", 0.8) - - # Build inverted index - kw_index = _build_keyword_index(chunks) - - # Pre-compute per-chunk data - chunk_data: Dict[ - str, Tuple[str, set] - ] = {} # chunk_id → (file_key, normalized_keywords) - chunk_content: Dict[str, str] = {} # chunk_id → content (for dedup) - for chunk in chunks: - cid = str(chunk.get("chunk_id") or chunk.get("know_id", "")) - if not cid: - continue - file_key = _extract_file_key(chunk.get("path", "")) - kws = _get_keywords(chunk) - normalized_kws = {_normalize_keyword(k) for k in kws if k} - normalized_kws.discard("") - chunk_data[cid] = (file_key, normalized_kws) - chunk_content[cid] = chunk.get("content") or chunk.get("text", "") - - # For each chunk, find candidates via keyword index - connections: Dict[str, List[Dict[str, Any]]] = defaultdict(list) - - for cid, (file_key, my_kws) in chunk_data.items(): - if not my_kws: - continue - - # Collect candidates and their shared keywords - candidate_shared: Dict[str, set] = defaultdict(set) # target_id → shared_kw set - - for kw in my_kws: - entries = kw_index.get(kw, []) - for target_id, target_file in entries: - if target_id == cid: - continue - if cross_only and target_file == file_key: - continue - candidate_shared[target_id].add(kw) - - # Score and filter candidates - scored: List[Tuple[str, float, set]] = [] - for target_id, shared_kws in candidate_shared.items(): - if len(shared_kws) < min_overlap: - continue - - target_data = chunk_data.get(target_id) - if not target_data: - continue - - _, target_kws = target_data - score = _compute_keyword_score( - shared_kws=shared_kws, - kws_a=my_kws, - kws_b=target_kws, - weight=kw_weight, - ) - if score >= min_score: - # Near-duplicate filter: skip pairs with high character overlap - if max_overlap < 1.0: - src_text = chunk_content.get(cid, "") - tgt_text = chunk_content.get(target_id, "") - if src_text and tgt_text: - char_ratio = SequenceMatcher(None, src_text, tgt_text).ratio() - if char_ratio >= max_overlap: - continue - scored.append((target_id, score, shared_kws)) - - # Sort by score descending, keep top-N - scored.sort(key=lambda x: x[1], reverse=True) - for target_id, score, shared_kws in scored[:max_conns]: - connections[cid].append( - { - "target": target_id, - "relation": "related", - "score": round(score, 4), - "keywords": sorted(shared_kws), - } - ) - - total_edges = sum(len(v) for v in connections.values()) - logger.info( - f"🔗 ConnectTo: {total_edges} connections found " - f"across {len(connections)} chunks " - f"(config: min_overlap={min_overlap}, threshold={min_score})" - ) - - return dict(connections) - - -# ─── Serialization ──────────────────────────────────────────────────────────── - - -def serialize_connections(connections: List[Dict[str, Any]]) -> str: - """ - Serialize a chunk's connection list to a JSON string - for storage in the `connectto` DataFrame column. - - Args: - connections: List of connection dicts. - - Returns: - JSON string, or empty string if no connections. - """ - if not connections: - return "" - return json.dumps(connections, ensure_ascii=False, separators=(",", ":")) - - -def deserialize_connections(raw: Any) -> List[Dict[str, Any]]: - """ - Deserialize the `connectto` column value back to a list of connections. - - Handles: - - JSON array string: '[{"target": "...", ...}]' - - Single target string: 'doc/section-a' - - Empty / NaN / None: returns [] - - Args: - raw: Raw value from DataFrame connectto column. - - Returns: - List of connection dicts. - """ - if raw is None: - return [] - - try: - import pandas as pd - - if pd.isna(raw): - return [] - except (ImportError, TypeError, ValueError): - pass - - raw_str = str(raw).strip() - if not raw_str: - return [] - - # Try JSON parse first - if raw_str.startswith("["): - try: - parsed = json.loads(raw_str) - if isinstance(parsed, list): - return parsed - except json.JSONDecodeError: - pass - - if raw_str: - return [ - {"target": raw_str, "relation": "related", "score": 1.0, "keywords": []} - ] - - return [] - - -# ─── LLM Relation Classification (stub) ────────────────────────────────────── - - -def classify_relation( - summary_a: str, - summary_b: str, - shared_keywords: List[str], - llm_client: Any = None, -) -> Dict[str, Any]: - """ - Classify the specific relation type between two related chunks using LLM. - - This function is a **stub** — the concrete classification prompt and LLM - call logic are TODO. Currently returns "related" for all pairs. - - Args: - summary_a: Summary text of chunk A. - summary_b: Summary text of chunk B. - shared_keywords: Keywords they share. - llm_client: Optional LLM client for making classification calls. - - Returns: - Dict with keys: - - "relation": str (from RELATION_REGISTRY) - - "reason": str (human-readable explanation) - - "confidence": float (0.0 ~ 1.0) - - TODO: Implement classification prompt: - Given two knowledge chunks: - [Chunk A]: {summary_a} - [Chunk B]: {summary_b} - Shared concepts: {shared_keywords} - - Classify their relationship: - - contradicts: A and B describe opposing facts - - causal: A and B have a cause-effect relationship - - extends: B extends or improves upon A - - supports: B provides evidence for A - - same_method: A and B use the same methodology - - same_data: A and B use the same dataset - - related: Related but none of the above - - other: Has a clear relationship not listed above (describe it) - - Return JSON: {"relation": "...", "reason": "...", "confidence": 0.0~1.0} - """ - return { - "relation": "related", - "reason": ( - f"Keyword overlap: {', '.join(shared_keywords)}" - if shared_keywords - else "Keyword overlap" - ), - "confidence": 1.0, - } diff --git a/apps/worker/app/services/connect_builder/graph_builder.py b/apps/worker/app/services/connect_builder/graph_builder.py deleted file mode 100644 index 4803034eb..000000000 --- a/apps/worker/app/services/connect_builder/graph_builder.py +++ /dev/null @@ -1,1538 +0,0 @@ -""" -Knowledge Graph Builder — KB-level knowledge graph assembler (v2.0). - -Assembles a file-level knowledge_graph.json from parsed chunks + connect_builder edges. -Deployed to ~/.knowhere/{kb_id}/ and grows incrementally as more files are parsed. - -Architecture: - - files: per-file summaries (chunks_count, types, top_keywords, importance) - - edges: cross-file relationships (aggregated from chunk-level connections) - - chunk_stats.json: per-chunk usage tracking (hit_count, last_hit, decay) - - Per-file chunks.json: full chunk data lives in subdirectories - -Usage: - # One-stop API (recommended) - graph = build_and_deploy(chunks, kb_id="my_kb", parsed_output_dir=add_dir) - - # Manual: first build - graph = build_knowledge_graph(all_chunks, connections, kb_id="my_kb") - - # Manual: incremental update - graph = update_knowledge_graph(existing_graph, new_chunks, existing_chunks) -""" - -import json -import math -import os -import re -from collections import defaultdict -from datetime import datetime, timezone -from typing import Any, Dict, List, Optional, Tuple - -from app.services.connect_builder.builder import ( - DEFAULT_CONFIG, - _compute_keyword_score, - _extract_file_key, - _get_keywords, - _normalize_keyword, -) -from loguru import logger - -from shared.utils.chunk_refs import CHUNK_REF_PATTERN - -# ─── Tree Construction ─────────────────────────────────────────────────────── - - -def _build_tree_from_paths(paths: List[str]) -> Dict[str, Any]: - """ - Rebuild hierarchical tree from chunk path list. - - Args: - paths: List of chunk paths, e.g. ["Default_Root/report.pdf/Section 1/1.1", ...] - - Returns: - Nested dict tree rooted at Default_Root. - """ - root: Dict[str, Any] = {} - for path in paths: - if not path: - continue - nodes = [n.strip() for n in path.split("/") if n.strip()] - current = root - for node in nodes: - if node not in current: - current[node] = {} - current = current[node] - return root - - -def _merge_tree(base: Dict[str, Any], addition: Dict[str, Any]) -> Dict[str, Any]: - """ - Deep-merge two tree dicts. Addition is merged INTO base (in-place). - - Args: - base: Existing tree. - addition: New tree to merge in. - - Returns: - The merged base dict (same reference). - """ - for key, value in addition.items(): - if key in base and isinstance(base[key], dict) and isinstance(value, dict): - _merge_tree(base[key], value) - else: - base[key] = value - return base - - -# ─── Node Extraction ───────────────────────────────────────────────────────── - - -def _chunks_to_nodes( - chunks: List[Dict[str, Any]], - content_preview_len: int = 200, -) -> List[Dict[str, Any]]: - """ - Extract node metadata from chunks for the knowledge graph. - - Args: - chunks: List of normalized chunk dicts. - content_preview_len: Max characters for content_preview. - - Returns: - List of node dicts with: id, type, path, summary, keywords, content_preview. - """ - nodes = [] - for chunk in chunks: - chunk_id = str(chunk.get("chunk_id") or chunk.get("know_id", "")) - if not chunk_id: - continue - - content = chunk.get("content") or chunk.get("text", "") - metadata = chunk.get("metadata", {}) - if not isinstance(metadata, dict): - metadata = {} - - # Extract keywords from metadata or top-level - keywords = metadata.get("keywords", []) - if not keywords: - keywords = chunk.get("keywords", []) - if isinstance(keywords, str): - keywords = [k.strip() for k in keywords.split(";") if k.strip()] - - node = { - "id": chunk_id, - "type": chunk.get("type", "text"), - "path": chunk.get("path", ""), - "summary": metadata.get("summary") or chunk.get("summary", ""), - "keywords": keywords, - "content_preview": content[:content_preview_len] if content else "", - } - nodes.append(node) - - return nodes - - -# ─── Edge Extraction ───────────────────────────────────────────────────────── - - -def _connections_to_edges( - connections: Dict[str, List[Dict[str, Any]]], -) -> List[Dict[str, Any]]: - """ - Convert connect_builder output to deduplicated edge list. - connect_builder produces bidirectional entries (A→B and B→A); - we deduplicate to keep only one edge per pair. - - Args: - connections: Output from build_connections(), mapping chunk_id → list of connections. - - Returns: - List of edge dicts: {source, target, relation, score, shared_keywords}. - """ - seen_pairs: set = set() - edges = [] - - for source_id, conn_list in connections.items(): - for conn in conn_list: - target_id = conn.get("target", "") - pair_key = tuple(sorted([source_id, target_id])) - if pair_key in seen_pairs: - continue - seen_pairs.add(pair_key) - - edges.append( - { - "source": source_id, - "target": target_id, - "relation": conn.get("relation", "related"), - "score": conn.get("score", 0.0), - "shared_keywords": conn.get("keywords", []), - } - ) - - return edges - - -def _merge_related_connections_into_chunks( - chunks: List[Dict[str, Any]], - connections: Dict[str, List[Dict[str, Any]]], -) -> None: - """Backfill related connections into chunk metadata without touching embeds.""" - if not chunks or not connections: - return - - chunk_map = { - str(chunk.get("chunk_id") or chunk.get("know_id", "")): chunk - for chunk in chunks - if chunk.get("chunk_id") or chunk.get("know_id") - } - - for chunk_id, conn_list in connections.items(): - chunk = chunk_map.get(str(chunk_id)) - if not chunk: - continue - metadata = chunk.setdefault("metadata", {}) - if not isinstance(metadata, dict): - metadata = {} - chunk["metadata"] = metadata - existing = metadata.get("connect_to", []) - if not isinstance(existing, list): - existing = [] - - merged = [] - seen = set() - for item in existing: - if not isinstance(item, dict): - continue - key = ( - str(item.get("target") or ""), - str(item.get("relation") or "related"), - str(item.get("ref") or ""), - ) - if key in seen: - continue - seen.add(key) - merged.append(item) - - for conn in conn_list: - if not isinstance(conn, dict): - continue - if conn.get("relation", "related") != "related": - continue - key = ( - str(conn.get("target") or ""), - "related", - "", - ) - if key in seen: - continue - seen.add(key) - merged.append( - { - "target": conn.get("target", ""), - "relation": "related", - "score": conn.get("score", 0.0), - "keywords": conn.get("keywords", []), - } - ) - - metadata["connect_to"] = merged - - -def _save_chunks_by_source_file(kb_dir: str, chunks: List[Dict[str, Any]]) -> None: - """Persist grouped chunks.json files after metadata backfill.""" - grouped: Dict[str, List[Dict[str, Any]]] = defaultdict(list) - for chunk in chunks: - source_file = chunk.get("_source_file") - if not source_file: - continue - cleaned = dict(chunk) - cleaned.pop("_source_file", None) - grouped[str(source_file)].append(cleaned) - - for source_file, source_chunks in grouped.items(): - output_path = os.path.join(kb_dir, source_file, "chunks.json") - _save_chunks(source_chunks, output_path) - - -# ─── Incremental Matching ──────────────────────────────────────────────────── - - -def _incremental_connections( - new_chunks: List[Dict[str, Any]], - existing_chunks: List[Dict[str, Any]], - config: Optional[Dict[str, Any]] = None, -) -> Dict[str, List[Dict[str, Any]]]: - """ - Match ONLY new_chunks ↔ existing_chunks (skip existing ↔ existing). - Reuses connect_builder scoring functions. - - Complexity: O(new × existing) instead of O(all²). - - Args: - new_chunks: Newly parsed file's chunks. - existing_chunks: All previously known chunks. - config: Optional config overrides (same keys as connect_builder.DEFAULT_CONFIG). - - Returns: - Dict mapping chunk_id → list of connection dicts (same format as build_connections). - """ - from difflib import SequenceMatcher - - cfg = {**DEFAULT_CONFIG, **(config or {})} - min_overlap = cfg["min_keyword_overlap"] - kw_weight = cfg["keyword_score_weight"] - max_conns = cfg["max_connections_per_chunk"] - min_score = cfg["min_score_threshold"] - cross_only = cfg["cross_file_only"] - max_content_overlap = cfg.get("max_content_overlap", 0.8) - - # Pre-compute keyword sets for new chunks - new_data: Dict[str, Tuple[str, set, str]] = {} # id → (file_key, kw_set, content) - for chunk in new_chunks: - cid = str(chunk.get("chunk_id") or chunk.get("know_id", "")) - if not cid: - continue - file_key = _extract_file_key(chunk.get("path", "")) - kws = _get_keywords(chunk) - normalized = {_normalize_keyword(k) for k in kws if k} - normalized.discard("") - content = chunk.get("content") or chunk.get("text", "") - new_data[cid] = (file_key, normalized, content) - - # Pre-compute keyword sets for existing chunks - existing_data: Dict[str, Tuple[str, set, str]] = {} - for chunk in existing_chunks: - cid = str(chunk.get("chunk_id") or chunk.get("know_id", "")) - if not cid: - continue - file_key = _extract_file_key(chunk.get("path", "")) - kws = _get_keywords(chunk) - normalized = {_normalize_keyword(k) for k in kws if k} - normalized.discard("") - content = chunk.get("content") or chunk.get("text", "") - existing_data[cid] = (file_key, normalized, content) - - # Build keyword index for existing chunks only - existing_kw_index: Dict[str, List[str]] = defaultdict(list) # kw → [chunk_id] - for cid, (_, kw_set, _) in existing_data.items(): - for kw in kw_set: - existing_kw_index[kw].append(cid) - - connections: Dict[str, List[Dict[str, Any]]] = defaultdict(list) - - # For each new chunk, find candidates in existing chunks - for new_id, (new_file, new_kws, new_content) in new_data.items(): - if not new_kws: - continue - - candidate_shared: Dict[str, set] = defaultdict(set) - for kw in new_kws: - for existing_id in existing_kw_index.get(kw, []): - if cross_only: - existing_file = existing_data[existing_id][0] - if existing_file == new_file: - continue - candidate_shared[existing_id].add(kw) - - # Score and filter - scored: List[Tuple[str, float, set]] = [] - for existing_id, shared_kws in candidate_shared.items(): - if len(shared_kws) < min_overlap: - continue - - existing_kws = existing_data[existing_id][1] - score = _compute_keyword_score( - shared_kws=shared_kws, - kws_a=new_kws, - kws_b=existing_kws, - weight=kw_weight, - ) - if score >= min_score: - # Near-duplicate filter - if max_content_overlap < 1.0: - existing_content = existing_data[existing_id][2] - if new_content and existing_content: - ratio = SequenceMatcher( - None, new_content, existing_content - ).ratio() - if ratio >= max_content_overlap: - continue - scored.append((existing_id, score, shared_kws)) - - scored.sort(key=lambda x: x[1], reverse=True) - for existing_id, score, shared_kws in scored[:max_conns]: - conn = { - "target": existing_id, - "relation": "related", - "score": round(score, 4), - "keywords": sorted(shared_kws), - } - connections[new_id].append(conn) - # Bidirectional: also add reverse edge - connections[existing_id].append( - { - "target": new_id, - "relation": "related", - "score": round(score, 4), - "keywords": sorted(shared_kws), - } - ) - - total = sum(len(v) for v in connections.values()) - logger.info( - f"🔗 Incremental connections: {total} new edges " - f"between {len(new_data)} new chunks and {len(existing_data)} existing chunks" - ) - - return dict(connections) - - -# ─── File-Level Aggregation (v2.0) ─────────────────────────────────────────── - -# Token filtering — same logic as text_utils._is_meaningful_token -_CN_EN_NUM_RE = re.compile(r"[\u4e00-\u9fff]|[A-Za-z]+|\d+(?:\.\d+)?") -_CHUNK_MARKER_RE = re.compile( - rf"{CHUNK_REF_PATTERN}|image-\d+|table-\d+", - re.IGNORECASE, -) - - -def _is_meaningful_token(token: str) -> bool: - """Check if a token is worth keeping (same logic as text_utils).""" - if not _CN_EN_NUM_RE.search(token): - return False - if len(token) == 1: - return False - if re.fullmatch(r"\d+(?:\.\d+)?", token): - return False - return True - - -def _extract_tokens_from_content(content: str) -> List[str]: - """Extract meaningful tokens from content using jieba (regex fallback).""" - content = _CHUNK_MARKER_RE.sub("", content) - # Strip HTML tags and entities (table chunks contain raw HTML) - content = re.sub(r"<[^>]+>", " ", content) - content = re.sub(r"&\w+;", " ", content) - try: - import jieba - - if hasattr(jieba, "lcut"): - raw = jieba.lcut(content) - else: - raw = list(jieba.cut(content)) - except ImportError: - raw = re.split(r"[\s,;，；。！？、\-/]+", content) - return [t for t in raw if _is_meaningful_token(t)] - - -def _get_chunk_keywords(chunk: Dict[str, Any]) -> List[str]: - """Get keywords for a chunk; falls back to tokens from content if empty.""" - keywords = _get_keywords(chunk) - meaningful = [k for k in keywords if _is_meaningful_token(k)] - if meaningful: - return meaningful - content = chunk.get("content") or chunk.get("text", "") - if not content: - return [] - tokens = _extract_tokens_from_content(content) - seen = set() - unique = [] - for t in tokens: - normalized = _normalize_keyword(t) - if normalized and normalized not in seen: - seen.add(normalized) - unique.append(normalized) - return unique - - -def _compute_tfidf_top_keywords( - file_chunks: Dict[str, List[Dict[str, Any]]], - top_k: int = 6, -) -> Dict[str, List[str]]: - """ - TF-IDF top keywords per file. - TF = chunks in file containing keyword. IDF = log(total_files / files_with_keyword). - """ - total_files = len(file_chunks) - if total_files == 0: - return {} - - file_kw_tf: Dict[str, Dict[str, int]] = {} - doc_freq: Dict[str, int] = defaultdict(int) - - for fk, chunks in file_chunks.items(): - kw_count: Dict[str, int] = defaultdict(int) - file_kw_set: set = set() - for chunk in chunks: - for kw in _get_chunk_keywords(chunk): - normalized = _normalize_keyword(kw) - if normalized: - kw_count[normalized] += 1 - file_kw_set.add(normalized) - file_kw_tf[fk] = dict(kw_count) - for kw in file_kw_set: - doc_freq[kw] += 1 - - result: Dict[str, List[str]] = {} - for fk, kw_count in file_kw_tf.items(): - scored = [] - for kw, tf in kw_count.items(): - if total_files == 1: - score = tf # Single-file KB: pure frequency - else: - idf = ( - math.log(total_files / doc_freq[kw]) - if doc_freq[kw] < total_files - else 0.1 - ) - score = tf * idf - scored.append((score, tf, kw)) - scored.sort(key=lambda x: (x[0], x[1]), reverse=True) - result[fk] = [kw for _, _, kw in scored[:top_k]] - - return result - - -def _compute_file_importance( - chunk_ids: List[str], - chunk_stats: Dict[str, Dict[str, Any]], - half_life_days: float = 30.0, - alpha: float = 0.7, - beta: float = 0.3, -) -> float: - """importance = α × usage_heat + β × freshness""" - if not chunk_ids: - return 0.0 - total_relevance = 0.0 - earliest_created = None - for cid in chunk_ids: - stat = chunk_stats.get(cid, {}) - hc = stat.get("hit_count", 0) - lh = stat.get("last_hit") - ca = stat.get("created_at") - if hc > 0 and lh: - total_relevance += relevance_score(hc, lh, half_life_days) - if ca and (earliest_created is None or ca < earliest_created): - earliest_created = ca - usage_heat = total_relevance / len(chunk_ids) - freshness = ( - relevance_score(1, earliest_created, half_life_days) - if earliest_created - else 1.0 - ) - return round(alpha * usage_heat + beta * freshness, 4) - - -def _aggregate_file_level_edges( - chunk_edges: List[Dict[str, Any]], - chunk_to_file: Dict[str, str], - chunk_paths: Optional[Dict[str, str]] = None, - max_top_connections: int = 10, -) -> List[Dict[str, Any]]: - """ - Aggregate chunk-level edges into file-level edges. - Shows top_connections with readable chunk names instead of raw keywords. - """ - if chunk_paths is None: - chunk_paths = {} - - pair_data: Dict[Tuple[str, str], Dict[str, List[Dict[str, Any]]]] = {} - for edge in chunk_edges: - src_id = str(edge.get("source", "") or "") - tgt_id = str(edge.get("target", "") or "") - sf = chunk_to_file.get(src_id, "") - tf = chunk_to_file.get(tgt_id, "") - if not sf or not tf or sf == tf: - continue - src_path = chunk_paths.get(src_id) or src_id - tgt_path = chunk_paths.get(tgt_id) or tgt_id - src_name = src_path.rsplit("/", 1)[-1] if "/" in src_path else src_path - tgt_name = tgt_path.rsplit("/", 1)[-1] if "/" in tgt_path else tgt_path - - if sf <= tf: - pk: Tuple[str, str] = (sf, tf) - connection = { - "source_chunk": src_name, - "source_id": src_id, - "target_chunk": tgt_name, - "target_id": tgt_id, - "relation": edge.get("relation", "related"), - "score": edge.get("score", 0), - } - else: - pk = (tf, sf) - connection = { - "source_chunk": tgt_name, - "source_id": tgt_id, - "target_chunk": src_name, - "target_id": src_id, - "relation": edge.get("relation", "related"), - "score": edge.get("score", 0), - } - - if pk not in pair_data: - pair_data[pk] = {"connections": []} - pair_data[pk]["connections"].append(connection) - - file_edges = [] - for (f1, f2), data in pair_data.items(): - conns = data["connections"] - # Sort by score desc, take top N - conns.sort(key=lambda x: x["score"], reverse=True) - scores = [c["score"] for c in conns] - file_edges.append( - { - "source": f1, - "target": f2, - "connection_count": len(conns), - "avg_score": round(sum(scores) / len(scores), 4) if scores else 0, - "top_connections": conns[:max_top_connections], - } - ) - file_edges.sort(key=lambda x: x["connection_count"], reverse=True) - return file_edges - - -# ─── Main API ──────────────────────────────────────────────────────────────── - - -def _get_source_file(chunk: Dict[str, Any]) -> str: - """ - Get the source document file for a chunk. - Uses `_source_file` tag (injected by build_and_deploy / _load_all_chunks_from_kb) - for correct grouping of images/tables with their parent document. - Falls back to `_extract_file_key` for backwards compatibility. - """ - sf = chunk.get("_source_file") - if sf: - return sf - return _extract_file_key(chunk.get("path", "")) - - -def build_knowledge_graph( - all_chunks: List[Dict[str, Any]], - connections: Dict[str, List[Dict[str, Any]]], - kb_id: str = "", - chunk_stats: Optional[Dict[str, Dict[str, Any]]] = None, - file_summaries: Optional[Dict[str, str]] = None, -) -> Dict[str, Any]: - """Build a file-level knowledge graph (v2.0).""" - if chunk_stats is None: - chunk_stats = {} - - file_chunks: Dict[str, List[Dict[str, Any]]] = defaultdict(list) - chunk_to_file: Dict[str, str] = {} - chunk_paths: Dict[str, str] = {} - for chunk in all_chunks: - cid = str(chunk.get("chunk_id") or chunk.get("know_id", "")) - fk = _get_source_file(chunk) - if fk: - file_chunks[fk].append(chunk) - if cid: - chunk_to_file[cid] = fk - chunk_paths[cid] = chunk.get("path", "") - - file_keywords = _compute_tfidf_top_keywords(file_chunks) - chunk_edges = _connections_to_edges(connections) - file_edges = _aggregate_file_level_edges(chunk_edges, chunk_to_file, chunk_paths) - - files_dict = {} - for fk, chunks in file_chunks.items(): - types_count: Dict[str, int] = defaultdict(int) - cids = [] - for c in chunks: - types_count[c.get("type", "text")] += 1 - cid = str(c.get("chunk_id") or c.get("know_id", "")) - if cid: - cids.append(cid) - files_dict[fk] = { - "chunks_count": len(chunks), - "types": dict(types_count), - "top_keywords": file_keywords.get(fk, []), - "top_summary": (file_summaries or {}).get(fk, ""), - "importance": _compute_file_importance(cids, chunk_stats), - "created_at": datetime.now(timezone.utc).isoformat(), - } - - total_chunks = sum(f["chunks_count"] for f in files_dict.values()) - graph = { - "version": "2.0", - "updated_at": datetime.now(timezone.utc).isoformat(), - "kb_id": kb_id, - "stats": { - "total_files": len(files_dict), - "total_chunks": total_chunks, - "total_cross_file_edges": len(file_edges), - }, - "files": files_dict, - "edges": file_edges, - } - logger.info( - f"📊 Knowledge graph built: " - f"{graph['stats']['total_files']} files, " - f"{graph['stats']['total_chunks']} chunks, " - f"{graph['stats']['total_cross_file_edges']} edges" - ) - return graph - - -def update_knowledge_graph( - existing_graph: Dict[str, Any], - new_chunks: List[Dict[str, Any]], - existing_chunks: List[Dict[str, Any]], - kb_id: str = "", - connect_config: Optional[Dict[str, Any]] = None, - chunk_stats: Optional[Dict[str, Dict[str, Any]]] = None, - file_summaries: Optional[Dict[str, str]] = None, - new_connections: Optional[Dict[str, List[Dict[str, Any]]]] = None, -) -> Dict[str, Any]: - """Incrementally update a file-level knowledge graph with new chunks.""" - if chunk_stats is None: - chunk_stats = {} - - all_combined = existing_chunks + new_chunks - file_chunks: Dict[str, List[Dict[str, Any]]] = defaultdict(list) - chunk_to_file: Dict[str, str] = {} - chunk_paths: Dict[str, str] = {} - for chunk in all_combined: - cid = str(chunk.get("chunk_id") or chunk.get("know_id", "")) - fk = _get_source_file(chunk) - if fk: - file_chunks[fk].append(chunk) - if cid: - chunk_to_file[cid] = fk - chunk_paths[cid] = chunk.get("path", "") - - file_keywords = _compute_tfidf_top_keywords(file_chunks) - - if new_connections is None: - new_connections = _incremental_connections( - new_chunks=new_chunks, - existing_chunks=existing_chunks, - config=connect_config, - ) - new_chunk_edges = _connections_to_edges(new_connections) - existing_file_edges = existing_graph.get("edges", []) - new_file_edges = _aggregate_file_level_edges( - new_chunk_edges, chunk_to_file, chunk_paths - ) - - # Merge file edges - merged_map: Dict[Tuple[str, str], Dict] = {} - for edge in existing_file_edges + new_file_edges: - pk = tuple(sorted([edge["source"], edge["target"]])) - if pk not in merged_map: - merged_map[pk] = edge - else: - old = merged_map[pk] - # Merge connections, dedup by chunk pair - all_conns = old.get("top_connections", []) + edge.get("top_connections", []) - seen = set() - deduped = [] - for c in all_conns: - pair = (c.get("source_chunk", ""), c.get("target_chunk", "")) - if pair not in seen: - seen.add(pair) - deduped.append(c) - deduped.sort(key=lambda x: x.get("score", 0), reverse=True) - tc = old["connection_count"] + edge["connection_count"] - scores = [c.get("score", 0) for c in deduped] - avg = sum(scores) / len(scores) if scores else 0 - merged_map[pk] = { - "source": pk[0], - "target": pk[1], - "connection_count": tc, - "avg_score": round(avg, 4), - "top_connections": deduped[:10], - } - all_file_edges = sorted( - merged_map.values(), key=lambda x: x["connection_count"], reverse=True - ) - - existing_files = existing_graph.get("files", {}) - files_dict = {} - new_file_count = 0 - for fk, chunks in file_chunks.items(): - types_count: Dict[str, int] = defaultdict(int) - cids = [] - for c in chunks: - types_count[c.get("type", "text")] += 1 - cid = str(c.get("chunk_id") or c.get("know_id", "")) - if cid: - cids.append(cid) - created_at = existing_files.get(fk, {}).get( - "created_at", datetime.now(timezone.utc).isoformat() - ) - if fk not in existing_files: - new_file_count += 1 - files_dict[fk] = { - "chunks_count": len(chunks), - "types": dict(types_count), - "top_keywords": file_keywords.get(fk, []), - "top_summary": (file_summaries or {}).get(fk, "") or existing_files.get(fk, {}).get("top_summary", ""), - "importance": _compute_file_importance(cids, chunk_stats), - "created_at": created_at, - } - - total_chunks = sum(f["chunks_count"] for f in files_dict.values()) - graph = { - "version": "2.0", - "updated_at": datetime.now(timezone.utc).isoformat(), - "kb_id": kb_id or existing_graph.get("kb_id", ""), - "stats": { - "total_files": len(files_dict), - "total_chunks": total_chunks, - "total_cross_file_edges": len(all_file_edges), - }, - "files": files_dict, - "edges": all_file_edges, - } - logger.info( - f"📊 Knowledge graph updated: " - f"+{new_file_count} files → " - f"total {graph['stats']['total_files']} files, " - f"{graph['stats']['total_chunks']} chunks, " - f"{graph['stats']['total_cross_file_edges']} edges" - ) - return graph - - -# ─── Configuration ──────────────────────────────────────────────────────────── - -KNOWHERE_HOME = os.path.expanduser(os.environ.get("KNOWHERE_HOME", "~/.knowhere")) - - -def _get_kb_dir(kb_id: str) -> str: - """Get the knowledge base directory path.""" - return os.path.join(KNOWHERE_HOME, kb_id) - - -def _get_kg_path(kb_id: str) -> str: - """Get the knowledge_graph.json path for a KB.""" - return os.path.join(_get_kb_dir(kb_id), "knowledge_graph.json") - - -def _get_stats_path(kb_id: str) -> str: - """Get the chunk_stats.json path for a KB.""" - return os.path.join(_get_kb_dir(kb_id), "chunk_stats.json") - - -def _empty_knowledge_graph(kb_id: str) -> Dict[str, Any]: - """Build an empty v2 knowledge graph for a KB with no local chunks.""" - return { - "version": "2.0", - "updated_at": datetime.now(timezone.utc).isoformat(), - "kb_id": kb_id, - "stats": { - "total_files": 0, - "total_chunks": 0, - "total_cross_file_edges": 0, - }, - "files": {}, - "edges": [], - } - - -# ─── Chunk Usage Tracking ───────────────────────────────────────────────────── - - -def load_chunk_stats(kb_id: str) -> Dict[str, Dict[str, Any]]: - """Load chunk usage stats from chunk_stats.json.""" - path = _get_stats_path(kb_id) - if not os.path.exists(path): - return {} - try: - with open(path, "r", encoding="utf-8") as f: - return json.load(f) - except (json.JSONDecodeError, IOError): - return {} - - -def record_chunk_hits( - kb_id: str, - chunk_ids: List[str], -) -> None: - """ - Record that chunks were accessed (returned in search results). - Updates hit_count and last_hit for each chunk. - - Args: - kb_id: Knowledge base ID. - chunk_ids: List of chunk IDs that were hit. - """ - stats = load_chunk_stats(kb_id) - now = datetime.now(timezone.utc).isoformat() - - for cid in chunk_ids: - if cid not in stats: - stats[cid] = { - "hit_count": 0, - "first_hit": now, - "last_hit": now, - "created_at": now, - } - stats[cid]["hit_count"] += 1 - stats[cid]["last_hit"] = now - - path = _get_stats_path(kb_id) - os.makedirs(os.path.dirname(path), exist_ok=True) - with open(path, "w", encoding="utf-8") as f: - json.dump(stats, f, ensure_ascii=False, indent=2) - - -def relevance_score( - hit_count: int, - last_hit_iso: str, - half_life_days: float = 30.0, -) -> float: - """ - Compute relevance score with exponential decay. - Higher hit_count + more recent access → higher score. - - Args: - hit_count: Number of times this chunk was accessed. - last_hit_iso: ISO timestamp of last access. - half_life_days: Days until relevance halves. - - Returns: - Decay-weighted score. - """ - try: - last_hit_dt = datetime.fromisoformat(last_hit_iso) - days_since = (datetime.now(timezone.utc) - last_hit_dt).total_seconds() / 86400 - except (ValueError, TypeError): - days_since = 0 - - decay = math.exp(-0.693 * days_since / half_life_days) - return hit_count * decay - - -# ─── File I/O ───────────────────────────────────────────────────────────────── - - -def save_knowledge_graph(graph: Dict[str, Any], output_path: str) -> str: - """Save knowledge graph to a JSON file.""" - os.makedirs(os.path.dirname(output_path), exist_ok=True) - with open(output_path, "w", encoding="utf-8") as f: - json.dump(graph, f, ensure_ascii=False, indent=2) - logger.info(f"💾 Knowledge graph saved: {output_path}") - return output_path - - -def load_knowledge_graph(path: str) -> Optional[Dict[str, Any]]: - """Load an existing knowledge graph from JSON file.""" - if not os.path.exists(path): - return None - try: - with open(path, "r", encoding="utf-8") as f: - return json.load(f) - except (json.JSONDecodeError, IOError) as e: - logger.warning(f"Failed to load knowledge graph from {path}: {e}") - return None - - -def _save_chunks(chunks: List[Dict[str, Any]], output_path: str) -> None: - """Save chunks data to a JSON file in the standard format.""" - os.makedirs(os.path.dirname(output_path), exist_ok=True) - with open(output_path, "w", encoding="utf-8") as f: - json.dump({"chunks": chunks}, f, ensure_ascii=False, indent=2) - - -def _load_chunks(path: str) -> List[Dict[str, Any]]: - """Load chunks from a stored chunks.json file.""" - if not os.path.exists(path): - return [] - try: - with open(path, "r", encoding="utf-8") as f: - data = json.load(f) - if isinstance(data, dict) and "chunks" in data: - return data["chunks"] - if isinstance(data, list): - return data - except (json.JSONDecodeError, IOError): - pass - return [] - - -def extract_chunks_from_graph(graph: Dict[str, Any]) -> List[Dict[str, Any]]: - """ - Reconstruct minimal chunk dicts from graph for incremental matching. - This is a last-resort fallback when subdirectory chunks.json files are unavailable. - v2.0: no chunk-level data in graph; returns empty list. - Legacy: falls back to node_index or nodes array. - """ - chunks = [] - # v2.0: files dict doesn't store chunk IDs, return empty - if graph.get("version", "").startswith("2."): - return chunks - # Legacy v1.x: handle node_index - node_index = graph.get("node_index", {}) - if node_index: - for chunk_id, file_key in node_index.items(): - chunks.append( - { - "chunk_id": chunk_id, - "path": file_key, - "content": "", - "metadata": {"keywords": []}, - } - ) - return chunks - # Legacy: handle old nodes array - for node in graph.get("nodes", []): - chunks.append( - { - "chunk_id": node["id"], - "path": node.get("path", ""), - "content": node.get("content_preview", ""), - "metadata": {"keywords": node.get("keywords", [])}, - } - ) - return chunks - - -# ─── Chunk ID Dedup ────────────────────────────────────────────────────────── - - -def _dedup_chunks_by_content( - new_chunks: List[Dict[str, Any]], - existing_chunks: List[Dict[str, Any]], -) -> List[Dict[str, Any]]: - """ - Filter new_chunks: discard any whose chunk_id already exists in existing_chunks. - - Since all parsers now generate deterministic know_id (content-hash based), - identical content always produces the same chunk_id. Simple set comparison - replaces the old strip+hash pipeline. - - Returns: - List of new chunks that have no chunk_id duplicate in existing_chunks. - """ - existing_ids = { - str(c.get("chunk_id") or c.get("know_id", "")) for c in existing_chunks - } - existing_ids.discard("") - - deduped = [] - skipped = 0 - for chunk in new_chunks: - cid = str(chunk.get("chunk_id") or chunk.get("know_id", "")) - if cid and cid in existing_ids: - skipped += 1 - else: - deduped.append(chunk) - - if skipped > 0: - logger.info( - f"📊 chunk dedup: {skipped} duplicate chunks skipped " - f"(by chunk_id), {len(deduped)} chunks to add" - ) - return deduped - - -def _load_all_chunks_from_kb(kb_dir: str) -> List[Dict[str, Any]]: - """ - Load all chunks from per-file chunks.json files under a KB directory. - Tags each chunk with _source_file = subdirectory name (= source document). - """ - all_chunks = [] - for entry in os.listdir(kb_dir): - entry_path = os.path.join(kb_dir, entry) - if not os.path.isdir(entry_path): - continue - chunks_file = os.path.join(entry_path, "chunks.json") - if os.path.isfile(chunks_file): - loaded = _load_chunks(chunks_file) - for chunk in loaded: - chunk["_source_file"] = entry - all_chunks.extend(loaded) - return all_chunks - - -def _source_files_from_chunks(chunks: List[Dict[str, Any]]) -> set[str]: - """Return the source-file set represented by loaded KB chunks.""" - return { - str(chunk.get("_source_file") or "").strip() - for chunk in chunks - if str(chunk.get("_source_file") or "").strip() - } - - -def _prune_chunk_stats(kb_id: str, chunks: List[Dict[str, Any]]) -> None: - """Remove chunk_stats entries whose chunks no longer exist on disk.""" - stats_path = _get_stats_path(kb_id) - if not os.path.exists(stats_path): - return - - stats = load_chunk_stats(kb_id) - live_chunk_ids = { - str(chunk.get("chunk_id") or chunk.get("know_id", "")) - for chunk in chunks - if chunk.get("chunk_id") or chunk.get("know_id") - } - pruned = {cid: data for cid, data in stats.items() if cid in live_chunk_ids} - if len(pruned) == len(stats): - return - - os.makedirs(os.path.dirname(stats_path), exist_ok=True) - with open(stats_path, "w", encoding="utf-8") as f: - json.dump(pruned, f, ensure_ascii=False, indent=2) - logger.info( - f"📊 Chunk stats pruned: {len(stats) - len(pruned)} stale chunks removed" - ) - - -def sync_knowledge_graph_with_local_files( - kb_id: str, - connect_config: Optional[Dict[str, Any]] = None, - summary_use_llm: bool = False, -) -> Dict[str, Any]: - """Synchronize knowledge_graph.json with current ~/.knowhere/{kb_id} files. - - This is intentionally a no-op when graph files match on-disk document - directories. If a user manually deletes a local parsed document directory, - the graph is rebuilt from remaining chunks and stale chunk_stats entries are - removed. - """ - from app.services.connect_builder.builder import build_connections - from app.services.connect_builder.summary_builder import enrich_doc_nav_summaries - - kb_dir = _get_kb_dir(kb_id) - kg_path = _get_kg_path(kb_id) - os.makedirs(kb_dir, exist_ok=True) - - existing_graph = load_knowledge_graph(kg_path) - chunks_on_disk = _load_all_chunks_from_kb(kb_dir) - disk_files = _source_files_from_chunks(chunks_on_disk) - graph_files = set((existing_graph or {}).get("files", {}).keys()) - - if existing_graph is not None and graph_files == disk_files: - _prune_chunk_stats(kb_id, chunks_on_disk) - return existing_graph - - removed_files = sorted(graph_files - disk_files) - added_files = sorted(disk_files - graph_files) - logger.info( - "📊 Syncing Knowledge Graph with local files: " - f"removed={removed_files}, added={added_files}" - ) - - if not chunks_on_disk: - graph = _empty_knowledge_graph(kb_id) - save_knowledge_graph(graph, kg_path) - _prune_chunk_stats(kb_id, []) - return graph - - try: - file_summaries = enrich_doc_nav_summaries( - kb_dir=kb_dir, - source_file=None, - use_llm=summary_use_llm, - ) - except Exception as e: - logger.warning(f"sync enrich_doc_nav_summaries failed: {e}") - file_summaries = {} - - stats = load_chunk_stats(kb_id) - connections = build_connections(chunks_on_disk, connect_config) - _merge_related_connections_into_chunks(chunks_on_disk, connections) - _save_chunks_by_source_file(kb_dir, chunks_on_disk) - graph = build_knowledge_graph( - all_chunks=chunks_on_disk, - connections=connections, - kb_id=kb_id, - chunk_stats=stats, - file_summaries=file_summaries, - ) - save_knowledge_graph(graph, kg_path) - _prune_chunk_stats(kb_id, chunks_on_disk) - return graph - - -# ─── MCP Auto-Registration ─────────────────────────────────────────────────── - - -def _get_mcp_server_path() -> str: - """Get the absolute path to the MCP server script. - - Points to the consolidated knowhere-mcp/server.py (unified server - with both Cloud API and local search tools). - """ - # Navigate from graph_builder.py → project root → knowhere-mcp/server.py - # graph_builder.py is at: apps/worker/app/services/connect_builder/ - project_root = os.path.normpath( - os.path.join( - os.path.dirname(os.path.abspath(__file__)), - "..", - "..", - "..", - "..", - "..", - ) - ) - return os.path.join(project_root, "knowhere-mcp", "server.py") - - -def _auto_register_mcp() -> None: - """ - Detect installed Agent products and auto-register the knowhere MCP server. - Only runs on first deploy (when ~/.knowhere/ is freshly created). - - Supported products: - - Cursor: ~/.cursor/mcp.json - - Claude Code: ~/.claude.json (project-level) or ~/.claude/claude_code_config.json - """ - mcp_server_path = os.path.normpath(_get_mcp_server_path()) - home = os.path.expanduser("~") - - knowhere_mcp_entry = { - "command": "python3", - "args": [mcp_server_path], - "env": { - "KNOWHERE_API_KEY": os.environ.get("KNOWHERE_API_KEY", ""), - }, - } - - registered = [] - - # ── Cursor ──────────────────────────────────────────────────────────── - cursor_mcp = os.path.join(home, ".cursor", "mcp.json") - if os.path.isdir(os.path.join(home, ".cursor")): - try: - existing = {} - if os.path.exists(cursor_mcp): - with open(cursor_mcp, "r") as f: - existing = json.load(f) - - servers = existing.get("mcpServers", {}) - # Update even if "knowhere" exists (to point to new server) - if "knowhere" not in servers or "mcp/knowhere_mcp_server" in str( - servers.get("knowhere", {}).get("args", []) - ): - servers["knowhere"] = knowhere_mcp_entry - existing["mcpServers"] = servers - with open(cursor_mcp, "w") as f: - json.dump(existing, f, indent=2) - registered.append("Cursor") - except Exception as e: - logger.debug(f"Cursor MCP registration skipped: {e}") - - # ── Claude Code ─────────────────────────────────────────────────────── - claude_config = os.path.join(home, ".claude.json") - if os.path.exists(claude_config) or os.path.isdir(os.path.join(home, ".claude")): - try: - existing = {} - if os.path.exists(claude_config): - with open(claude_config, "r") as f: - existing = json.load(f) - - servers = existing.get("mcpServers", {}) - if "knowhere" not in servers or "mcp/knowhere_mcp_server" in str( - servers.get("knowhere", {}).get("args", []) - ): - servers["knowhere"] = knowhere_mcp_entry - existing["mcpServers"] = servers - with open(claude_config, "w") as f: - json.dump(existing, f, indent=2) - registered.append("Claude Code") - except Exception as e: - logger.debug(f"Claude Code MCP registration skipped: {e}") - - if registered: - logger.info(f"🔌 MCP auto-registered for: {', '.join(registered)}") - else: - logger.debug("No Agent products detected for MCP auto-registration") - - -# ─── doc_nav section extraction for GraphNode persistence ──────────────────── - - - - - -# ─── One-Stop API ───────────────────────────────────────────────────────────── - - -def build_and_deploy( - chunks: List[Dict[str, Any]], - kb_id: str, - parsed_output_dir: Optional[str] = None, - connect_config: Optional[Dict[str, Any]] = None, - rebuild_all: bool = True, - summary_use_llm: bool = False, -) -> Dict[str, Any]: - """ - One-stop knowledge graph build/update + deploy to ~/.knowhere/ + MCP register. - - This is the main entry point for callers (parse services, debug scripts, etc). - Callers just provide chunks + kb_id; everything else is automatic. - - Flow: - 1. If parsed_output_dir provided → copy full parsed output to ~/.knowhere/{kb_id}/data/ - 2. Check if ~/.knowhere/{kb_id}/knowledge_graph.json exists - - No → build_knowledge_graph() (full build) - - rebuild_all=True → scan KB dir for existing files, merge with new chunks - - rebuild_all=False → only use the new chunks (ignore previous files) - - Yes → update_knowledge_graph() (incremental) - 3. Save knowledge_graph.json to ~/.knowhere/{kb_id}/ - 4. On first-ever deploy → _auto_register_mcp() - - Args: - chunks: Parsed chunks from the current file. - kb_id: Knowledge base identifier (e.g. dataset name). - parsed_output_dir: Path to the parsed output directory (add_dir) containing - images, tables, doc_nav.json etc. If provided, its contents are - copied to ~/.knowhere/{kb_id}/data/{dirname}/. - connect_config: Optional config overrides for connect_builder. - rebuild_all: When knowledge_graph.json is missing, whether to scan the - KB directory for existing chunk data and include them in the full - rebuild. Defaults to True. Set to False to only process the new - chunks (legacy behavior). - summary_use_llm: If True, use LLM to generate coherent hierarchical - summaries (slow, costs API tokens). If False (default), use - lightweight title enumeration (e.g. "This section covers: Section 1, - Section 2"). Only affects `top_summary` and `summary` fields. - - Returns: - The knowledge graph dict. - """ - import shutil - - from app.services.connect_builder.builder import build_connections - - kg_path = _get_kg_path(kb_id) - kb_dir = _get_kb_dir(kb_id) - - # Detect if this is a first-ever deploy (for MCP registration) - first_deploy = not os.path.exists(KNOWHERE_HOME) - - # Ensure directory exists - os.makedirs(kb_dir, exist_ok=True) - - # Load existing state BEFORE deploy (to avoid counting new file's chunks twice) - # Determine source_file early so we can exclude it from existing_chunks - source_file = ( - os.path.basename(parsed_output_dir) - if parsed_output_dir and os.path.isdir(parsed_output_dir) - else None - ) - existing_graph = load_knowledge_graph(kg_path) - - # Load chunks once and reuse — avoids the double-load where - # sync_knowledge_graph_with_local_files internally calls - # _load_all_chunks_from_kb and then we call it again. - all_on_disk: List[Dict[str, Any]] = [] - if existing_graph is not None: - all_on_disk = _load_all_chunks_from_kb(kb_dir) - disk_files = _source_files_from_chunks(all_on_disk) - graph_files = set(existing_graph.get("files", {}).keys()) - - if graph_files != disk_files: - # Disk state diverged from graph → full sync rebuild. - # sync will reload chunks internally (it backfills connect_to - # metadata), so we reload afterwards to pick up changes. - existing_graph = sync_knowledge_graph_with_local_files( - kb_id=kb_id, - connect_config=connect_config, - summary_use_llm=summary_use_llm, - ) - all_on_disk = _load_all_chunks_from_kb(kb_dir) - else: - # Files match — fast-path: just prune stale chunk_stats. - _prune_chunk_stats(kb_id, all_on_disk) - - if existing_graph is not None: - if not all_on_disk: - existing_chunks = extract_chunks_from_graph(existing_graph) - else: - # Exclude chunks from the current source file — they may already - # be on disk if parsed_output_dir is inside kb_dir (debug_parse). - # Without this filter, _dedup_chunks_by_content would treat them - # as "existing" and skip the incremental update entirely. - existing_chunks = ( - [c for c in all_on_disk if c.get("_source_file") != source_file] - if source_file - else all_on_disk - ) - else: - existing_chunks = [] - - # ── Deploy parsed output (images, tables, hierarchy, etc.) ── - if parsed_output_dir and os.path.isdir(parsed_output_dir) and source_file is not None: - deploy_target = os.path.join(kb_dir, source_file) - - # Skip copy if parsed output is already in the target location - parsed_abs = os.path.normpath(os.path.abspath(parsed_output_dir)) - target_abs = os.path.normpath(os.path.abspath(deploy_target)) - if parsed_abs == target_abs: - logger.info( - f"📂 Parsed output already at target: {deploy_target} (skip copy)" - ) - else: - if os.path.exists(deploy_target): - shutil.rmtree(deploy_target) - shutil.copytree(parsed_output_dir, deploy_target) - # Delete ZIP files from deployed directory (no longer needed) - import glob - - for zip_file in glob.glob(os.path.join(deploy_target, "*.zip")): - os.remove(zip_file) - logger.info(f"📂 Parsed output deployed: {deploy_target}") - - # Tag all chunks with source document for correct file-level grouping - if source_file: - for chunk in chunks: - chunk["_source_file"] = source_file - - # ── Generate hierarchical summaries ── - from app.services.connect_builder.summary_builder import enrich_doc_nav_summaries - try: - file_summaries = enrich_doc_nav_summaries( - kb_dir=kb_dir, - source_file=source_file, - use_llm=summary_use_llm, - ) - except Exception as e: - logger.warning(f"doc_nav summary enrichment failed: {e}") - file_summaries = {} - - # Load chunk_stats for importance calculation - stats = load_chunk_stats(kb_id) - - if existing_graph is None: - # ── First build: full ── - if rebuild_all: - # Scan KB dir for existing chunk data (deploy already happened, - # so the new file's chunks are on disk if parsed_output_dir was given). - all_on_disk = _load_all_chunks_from_kb(kb_dir) - if source_file and all_on_disk: - # New file already deployed → all_on_disk includes it, no merge needed - all_chunks = all_on_disk - else: - # New file not deployed to disk (no parsed_output_dir), - # or KB dir was empty → merge in-memory chunks with disk data. - # Dedup by chunk_id to prevent double-counting. - seen_ids = { - str(c.get("chunk_id") or c.get("know_id", "")) for c in all_on_disk - } - extra = [ - c - for c in chunks - if str(c.get("chunk_id") or c.get("know_id", "")) not in seen_ids - ] - all_chunks = all_on_disk + extra - # Full rebuild: generate summaries for ALL files, not just source_file - from app.services.connect_builder.summary_builder import enrich_doc_nav_summaries as _enrich_nav - try: - all_nav_summaries = _enrich_nav( - kb_dir=kb_dir, - source_file=None, - use_llm=summary_use_llm, - ) - file_summaries.update(all_nav_summaries) - except Exception as e: - logger.warning(f"Full rebuild enrich_doc_nav_summaries failed: {e}") - logger.info( - f"📊 rebuild Knowledge Graph " - f"(rebuild_all=True, {len(all_chunks)} chunks from KB dir) ..." - ) - else: - all_chunks = chunks - logger.info( - "📊 rebuild Knowledge Graph (rebuild_all=False, new chunks only) ..." - ) - - connections = build_connections(all_chunks, connect_config) - _merge_related_connections_into_chunks(all_chunks, connections) - _save_chunks_by_source_file(kb_dir, all_chunks) - stats_chunks = all_chunks - graph = build_knowledge_graph( - all_chunks=all_chunks, - connections=connections, - kb_id=kb_id, - chunk_stats=stats, - file_summaries=file_summaries, - ) - else: - # ── Incremental update ── - # Content-hash dedup: discard new chunks identical to existing ones - # to preserve established graph edges and relationships. - deduped_new = _dedup_chunks_by_content(chunks, existing_chunks) - if len(deduped_new) == 0: - logger.info( - "📊 All new chunks are duplicates of existing data, " - "skipping incremental update" - ) - stats_chunks = existing_chunks - graph = existing_graph - # Still inject summaries even if chunks unchanged - for fk, fdata in graph.get("files", {}).items(): - if file_summaries and fk in file_summaries and not fdata.get("top_summary"): - fdata["top_summary"] = file_summaries[fk] - else: - logger.info( - f"📊 incremental update Knowledge Graph " - f"({len(deduped_new)} new, {len(chunks) - len(deduped_new)} skipped) ..." - ) - related_connections = _incremental_connections( - new_chunks=deduped_new, - existing_chunks=existing_chunks, - config=connect_config, - ) - _merge_related_connections_into_chunks( - existing_chunks + deduped_new, related_connections - ) - _save_chunks_by_source_file(kb_dir, existing_chunks + deduped_new) - stats_chunks = existing_chunks + deduped_new - graph = update_knowledge_graph( - existing_graph=existing_graph, - new_chunks=deduped_new, - existing_chunks=existing_chunks, - kb_id=kb_id, - connect_config=connect_config, - chunk_stats=stats, - file_summaries=file_summaries, - new_connections=related_connections, - ) - - # Save graph - save_knowledge_graph(graph, kg_path) - - # Initialize chunk_stats.json with created_at for all new chunks - stats_path = _get_stats_path(kb_id) - now = datetime.now(timezone.utc).isoformat() - updated = False - for chunk in stats_chunks: - cid = str(chunk.get("chunk_id") or chunk.get("know_id", "")) - if cid and cid not in stats: - stats[cid] = { - "hit_count": 0, - "first_hit": None, - "last_hit": None, - "created_at": now, - } - updated = True - if updated: - os.makedirs(os.path.dirname(stats_path), exist_ok=True) - with open(stats_path, "w", encoding="utf-8") as f: - json.dump(stats, f, ensure_ascii=False, indent=2) - logger.info(f"📊 Chunk stats initialized: {len(stats)} chunks tracked") - - logger.info( - f"✅ Knowledge Graph deployed to {kb_dir}: " - f"{graph['stats']['total_files']} files, " - f"{graph['stats']['total_chunks']} chunks, " - f"{graph['stats']['total_cross_file_edges']} edges" - ) - - # Auto-register MCP on first deploy - if first_deploy: - try: - _auto_register_mcp() - except Exception as e: - logger.debug(f"MCP auto-registration skipped: {e}") - - return graph diff --git a/apps/worker/app/services/connect_builder/summary_builder.py b/apps/worker/app/services/connect_builder/summary_builder.py index cac4b38bb..ae020e429 100644 --- a/apps/worker/app/services/connect_builder/summary_builder.py +++ b/apps/worker/app/services/connect_builder/summary_builder.py @@ -7,9 +7,7 @@ Usage (standalone): from app.services.connect_builder.summary_builder import enrich_doc_nav_summaries - enrich_doc_nav_summaries(kb_dir, source_file="report.pdf") - -Called by graph_builder.build_and_deploy() after file deploy, before KG build. + enrich_doc_nav_summaries(document_workspace_dir, source_file="report.pdf") """ import json @@ -30,16 +28,6 @@ NON_LLM_TOP_SUMMARY_MAX_SECTIONS = 20 NON_LLM_TOP_SUMMARY_MAX_DEPTH = 2 -_TREE_EXCLUDED_TITLES = {"root", "images", "tables"} -_TREE_TITLE_MAX_TOKENS_START = 20 -_TREE_TITLE_MAX_TOKENS_END = 5 -_TITLE_ENUM_PREFIXES = ( - "this section covers:", - "this section includes", - "this document covers:", - "this document includes", -) - # ─── LLM Interface ─────────────────────────────────────────────────────────── @@ -52,7 +40,7 @@ def _llm_summarize(snippets_text: str, node_name: str, max_tokens: int = 100) -> """ try: from shared.services.ai.prompt_service import build_prompt, _detect_text_language - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client # Deterministic language lock — see prompt_service._language_directive detected_lang = _detect_text_language(snippets_text) @@ -131,10 +119,9 @@ def ensure_doc_nav_json( if os.path.exists(nav_path) and not overwrite: return nav_path - # Re-use ZipResultService's builder to keep the format canonical - from shared.services.storage.zip_result_service import ZipResultService - svc = ZipResultService() - doc_nav = svc._build_doc_nav(chunks, source_file_name) + from shared.services.storage.zip_result_schema import ZipResultSchemaBuilder + + doc_nav = ZipResultSchemaBuilder().build_doc_nav(chunks, source_file_name) _save_doc_nav(file_dir, doc_nav) return nav_path @@ -276,7 +263,7 @@ def _build_nav_top_summary( def enrich_doc_nav_summaries( - kb_dir: str, + document_workspace_dir: str, source_file: Optional[str] = None, force: bool = False, use_llm: bool = True, @@ -284,7 +271,7 @@ def enrich_doc_nav_summaries( """Enrich doc_nav.json with bottom-up recursive summaries. Args: - kb_dir: Absolute path to the KB directory. + document_workspace_dir: Absolute path to the temporary document workspace. source_file: If given, only process this file. Otherwise process all. force: If True, regenerate even if summaries already exist. use_llm: If True, use LLM for multi-child aggregation. @@ -300,13 +287,13 @@ def enrich_doc_nav_summaries( else: targets = [ entry - for entry in os.listdir(kb_dir) - if os.path.isdir(os.path.join(kb_dir, entry)) + for entry in os.listdir(document_workspace_dir) + if os.path.isdir(os.path.join(document_workspace_dir, entry)) and not entry.startswith(".") ] for file_name in targets: - file_dir = os.path.join(kb_dir, file_name) + file_dir = os.path.join(document_workspace_dir, file_name) doc_nav = _load_doc_nav(file_dir) if doc_nav is None: logger.debug(f"No {DOC_NAV_FILENAME} for {file_name}, skipping") @@ -347,7 +334,7 @@ def build_section_summary_lookup(file_dir: str) -> Dict[str, str]: """Build a flat {section_path: summary} dict from all nodes in doc_nav.json. Keys use the DocumentSection.section_path format produced by - ``section_path_from_chunk_path`` (strips kb_root + filename prefix, + ``section_path_from_chunk_path`` (strips the filename prefix, joins remaining parts with ``" / "``). Traverses the full section tree at all depths. Used by the publication @@ -355,24 +342,28 @@ def build_section_summary_lookup(file_dir: str) -> Dict[str, str]: Args: file_dir: Absolute path to the file-level directory - (e.g. ~/.knowhere/{kb_id}/{source_file_name}/). + inside the task-scoped parse workspace. Returns: Dict mapping section_path → summary string (empty dict on any error). """ - from shared.services.retrieval.lexical_text import section_path_from_chunk_path + from shared.services.retrieval.search.lexical_text import section_path_from_chunk_path doc_nav = _load_doc_nav(file_dir) if doc_nav is None: return {} lookup: Dict[str, str] = {} + source_file_name = str(doc_nav.get("file_name") or "") def _walk(node: Dict[str, Any]) -> None: nav_path = node.get("path", "") summary = node.get("summary", "") if nav_path and summary: - section_path = section_path_from_chunk_path(nav_path) + section_path = section_path_from_chunk_path( + nav_path, + source_file_name=source_file_name, + ) if section_path and section_path != "Root": lookup[section_path] = summary for child in node.get("children", []): diff --git a/apps/worker/app/services/document_agent/tools/classify_special_pages.py b/apps/worker/app/services/document_agent/tools/classify_special_pages.py index 564307b21..0b8efa6c8 100644 --- a/apps/worker/app/services/document_agent/tools/classify_special_pages.py +++ b/apps/worker/app/services/document_agent/tools/classify_special_pages.py @@ -141,7 +141,7 @@ def classify_special_pages( try: from shared.core.config import settings - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client effective_model = model or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL client = get_openai_client(model=effective_model) diff --git a/apps/worker/app/services/document_agent/tools/probe_sample_pages.py b/apps/worker/app/services/document_agent/tools/probe_sample_pages.py index 1ac8018e8..b5b7e1122 100644 --- a/apps/worker/app/services/document_agent/tools/probe_sample_pages.py +++ b/apps/worker/app/services/document_agent/tools/probe_sample_pages.py @@ -6,7 +6,7 @@ import statistics from typing import Any, Literal -from app.services.document_parser.pymupdf_subprocess import run_in_child_process, worker +from app.services.document_parser.formats.pdf.pymupdf_subprocess import run_in_child_process, worker SampleStrategy = Literal["stratified", "uniform", "key_pages"] diff --git a/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py b/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py index 7b4de6396..b073236a5 100644 --- a/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py +++ b/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py @@ -7,7 +7,7 @@ import tempfile from typing import Any -from app.services.document_parser.pymupdf_subprocess import run_in_child_process, worker +from app.services.document_parser.formats.pdf.pymupdf_subprocess import run_in_child_process, worker from loguru import logger from openai.types.chat import ( ChatCompletionContentPartImageParam, @@ -71,7 +71,7 @@ def _call_vlm( max_tokens: int = 900, ) -> tuple[str, dict[str, int]]: from shared.core.config import settings - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client effective_model = model or settings.IMAGE_MODEL or "qwen3.5-flash" client = get_openai_client(model=effective_model) diff --git a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py index 2eb834ff4..1db72592b 100644 --- a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py +++ b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py @@ -295,7 +295,7 @@ def propose_shard_plan( if use_llm and page_count > max_pages_per_shard: try: from shared.core.config import settings - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client effective_model = model or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL client = get_openai_client(model=effective_model) diff --git a/apps/worker/app/services/document_ingestion/__init__.py b/apps/worker/app/services/document_ingestion/__init__.py new file mode 100644 index 000000000..191bf4665 --- /dev/null +++ b/apps/worker/app/services/document_ingestion/__init__.py @@ -0,0 +1 @@ +"""Worker-side Document Ingestion modules.""" diff --git a/apps/worker/app/services/common/job_start_service.py b/apps/worker/app/services/document_ingestion/job_state_gate.py similarity index 78% rename from apps/worker/app/services/common/job_start_service.py rename to apps/worker/app/services/document_ingestion/job_state_gate.py index 150b852c4..4c9b232f3 100644 --- a/apps/worker/app/services/common/job_start_service.py +++ b/apps/worker/app/services/document_ingestion/job_state_gate.py @@ -1,9 +1,9 @@ """ -Worker-side job start gating. +Worker-side Document Ingestion state gate. Keeps the worker-specific policy for when a parse task is allowed to move a -job into ``running`` while delegating the actual state transition to the -shared sync state machine service. +Job into ``running`` while delegating the actual transition to the shared sync +state machine service. """ from typing import Any @@ -24,11 +24,7 @@ def mark_job_running(job_id: str, redis_service: Any) -> bool: - """Transition a job from pending to running before parse execution. - - Returns ``True`` when parsing should proceed. Returns ``False`` when the - job is already terminal and the task should skip quietly. - """ + """Transition a Job from pending to running before parse execution.""" state_machine = SyncStateMachineService(redis_service) with get_sync_db_context() as db: @@ -45,9 +41,6 @@ def mark_job_running(job_id: str, redis_service: Any) -> bool: current_state = job.status if current_state == JobStatus.RUNNING.value: - # Likely a broker redelivery while the original worker is still - # processing. Let the caller proceed to RedisJobLock, which gates - # actual execution. logger.info( f"Job already running (likely redelivery), deferring to lock: {job_id}" ) @@ -59,7 +52,7 @@ def mark_job_running(job_id: str, redis_service: Any) -> bool: f"Parse task started before upload transition completed for job {job_id}; " f"current_state={current_state}" ), - retry_after=settings.KB_TASK_RETRY_COUNTDOWN, + retry_after=settings.DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN, user_message="Job is not ready for processing yet. Retrying shortly.", ) @@ -79,19 +72,20 @@ def mark_job_running(job_id: str, redis_service: Any) -> bool: ), ) - if not state_machine.transition( + outcome = state_machine.transition_outcome( db, job_id, JobStatus.RUNNING.value, "start_processing", operator_type="system", - ): + ) + if not outcome.succeeded: raise UnavailableException( internal_message=( f"Failed to transition job {job_id} from pending to running; " - "the state may have changed concurrently" + f"reason={outcome.reason}" ), - retry_after=settings.KB_TASK_RETRY_COUNTDOWN, + retry_after=settings.DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN, user_message="Job state is still settling. Retrying shortly.", ) diff --git a/apps/worker/app/services/document_ingestion/page_estimator.py b/apps/worker/app/services/document_ingestion/page_estimator.py new file mode 100644 index 000000000..f2ed89cdb --- /dev/null +++ b/apps/worker/app/services/document_ingestion/page_estimator.py @@ -0,0 +1,175 @@ +""" +Page estimator for worker-side Document Ingestion billing. + +Calculates page counts for billing based on: +- PDF: Physical page count from metadata +- PPTX: Slide count +- Text-based (DOC, DOCX, TXT, MD, JSON): Word-based estimation using count_cn_en +- Spreadsheet-based (XLS, XLSX): Row-based estimation +""" + +import math +import tempfile +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path + +from shared.core.logging import logger +from shared.utils.text_utils import count_cn_en + +WORDS_PER_PAGE = 500 +ROWS_PER_PAGE = 50 + + +@dataclass(frozen=True) +class WorkloadEstimate: + page_count: int + method: str + fallback_reason: str | None = None + + @property + def used_fallback(self) -> bool: + return self.fallback_reason is not None + + +class PageEstimator: + """Estimate page count for billing purposes.""" + + @classmethod + def estimate(cls, file_path: str) -> int: + """Estimate the billable page count for a file.""" + return cls.estimate_workload(file_path).page_count + + @classmethod + def estimate_workload(cls, file_path: str) -> WorkloadEstimate: + """Estimate billable workload while keeping fallback policy explicit.""" + path = Path(file_path) + suffix = path.suffix.lower() + + if suffix == ".pdf": + return cls._estimate_with_fallback("pdf_metadata", cls._count_pdf, file_path) + if suffix == ".pptx": + return cls._estimate_with_fallback("pptx_slides", cls._count_pptx, file_path) + if suffix == ".doc": + return cls._estimate_with_fallback("doc_conversion", cls._count_doc, file_path) + if suffix == ".docx": + return cls._estimate_with_fallback("docx_words", cls._count_docx, file_path) + if suffix == ".xls": + return cls._estimate_with_fallback("xls_conversion", cls._count_xls, file_path) + if suffix == ".xlsx": + return cls._estimate_with_fallback("xlsx_rows", cls._count_xlsx, file_path) + if suffix in [".txt", ".md", ".json", ".fragment"]: + return cls._estimate_with_fallback("text_words", cls._count_text, file_path) + if suffix in [".png", ".jpg", ".jpeg"]: + return WorkloadEstimate(page_count=1, method="image_default") + + fallback_reason = f"unknown_file_type:{suffix or ''}" + logger.warning( + f"Unknown file type for billing: {suffix}, defaulting to 1 page" + ) + return WorkloadEstimate( + page_count=1, + method="unknown_file_type", + fallback_reason=fallback_reason, + ) + + @classmethod + def _estimate_with_fallback( + cls, + method: str, + estimator: Callable[[str], int], + file_path: str, + ) -> WorkloadEstimate: + try: + page_count = max(1, estimator(file_path)) + return WorkloadEstimate(page_count=page_count, method=method) + except ImportError as exc: + fallback_reason = f"missing_dependency:{exc.name or exc}" + logger.warning( + f"Missing dependency for {method} page estimation, defaulting to 1 page: {exc}" + ) + except Exception as exc: + fallback_reason = f"{method}_error:{type(exc).__name__}" + logger.error(f"Error estimating pages for {file_path}: {exc}") + + return WorkloadEstimate( + page_count=1, + method=method, + fallback_reason=fallback_reason, + ) + + @classmethod + def _count_pdf(cls, file_path: str) -> int: + """Estimate pages for PDF using physical page count.""" + from pypdf import PdfReader + + reader = PdfReader(file_path) + return len(reader.pages) + + @classmethod + def _count_pptx(cls, file_path: str) -> int: + """Estimate pages for PPTX using slide count.""" + from pptx import Presentation + + presentation = Presentation(file_path) + return len(presentation.slides) + + @classmethod + def _count_docx(cls, file_path: str) -> int: + """Estimate pages for DOCX using word-based counting.""" + from docx import Document + + document = Document(file_path) + total_text = "" + + for paragraph in document.paragraphs: + total_text += paragraph.text + " " + + for table in document.tables: + for row in table.rows: + for cell in row.cells: + total_text += cell.text + " " + + word_count = count_cn_en(total_text) + return math.ceil(word_count / WORDS_PER_PAGE) + + @classmethod + def _count_doc(cls, file_path: str) -> int: + """Estimate pages for DOC by converting it to DOCX first.""" + from app.services.document_parser.conversion.legacy_converter import doc_to_docx + + with tempfile.TemporaryDirectory(prefix="page-estimator-doc-") as temp_dir: + converted_path, _ = doc_to_docx(file_path, temp_dir) + return cls._count_docx(converted_path) + + @classmethod + def _count_xlsx(cls, file_path: str) -> int: + """Estimate pages for XLSX using row count.""" + import pandas as pd + + workbook = pd.ExcelFile(file_path) + total_rows = 0 + + for sheet_name in workbook.sheet_names: + dataframe = pd.read_excel(workbook, sheet_name=sheet_name) + total_rows += len(dataframe) + + return math.ceil(total_rows / ROWS_PER_PAGE) + + @classmethod + def _count_xls(cls, file_path: str) -> int: + """Estimate pages for XLS by converting it to XLSX first.""" + from app.services.document_parser.conversion.legacy_converter import xls_to_xlsx + + with tempfile.TemporaryDirectory(prefix="page-estimator-xls-") as temp_dir: + converted_path, _ = xls_to_xlsx(file_path, temp_dir) + return cls._count_xlsx(converted_path) + + @classmethod + def _count_text(cls, file_path: str) -> int: + """Estimate pages for text-like files using word-based counting.""" + with open(file_path, "r", encoding="utf-8", errors="ignore") as file: + content = file.read() + + word_count = count_cn_en(content) + return math.ceil(word_count / WORDS_PER_PAGE) diff --git a/apps/worker/app/services/document_ingestion/parse_execution.py b/apps/worker/app/services/document_ingestion/parse_execution.py new file mode 100644 index 000000000..594ff60fe --- /dev/null +++ b/apps/worker/app/services/document_ingestion/parse_execution.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from app.services.document_ingestion.processing_context import ParseJobContext +from app.services.document_ingestion.source_preparation import PreparedSourceFile +from app.services.document_parser import parse_service +from app.services.document_parser.orchestration.parse_output import ParseOutput +from app.services.document_parser.support.stage_profiler import stage_timer +from loguru import logger + +from shared.models.schemas.job_metadata import JobMetadataHelper + + +def execute_document_parse( + *, + job_id: str, + job_context: ParseJobContext, + prepared_source: PreparedSourceFile, + output_dir: str, +) -> ParseOutput: + """Run the parser adapter for a prepared source file.""" + doc_type = JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "doc_type", + "auto", + ) + logger.info( + f"Start parse: job_id={job_id}, " + f"filename={prepared_source.source_file_name}, " + f"internal_filename={prepared_source.internal_parse_name}, type={doc_type}" + ) + + with stage_timer( + "worker.parse.document", + job_id=job_id, + filename=prepared_source.source_file_name, + doc_type=doc_type, + ): + parse_output = parse_service.checkerboard_parse_output( + file_full_path=prepared_source.local_file_path, + filename=prepared_source.source_file_name, + output_dir=output_dir, + job_id=job_id, + internal_output_filename=prepared_source.internal_parse_name, + doc_type=doc_type, + smart_title_parse=JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "smart_title_parse", + True, + ), + summary_image=JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "summary_image", + True, + ), + summary_table=JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "summary_table", + True, + ), + summary_txt=JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "summary_txt", + True, + ), + add_frag_desc=JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "add_frag_desc", + "", + ), + s3_key=job_context.s3_key, + ) + + logger.info( + "File parsing completed: " + f"job_id={job_id}, output_dir={parse_output.output_dir}, " + f"chunks={parse_output.rows_count}" + ) + return parse_output diff --git a/apps/worker/app/services/document_ingestion/parse_result_package.py b/apps/worker/app/services/document_ingestion/parse_result_package.py new file mode 100644 index 000000000..c26fdf186 --- /dev/null +++ b/apps/worker/app/services/document_ingestion/parse_result_package.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import pandas as pd +from app.services.document_parser.orchestration.parse_output import ParseOutput +from loguru import logger + +from shared.core.exceptions.domain_exceptions import WorkerHandlingException +from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks + + +@dataclass(frozen=True) +class ParseArtifact: + add_dir: str + dataframe: pd.DataFrame + + @property + def contents_count(self) -> int: + return len(self.dataframe) + + +@dataclass(frozen=True) +class ParseResultPackage: + artifact: ParseArtifact + chunks: list[dict[str, Any]] + + +@dataclass(frozen=True) +class GeneratedResultPackage: + zip_file_path: str + checksum_value: str + statistics: dict[str, Any] + zip_size: int + + +def build_parse_result_package( + *, + job_id: str, + filename: str, + parse_output: ParseOutput, +) -> ParseResultPackage: + artifact = _build_parse_artifact( + job_id=job_id, + filename=filename, + parse_output=parse_output, + ) + chunks = dataframe_to_chunks(artifact.dataframe) + return ParseResultPackage(artifact=artifact, chunks=chunks) + + +def build_generated_result_package( + zip_file_path: str, + checksum: dict[str, Any] | str | None, + statistics: dict[str, Any], + zip_size: int, +) -> GeneratedResultPackage: + checksum_value = ( + str(checksum.get("value", "")) + if isinstance(checksum, dict) + else str(checksum or "") + ) + return GeneratedResultPackage( + zip_file_path=zip_file_path, + checksum_value=checksum_value, + statistics=statistics, + zip_size=zip_size, + ) + + +def _build_parse_artifact( + *, + job_id: str, + filename: str, + parse_output: ParseOutput, +) -> ParseArtifact: + parsed_contents_df = parse_output.parsed_df + if parsed_contents_df is None: + raise WorkerHandlingException( + user_message="We could not extract content from your file", + internal_message="File parsing failed, no content returned from parser", + ) + + if parsed_contents_df.empty: + logger.warning( + f"No content returned from file parsing: job_id={job_id}, filename={filename}" + ) + + return ParseArtifact(add_dir=parse_output.output_dir, dataframe=parsed_contents_df) diff --git a/apps/worker/app/services/document_ingestion/processing_billing.py b/apps/worker/app/services/document_ingestion/processing_billing.py new file mode 100644 index 000000000..f0d313fa4 --- /dev/null +++ b/apps/worker/app/services/document_ingestion/processing_billing.py @@ -0,0 +1,125 @@ +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime + +from app.services.document_ingestion.page_estimator import WorkloadEstimate +from app.services.document_ingestion.processing_context import ParseJobContext +from loguru import logger +from sqlalchemy import select + +from shared.core.database_sync import get_sync_db_context +from shared.core.exceptions.domain_exceptions import ( + InsufficientCreditsException, + NotFoundException, +) +from shared.models.database.job import Job +from shared.services.billing.work_billing_service import WorkBillingService + + +@dataclass(frozen=True) +class ParseJobBillingSnapshot: + billing_amount_micro_dollars: int + billing_credits: float + billing_status: str + + +def charge_parse_job_pages( + *, + job_id: str, + filename: str | None, + job_user_id: str | None, + workload_estimate: WorkloadEstimate, +) -> ParseJobBillingSnapshot: + page_count = workload_estimate.page_count + if not job_user_id: + raise NotFoundException( + resource="JobInfo", + resource_id="user_id", + internal_message=f"Missing user_id in job info for job_id={job_id}", + ) + + billing_service = WorkBillingService() + billing_filename = filename or "" + billing_status: str + billing_amount_micro_dollars: int + billing_credits: float + + with get_sync_db_context() as db: + job_result = db.execute(select(Job).where(Job.job_id == job_id).with_for_update()) + job = job_result.scalar_one_or_none() + + if job and getattr(job, "billing_status", "") == "charged": + logger.info(f"Job already charged: {job_id}") + billing_status = "charged" + billing_amount_micro_dollars = int(job.credits_charged or 0) + billing_credits = billing_amount_micro_dollars / 1_000_000 + else: + try: + billing_result = billing_service.charge_for_pages( + session=db, + user_id=job_user_id, + page_count=page_count, + filename=billing_filename, + ) + except InsufficientCreditsException: + logger.warning(f"Billing failed: job_id={job_id}, user_id={job_user_id}") + billing_amount = billing_service.estimate_page_charge( + page_count=page_count + ) + if job: + job.page_count = page_count + job.credits_charged = billing_amount.amount_micro_dollars + job.billing_status = "billing_failed" + db.commit() + + raise InsufficientCreditsException( + user_message=( + "Insufficient credits to process this document " + f"({page_count} pages required, cost: " + f"{billing_amount.credits})." + ), + required_credits=billing_amount.credits, + internal_message=( + f"job_id={job_id}, user_id={job_user_id}, " + f"page_count={page_count}" + ), + ) + + billing_status = billing_result.billing_status + billing_amount_micro_dollars = billing_result.amount_micro_dollars + billing_credits = billing_result.credits + if job: + job.page_count = page_count + job.credits_charged = billing_amount_micro_dollars + job.billing_status = billing_status + + return ParseJobBillingSnapshot( + billing_amount_micro_dollars=billing_amount_micro_dollars, + billing_credits=billing_credits, + billing_status=billing_status, + ) + + +def record_processing_start( + *, + job_id: str, + job_context: ParseJobContext, + billing_snapshot: ParseJobBillingSnapshot, + processing_started_at: datetime, + workload_estimate: WorkloadEstimate, +) -> None: + metadata_updates = { + "page_count": workload_estimate.page_count, + "billing_status": billing_snapshot.billing_status, + "billing_amount_micro_dollars": billing_snapshot.billing_amount_micro_dollars, + "billing_credits": billing_snapshot.billing_credits, + "processing_started_at": processing_started_at.isoformat(), + "workload_estimate_method": workload_estimate.method, + } + if workload_estimate.fallback_reason is not None: + metadata_updates["workload_estimate_fallback_reason"] = ( + workload_estimate.fallback_reason + ) + job_context.metadata_service.update_metadata(job_id, metadata_updates) + job_context.job_metadata.update(metadata_updates) diff --git a/apps/worker/app/services/document_ingestion/processing_context.py b/apps/worker/app/services/document_ingestion/processing_context.py new file mode 100644 index 000000000..51a40cd43 --- /dev/null +++ b/apps/worker/app/services/document_ingestion/processing_context.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Any + +from loguru import logger +from sqlalchemy import select + +from shared.core.config import settings +from shared.core.database_sync import get_sync_db_context +from shared.core.exceptions.domain_exceptions import ( + NotFoundException, + ValidationException, +) +from shared.models.database.job import Job +from shared.services.redis.redis_sync_service import ( + SyncJobInfoRedisService, + SyncJobMetadataService, +) +from shared.services.storage.job_file_storage import JobFileStorage + + +@dataclass(frozen=True) +class ParseJobContext: + job_metadata: dict[str, object] + job_user_id: str | None + metadata_service: SyncJobMetadataService + redis_service: Any + s3_key: str + + +def load_parse_job_context( + job_id: str, + requested_user_id: str | None, + redis_service: Any, +) -> ParseJobContext: + job_info_service = SyncJobInfoRedisService(redis_service) + job_info = job_info_service.get_job_info(job_id) + + if not job_info: + logger.warning( + f"JobInfo not found in Redis for job_id={job_id}; falling back to database" + ) + with get_sync_db_context() as fallback_db: + job_row = fallback_db.execute( + select(Job).where(Job.job_id == job_id) + ).scalar_one_or_none() + + if not job_row or not job_row.s3_key: + raise NotFoundException( + resource="JobInfo", + resource_id=job_id, + internal_message="job info not found in Redis or database", + ) + + s3_key: str = job_row.s3_key + job_user_id: str | None = ( + str(job_row.user_id) if job_row.user_id else requested_user_id + ) + logger.info(f"Recovered JobInfo from database: job_id={job_id}, s3_key={s3_key}") + else: + raw_s3_key = job_info.get("s3_key") + if not isinstance(raw_s3_key, str) or not raw_s3_key: + raise NotFoundException( + resource="JobInfo", + resource_id="s3_key", + internal_message="Missing s3_key in job_info", + ) + + s3_key = raw_s3_key + raw_job_user_id = job_info.get("user_id") + job_user_id = ( + raw_job_user_id if isinstance(raw_job_user_id, str) else requested_user_id + ) + + metadata_service = SyncJobMetadataService(redis_service) + raw_job_metadata = metadata_service.get_metadata(job_id) + if not isinstance(raw_job_metadata, dict) or not raw_job_metadata: + raise NotFoundException( + resource="JobMetadata", + resource_id=job_id, + internal_message=f"Job metadata not found for job_id={job_id}", + ) + + return ParseJobContext( + job_metadata=dict(raw_job_metadata), + job_user_id=job_user_id, + metadata_service=metadata_service, + redis_service=redis_service, + s3_key=s3_key, + ) + + +def assert_source_file_within_size_limit(s3_key: str) -> None: + file_info = JobFileStorage().verify_upload_exists(s3_key) + if not file_info.get("exists"): + raise NotFoundException( + resource="S3File", + resource_id=s3_key, + internal_message=f"S3 file not found: {s3_key}", + ) + + logger.info(f"S3 file verified: {s3_key}") + + file_size = file_info.get("size", 0) + file_extension = os.path.splitext(s3_key)[1].lower() + if file_size > settings.MAX_FILE_SIZE: + limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024) + raise ValidationException( + user_message=f"File size exceeds limit (max {limit_mb}MB for {file_extension})", + violations=[ + { + "field": "file_size", + "description": ( + f"Size {file_size} bytes exceeds limit of " + f"{settings.MAX_FILE_SIZE} bytes" + ), + } + ], + ) diff --git a/apps/worker/app/services/document_ingestion/processing_run.py b/apps/worker/app/services/document_ingestion/processing_run.py new file mode 100644 index 000000000..995990379 --- /dev/null +++ b/apps/worker/app/services/document_ingestion/processing_run.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any + +from app.services.document_ingestion.job_state_gate import mark_job_running +from app.services.document_ingestion.page_estimator import PageEstimator +from app.services.document_ingestion.parse_result_package import ( + build_parse_result_package, +) +from app.services.document_ingestion.parse_execution import execute_document_parse +from app.services.document_ingestion.processing_billing import ( + charge_parse_job_pages, + record_processing_start, +) +from app.services.document_ingestion.processing_context import ( + ParseJobContext, + assert_source_file_within_size_limit, + load_parse_job_context, +) +from app.services.document_ingestion.source_preparation import prepare_source_file +from app.services.document_ingestion.success_finalization import finalize_parse_success +from app.services.document_ingestion.workspace import ( + TemporaryParseWorkspace, + cleanup_task_workspace, + download_s3_file_to_temp, +) +from loguru import logger + +from shared.services.jobs.lifecycle.service import get_sync_job_lifecycle_service +from shared.services.redis.distributed_lock import RedisJobLock +from shared.services.redis.redis_sync_service import ( + SyncRedisServiceFactory, +) +from shared.services.storage.result_storage import get_result_storage + + +class DocumentProcessingRun: + """Run worker-side Document Ingestion for an uploaded file Job.""" + + def execute(self, job_id: str, user_id: str | None) -> dict[str, object]: + logger.info(f"Parse started: job_id={job_id}, user_id={user_id}") + lifecycle_service = get_sync_job_lifecycle_service() + + redis_service = SyncRedisServiceFactory.get_service() + job_context = load_parse_job_context(job_id, user_id, redis_service) + assert_source_file_within_size_limit(job_context.s3_key) + + should_process = mark_job_running(job_id, job_context.redis_service) + if not should_process: + logger.warning(f"Skipping parse_task for inactive job: job_id={job_id}") + return { + "status": "skipped", + "job_id": job_id, + "reason": "job_already_terminal", + } + + with RedisJobLock(job_context.redis_service, job_id): + task_workspace = TemporaryParseWorkspace.create(job_id) + try: + result = _run_parse_job( + job_id=job_id, + job_context=job_context, + lifecycle_service=lifecycle_service, + task_workspace=task_workspace, + ) + finally: + task_workspace.cleanup(cleanup_task_workspace) + + return result + + +def _run_parse_job( + *, + job_id: str, + job_context: ParseJobContext, + lifecycle_service: Any, + task_workspace: TemporaryParseWorkspace, +) -> dict[str, object]: + lifecycle_service.update_progress(job_id, progress=10, message="Parsing document...") + + prepared_source = prepare_source_file( + job_id=job_id, + job_context=job_context, + input_dir=task_workspace.input_dir, + download_source_file=download_s3_file_to_temp, + ) + + workload_estimate = PageEstimator.estimate_workload(prepared_source.local_file_path) + page_count = workload_estimate.page_count + logger.info( + "Workload estimation: " + f"job_id={job_id}, page_count={page_count}, " + f"method={workload_estimate.method}, " + f"fallback_reason={workload_estimate.fallback_reason}" + ) + + processing_started_at = datetime.now(timezone.utc) + billing_snapshot = charge_parse_job_pages( + job_id=job_id, + filename=prepared_source.source_file_name, + job_user_id=job_context.job_user_id, + workload_estimate=workload_estimate, + ) + record_processing_start( + job_id=job_id, + job_context=job_context, + billing_snapshot=billing_snapshot, + processing_started_at=processing_started_at, + workload_estimate=workload_estimate, + ) + + parse_output = execute_document_parse( + job_id=job_id, + job_context=job_context, + prepared_source=prepared_source, + output_dir=task_workspace.output_dir, + ) + + lifecycle_service.update_progress( + job_id, + progress=30, + message="Parse completed, preparing chunks...", + ) + result_package = build_parse_result_package( + job_id=job_id, + filename=prepared_source.source_file_name, + parse_output=parse_output, + ) + + lifecycle_service.update_progress( + job_id, + progress=70, + message="Chunks ready, generating zip...", + ) + logger.info( + f"Chunks prepared: job_id={job_id}, count={len(result_package.chunks)}" + ) + + return finalize_parse_success( + result_package=result_package, + job_context=job_context, + job_id=job_id, + lifecycle_service=lifecycle_service, + processing_started_at=processing_started_at, + task_workspace_dir=task_workspace.root_dir, + result_storage_factory=get_result_storage, + ) diff --git a/apps/worker/app/services/document_ingestion/service.py b/apps/worker/app/services/document_ingestion/service.py new file mode 100644 index 000000000..292bb155f --- /dev/null +++ b/apps/worker/app/services/document_ingestion/service.py @@ -0,0 +1,10 @@ +from __future__ import annotations + +from app.services.document_ingestion.processing_run import DocumentProcessingRun + +__all__ = ["parse_uploaded_file_job"] + + +def parse_uploaded_file_job(job_id: str, user_id: str | None) -> dict[str, object]: + """Run worker-side Document Ingestion for an uploaded file Job.""" + return DocumentProcessingRun().execute(job_id, user_id) diff --git a/apps/worker/app/services/document_ingestion/source_preparation.py b/apps/worker/app/services/document_ingestion/source_preparation.py new file mode 100644 index 000000000..cc3c1bd79 --- /dev/null +++ b/apps/worker/app/services/document_ingestion/source_preparation.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +import os +from collections.abc import Callable +from dataclasses import dataclass + +from app.services.document_ingestion.processing_context import ParseJobContext +from app.services.document_ingestion.workspace import download_s3_file_to_temp +from app.services.document_parser.support.internal_parse_name import ( + prepare_internal_parse_input, +) +from loguru import logger + +from shared.models.schemas.job_metadata import JobMetadataHelper + +DownloadSourceFile = Callable[[str, str, str], str] + + +@dataclass(frozen=True) +class PreparedSourceFile: + source_file_name: str + internal_parse_name: str + local_file_path: str + file_extension: str + + +def prepare_source_file( + *, + job_id: str, + job_context: ParseJobContext, + input_dir: str, + download_source_file: DownloadSourceFile = download_s3_file_to_temp, +) -> PreparedSourceFile: + """Download and normalize the source file before parser execution.""" + source_file_name = JobMetadataHelper.get_source_file_name( + job_context.job_metadata, + ) or os.path.basename(job_context.s3_key) + file_extension = ( + os.path.splitext(job_context.s3_key)[1].lower() if job_context.s3_key else "" + ) + + local_file_path = download_source_file( + job_context.s3_key, + file_extension, + input_dir, + ) + logger.info(f"File downloaded: job_id={job_id}, local_path={local_file_path}") + + prepared_parse_input = prepare_internal_parse_input( + local_file_path, + source_file_name, + fallback_ext=file_extension, + prefer_fallback_ext=True, + ) + logger.info( + f"File prepared for parsing: job_id={job_id}, " + f"internal_filename={prepared_parse_input.internal_filename}, " + f"local_path={prepared_parse_input.file_path}" + ) + + return PreparedSourceFile( + source_file_name=source_file_name, + internal_parse_name=prepared_parse_input.internal_filename, + local_file_path=prepared_parse_input.file_path, + file_extension=file_extension, + ) diff --git a/apps/worker/app/services/document_ingestion/success_finalization.py b/apps/worker/app/services/document_ingestion/success_finalization.py new file mode 100644 index 000000000..62397d56d --- /dev/null +++ b/apps/worker/app/services/document_ingestion/success_finalization.py @@ -0,0 +1,236 @@ +from __future__ import annotations + +import os +from collections.abc import Callable +from datetime import datetime, timezone +from typing import Any + +from app.services.connect_builder.summary_builder import ( + build_section_summary_lookup, + enrich_doc_nav_summaries, + ensure_doc_nav_json, + load_nav_top_summary, +) +from app.services.document_ingestion.parse_result_package import ( + GeneratedResultPackage, + ParseArtifact, + ParseResultPackage, + build_generated_result_package, +) +from app.services.document_ingestion.processing_context import ParseJobContext +from loguru import logger + +from shared.models.schemas.job_metadata import JobMetadataHelper +from shared.services.storage.result_storage import ResultStorage, get_result_storage +from shared.services.storage.zip_result_service import ZipResultService + +ResultStorageFactory = Callable[[], ResultStorage] + + +def finalize_parse_success( + *, + result_package: ParseResultPackage, + job_context: ParseJobContext, + job_id: str, + lifecycle_service: Any, + processing_started_at: datetime, + task_workspace_dir: str, + result_storage_factory: ResultStorageFactory = get_result_storage, +) -> dict[str, object]: + """Package, upload, and publish a successful parser result.""" + source_file_name = _resolve_source_file_name(job_context) + document_top_summary, section_summaries = _enrich_document_navigation( + artifact=result_package.artifact, + chunks=result_package.chunks, + job_context=job_context, + source_file_name=source_file_name, + ) + _attach_document_top_summary(result_package.chunks, document_top_summary) + + lifecycle_service.update_progress( + job_id, + progress=80, + message="Generating ZIP package...", + ) + _record_processing_completion( + job_id=job_id, + job_context=job_context, + processing_started_at=processing_started_at, + ) + generated_package = _generate_result_package( + result_package=result_package, + job_context=job_context, + job_id=job_id, + source_file_name=source_file_name, + task_workspace_dir=task_workspace_dir, + ) + + lifecycle_service.update_progress( + job_id, + progress=90, + message="Uploading results to S3...", + ) + result_s3_key = _upload_result_package( + result_package=result_package, + generated_package=generated_package, + job_id=job_id, + result_storage_factory=result_storage_factory, + ) + stored_count = 0 + + finalization_response = lifecycle_service.finalize_job_success( + job_id=job_id, + chunks=result_package.chunks, + result_s3_key=result_s3_key, + checksum=generated_package.checksum_value, + zip_size=generated_package.zip_size, + stored_count=stored_count, + delivery_mode="url", + section_summaries=section_summaries, + ) + if finalization_response.get("status") != "success": + logger.error( + f"Worker processing finalization failed: job_id={job_id}, " + f"response={finalization_response}" + ) + return dict(finalization_response) + + lifecycle_service.update_progress(job_id, progress=100, message="Task complete!") + logger.info( + f"Worker processing complete: job_id={job_id}, result_s3_key={result_s3_key}" + ) + + return { + "status": "success", + "job_id": job_id, + "add_dir": None, + "vectors_count": 0, + "contents_count": result_package.artifact.contents_count, + "stored_count": stored_count, + "delivery_mode": "url", + "result_s3_key": result_s3_key, + } + + +def _resolve_source_file_name(job_context: ParseJobContext) -> str: + source_file_name = JobMetadataHelper.get_source_file_name( + job_context.job_metadata, + ) or JobMetadataHelper.get_source_url(job_context.job_metadata) + if not source_file_name: + source_file_name = os.path.basename(job_context.s3_key) + if isinstance(source_file_name, str) and "/" in source_file_name: + source_file_name = os.path.basename(source_file_name) + return str(source_file_name) + + +def _enrich_document_navigation( + *, + artifact: ParseArtifact, + chunks: list[dict[str, Any]], + job_context: ParseJobContext, + source_file_name: str, +) -> tuple[str, dict[str, str]]: + document_top_summary = "" + section_summaries: dict[str, str] = {} + add_dir = artifact.add_dir + parsed_contents_df = artifact.dataframe + if add_dir and source_file_name: + if "path" in parsed_contents_df.columns: + ensure_doc_nav_json( + str(add_dir), + chunks, + source_file_name=source_file_name, + ) + try: + document_root_for_enrich = os.path.dirname(str(add_dir)) + summary_use_llm = JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "summary_use_llm", + False, + ) + enrich_doc_nav_summaries( + document_root_for_enrich, + source_file=source_file_name, + use_llm=summary_use_llm, + ) + section_summaries = build_section_summary_lookup(str(add_dir)) + except Exception as exc: + logger.warning(f"doc_nav enrichment failed (non-fatal): {exc}") + document_top_summary = load_nav_top_summary(str(add_dir), source_file_name) + return document_top_summary, section_summaries + + +def _attach_document_top_summary( + chunks: list[dict[str, Any]], + document_top_summary: str, +) -> None: + if not document_top_summary: + return + + for chunk in chunks: + metadata = chunk.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + chunk["metadata"] = metadata + metadata["document_top_summary"] = document_top_summary + + +def _record_processing_completion( + *, + job_id: str, + job_context: ParseJobContext, + processing_started_at: datetime, +) -> None: + processing_completed_at = datetime.now(timezone.utc) + processing_timing_updates = { + "processing_completed_at": processing_completed_at.isoformat(), + "processing_duration_ms": max( + 0, + int((processing_completed_at - processing_started_at).total_seconds() * 1000), + ), + } + job_context.metadata_service.update_metadata(job_id, processing_timing_updates) + job_context.job_metadata.update(processing_timing_updates) + + +def _generate_result_package( + *, + result_package: ParseResultPackage, + job_context: ParseJobContext, + job_id: str, + source_file_name: str, + task_workspace_dir: str, +) -> GeneratedResultPackage: + data_id = JobMetadataHelper.get_data_id(job_context.job_metadata) + zip_service = ZipResultService() + return build_generated_result_package( + *zip_service.generate_zip_package( + job_id=job_id, + chunks=result_package.chunks, + add_dir=str(result_package.artifact.add_dir) + if result_package.artifact.add_dir + else "", + source_file_name=source_file_name, + data_id=data_id, + job_metadata=job_context.job_metadata, + parsed_df=result_package.artifact.dataframe, + temp_dir=task_workspace_dir, + ) + ) + + +def _upload_result_package( + *, + result_package: ParseResultPackage, + generated_package: GeneratedResultPackage, + job_id: str, + result_storage_factory: ResultStorageFactory, +) -> str: + result_bundle = result_storage_factory().upload( + job_id=job_id, + result_dir=str(result_package.artifact.add_dir) + if result_package.artifact.add_dir + else "", + zip_file_path=generated_package.zip_file_path, + ) + return result_bundle.zip_key diff --git a/apps/worker/app/core/tasks/task_utils.py b/apps/worker/app/services/document_ingestion/workspace.py similarity index 54% rename from apps/worker/app/core/tasks/task_utils.py rename to apps/worker/app/services/document_ingestion/workspace.py index dafe3dd0e..e8a0188dc 100644 --- a/apps/worker/app/core/tasks/task_utils.py +++ b/apps/worker/app/services/document_ingestion/workspace.py @@ -1,17 +1,48 @@ +"""Task-scoped workspace helpers for worker-side Document Ingestion.""" + import os import shutil import tempfile +from collections.abc import Callable +from dataclasses import dataclass -import requests from loguru import logger from shared.core.config import settings from shared.core.exceptions.domain_exceptions import ( FileSystemException, - StorageServiceException, SystemSettingInvalidException, SystemSettingMissingException, ) +from shared.services.storage.job_file_storage import JobFileStorage + +CleanupTaskWorkspace = Callable[[str | None], bool] + + +@dataclass(frozen=True) +class TemporaryParseWorkspace: + """Task-local folders for parser input, parser output, and ZIP generation.""" + + root_dir: str + input_dir: str + output_dir: str + + @classmethod + def create(cls, job_id: str) -> "TemporaryParseWorkspace": + root_dir = create_task_workspace(job_id) + input_dir = os.path.join(root_dir, "input") + output_dir = os.path.join(root_dir, "output") + os.makedirs(input_dir, exist_ok=True) + os.makedirs(output_dir, exist_ok=True) + logger.info(f"Task workspace ready: job_id={job_id}, workspace={root_dir}") + return cls(root_dir=root_dir, input_dir=input_dir, output_dir=output_dir) + + def cleanup( + self, + cleanup_workspace: CleanupTaskWorkspace | None = None, + ) -> bool: + resolved_cleanup = cleanup_workspace or cleanup_task_workspace + return resolved_cleanup(self.root_dir) def cleanup_temp_file(file_path: str | None) -> None: @@ -28,10 +59,7 @@ def cleanup_temp_file(file_path: str | None) -> None: def cleanup_task_workspace(workspace_dir: str | None) -> bool: """Best-effort cleanup for a task-scoped temporary workspace.""" - if not workspace_dir: - return False - - if not os.path.isdir(workspace_dir): + if not workspace_dir or not os.path.isdir(workspace_dir): return False try: @@ -60,7 +88,10 @@ def create_task_workspace(job_id: str) -> str: try: os.makedirs(temp_root, exist_ok=True) - return tempfile.mkdtemp(prefix=f"kb_task_{job_id}_", dir=temp_root) + return tempfile.mkdtemp( + prefix=f"document_ingestion_task_{job_id}_", + dir=temp_root, + ) except (OSError, PermissionError) as exc: raise FileSystemException( user_message="System error preparing temporary storage", @@ -70,32 +101,11 @@ def create_task_workspace(job_id: str) -> str: ) from exc -def download_s3_file_to_temp(file_url: str, file_ext: str, temp_dir: str) -> str: - """Download the source file from object storage into a task workspace file.""" - local_temp_path = None - - try: - os.makedirs(temp_dir, exist_ok=True) - with tempfile.NamedTemporaryFile( - delete=False, suffix=file_ext, dir=temp_dir - ) as tmp_file: - local_temp_path = tmp_file.name - with requests.get( - file_url, - timeout=120, - stream=True, - headers={"User-Agent": "Knowhere-Worker/1.0"}, - ) as response: - response.raise_for_status() - for chunk in response.iter_content(chunk_size=65536): - if chunk: - tmp_file.write(chunk) - except requests.RequestException as exc: - cleanup_temp_file(local_temp_path) - raise StorageServiceException( - internal_message=f"Failed to download source file from object storage: {exc}", - operation="download_source_file", - original_exception=exc, - ) from exc - - return local_temp_path +def download_s3_file_to_temp(s3_key: str, file_ext: str, temp_dir: str) -> str: + """Download the source file from object storage into the task workspace.""" + storage = JobFileStorage() + return storage.download_upload_to_temp( + s3_key, + suffix=file_ext, + temp_dir=temp_dir, + ) diff --git a/apps/worker/app/services/document_parser/__init__.py b/apps/worker/app/services/document_parser/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/apps/worker/app/services/document_parser/image_compressor.py b/apps/worker/app/services/document_parser/assets/image_compressor.py similarity index 100% rename from apps/worker/app/services/document_parser/image_compressor.py rename to apps/worker/app/services/document_parser/assets/image_compressor.py diff --git a/apps/worker/app/services/document_parser/assets/inline_asset.py b/apps/worker/app/services/document_parser/assets/inline_asset.py new file mode 100644 index 000000000..9e792ee1d --- /dev/null +++ b/apps/worker/app/services/document_parser/assets/inline_asset.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from app.services.document_parser.support.parser_rows import ParsedRow + + +def build_image_asset_row( + *, + content: str, + relative_path: str, + summary: str, + know_id: str, + addtime: str, + page_nums: str = "", +) -> ParsedRow: + return ParsedRow( + content=content, + path=relative_path, + type="image", + keywords="", + summary=summary, + know_id=know_id, + tokens="", + connectto="", + addtime=addtime, + page_nums=page_nums, + ) + + +def build_table_asset_row( + *, + content: str, + relative_path: str, + summary: str, + keywords: str, + know_id: str, + addtime: str, + page_nums: str = "", +) -> ParsedRow: + return ParsedRow( + content=content, + path=relative_path, + type="table", + keywords=keywords, + summary=summary, + know_id=know_id, + tokens="", + connectto="", + addtime=addtime, + page_nums=page_nums, + ) diff --git a/apps/worker/app/services/document_parser/legacy_converter.py b/apps/worker/app/services/document_parser/conversion/legacy_converter.py similarity index 97% rename from apps/worker/app/services/document_parser/legacy_converter.py rename to apps/worker/app/services/document_parser/conversion/legacy_converter.py index f4b18fdc2..a602d67b8 100644 --- a/apps/worker/app/services/document_parser/legacy_converter.py +++ b/apps/worker/app/services/document_parser/conversion/legacy_converter.py @@ -4,7 +4,7 @@ import tempfile from pathlib import Path -from app.services.document_parser.parser_log_utils import truncate_log_value +from app.services.document_parser.support.parser_log_utils import truncate_log_value from shared.core.exceptions.domain_exceptions import LibreOfficeServiceException diff --git a/apps/worker/app/services/document_parser/atlas_classifier.py b/apps/worker/app/services/document_parser/formats/atlas/classifier.py similarity index 97% rename from apps/worker/app/services/document_parser/atlas_classifier.py rename to apps/worker/app/services/document_parser/formats/atlas/classifier.py index 1ee8dc4ff..87c43e862 100644 --- a/apps/worker/app/services/document_parser/atlas_classifier.py +++ b/apps/worker/app/services/document_parser/formats/atlas/classifier.py @@ -17,7 +17,7 @@ import tempfile from typing import Optional -from app.services.document_parser.pymupdf_subprocess import run_in_child_process, worker +from app.services.document_parser.formats.pdf.pymupdf_subprocess import run_in_child_process, worker from loguru import logger from openai.types.chat import ( ChatCompletionContentPartImageParam, @@ -102,7 +102,7 @@ def _png_to_data_url(path: str) -> Optional[str]: def _call_vlm(image_data_urls: list[str]) -> bool: """Call VLM with preview images. Returns True if atlas, False otherwise.""" from shared.core.config import settings - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client model = settings.IMAGE_MODEL or "qwen-vl-plus" client = get_openai_client(model=model) diff --git a/apps/worker/app/services/document_parser/atlas_parser.py b/apps/worker/app/services/document_parser/formats/atlas/parser.py similarity index 93% rename from apps/worker/app/services/document_parser/atlas_parser.py rename to apps/worker/app/services/document_parser/formats/atlas/parser.py index a379d9c51..a3210d67e 100644 --- a/apps/worker/app/services/document_parser/atlas_parser.py +++ b/apps/worker/app/services/document_parser/formats/atlas/parser.py @@ -19,13 +19,11 @@ from concurrent.futures import ThreadPoolExecutor, as_completed import pandas as pd -from app.services.common.kb_utils import ( - gen_str_codes, - get_str_time, - process_dup_paths_df, -) -from app.services.document_parser.pymupdf_subprocess import run_in_child_process, worker -from app.services.document_parser.toc_parser import detect_tocs_in_texts +from app.services.document_parser.tables.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.support.parser_rows import ParsedRow, ParsedRowsBuilder +from app.services.document_parser.formats.pdf.pymupdf_subprocess import run_in_child_process, worker +from app.services.document_parser.structure.toc_parser import detect_tocs_in_texts from loguru import logger from shared.core.config import settings @@ -119,7 +117,7 @@ def _vlm_extract_page_info(output_dir: str, img_name: str) -> str: Returns extracted info string, or empty string on failure. """ - from app.services.document_parser.image_parser import _get_vision_client, ask_image + from app.services.document_parser.formats.image.parser import _get_vision_client, ask_image img_path = os.path.join(output_dir, "images", img_name) if not os.path.exists(img_path): @@ -145,8 +143,8 @@ def _vlm_extract_page_info(output_dir: str, img_name: str) -> str: if vlm_path != img_path: try: os.remove(vlm_path) - except OSError: - pass + except OSError as exc: + logger.debug(f"Failed to remove compressed atlas image {vlm_path}: {exc}") if result: return result.strip() @@ -376,7 +374,7 @@ def _vlm_task(page_num, img_name): ) # ── Phase 3: Build chunks ── - df_list = [] + rows_builder = ParsedRowsBuilder() custom_md_lines = [] skipped_null = 0 used_image_names: set[str] = set() @@ -434,26 +432,20 @@ def _vlm_task(page_num, img_name): else: chunk_path = safe_title - # Build df row (11 columns) # Atlas chunks are image-primary: use IMAGE marker directly. # find_matches_parsing() prepends "PTXT\n" which causes downstream # chunk type classifier to misclassify as "text" instead of "image". - match_type = "image" tokens = tokenize2stw_remove([content], stopwords) - df_list.append( - [ - content, # content - chunk_path, # path - match_type, # type - len(content), # length - "", # keywords - "", # summary - know_id, # know_id - tokens, # tokens - "", # connectto - time_stamp, # addtime - str(page_num), # page_nums - ] + rows_builder.append( + ParsedRow( + content=content, + path=chunk_path, + type="image", + know_id=know_id, + addtime=time_stamp, + tokens=tokens, + page_nums=str(page_num), + ) ) # Build custom md line @@ -471,8 +463,9 @@ def _vlm_task(page_num, img_name): with open(custom_md_path, "w", encoding="utf-8") as f: f.write("\n".join(custom_md_lines)) + df = rows_builder.to_dataframe() logger.info( - f"📐 Atlas pipeline complete: {len(df_list)} chunks created " + f"📐 Atlas pipeline complete: {len(df)} chunks created " f"(skipped {len(toc_page_set)} TOC + {skipped_null} null pages, total {total_pages})" ) @@ -480,9 +473,5 @@ def _vlm_task(page_num, img_name): # Currently atlas chunks are flat with no parent-child relationships. # Future: integrate with hierarchy builder for unified schema. - # ── Build DataFrame ── - all_cols = settings.ALL_DF_COLS.split(",") - - df = pd.DataFrame(df_list, columns=all_cols) df = process_dup_paths_df(df) return df diff --git a/apps/worker/app/services/document_parser/formats/docx/asset_accumulator.py b/apps/worker/app/services/document_parser/formats/docx/asset_accumulator.py new file mode 100644 index 000000000..f69526e95 --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/docx/asset_accumulator.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +from app.services.document_parser.formats.docx.asset_store import DocxAssetStore + + +ImageHandler = Callable[ + [ + list[list[object]], + dict[str, Any], + DocxAssetStore, + list[dict[str, Any]], + str, + int, + bool, + dict[str, dict[str, str]], + ], + tuple[list[dict[str, Any]], list[list[object]], bool], +] + +TableHandler = Callable[..., tuple[list[dict[str, Any]], list[list[object]], int]] + + +@dataclass +class DocxAssetAccumulator: + asset_store: DocxAssetStore + should_summary_image: bool + should_summary_table: bool + image_handler: ImageHandler + table_handler: TableHandler + _rows: list[list[object]] = field(default_factory=list) + _image_count: int = 0 + _table_count: int = 0 + _seen_images: dict[str, dict[str, str]] = field(default_factory=dict) + + def append_image( + self, + image_meta: dict[str, Any], + headings_stack: list[dict[str, Any]], + current_heading: str, + ) -> list[dict[str, Any]]: + headings_stack, self._rows, is_new_image = self.image_handler( + self._rows, + image_meta, + self.asset_store, + headings_stack, + current_heading, + self._image_count, + self.should_summary_image, + self._seen_images, + ) + if is_new_image: + self._image_count += 1 + return headings_stack + + def append_table( + self, + block: Any, + headings_stack: list[dict[str, Any]], + current_heading: str, + cell_images: Any, + ) -> list[dict[str, Any]]: + headings_stack, self._rows, self._image_count = self.table_handler( + self._rows, + block, + self.asset_store, + headings_stack, + current_heading, + self._table_count, + summary_table=self.should_summary_table, + summary_image=self.should_summary_image, + cell_images=cell_images, + img_count=self._image_count, + seen_images=self._seen_images, + ) + self._table_count += 1 + return headings_stack + + def rows(self) -> list[list[object]]: + return self._rows diff --git a/apps/worker/app/services/document_parser/formats/docx/asset_store.py b/apps/worker/app/services/document_parser/formats/docx/asset_store.py new file mode 100644 index 000000000..f97f8b0a8 --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/docx/asset_store.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +import os +import shutil +from dataclasses import dataclass + + +@dataclass(frozen=True) +class StoredDocxAsset: + absolute_path: str + relative_path: str + name: str + extension: str + + +class DocxAssetStore: + def __init__(self, output_dir: str) -> None: + self.output_dir = output_dir + self.image_dir = os.path.join(output_dir, "images") + self.table_dir = os.path.join(output_dir, "tables") + + def reset(self) -> None: + self._reset_dir(self.table_dir) + self._reset_dir(self.image_dir) + + def write_image(self, name: str, extension: str, data: bytes) -> StoredDocxAsset: + absolute_path = os.path.join(self.image_dir, f"{name}{extension}") + with open(absolute_path, "wb") as image_file: + image_file.write(data) + return StoredDocxAsset( + absolute_path=absolute_path, + relative_path=f"images/{name}{extension}", + name=name, + extension=extension, + ) + + def rename_image(self, asset: StoredDocxAsset, new_name: str) -> StoredDocxAsset: + new_absolute_path = os.path.join(self.image_dir, f"{new_name}{asset.extension}") + if asset.absolute_path != new_absolute_path: + os.rename(asset.absolute_path, new_absolute_path) + return StoredDocxAsset( + absolute_path=new_absolute_path, + relative_path=f"images/{new_name}{asset.extension}", + name=new_name, + extension=asset.extension, + ) + + def write_table(self, name: str, html: str) -> StoredDocxAsset: + absolute_path = os.path.join(self.table_dir, f"{name}.html") + with open(absolute_path, "w", encoding="utf-8") as table_file: + table_file.write(html) + return StoredDocxAsset( + absolute_path=absolute_path, + relative_path=f"tables/{name}.html", + name=name, + extension=".html", + ) + + @staticmethod + def _reset_dir(directory: str) -> None: + if os.path.isdir(directory): + shutil.rmtree(directory) + os.makedirs(directory, exist_ok=True) diff --git a/apps/worker/app/services/document_parser/formats/docx/block_stream.py b/apps/worker/app/services/document_parser/formats/docx/block_stream.py new file mode 100644 index 000000000..dafe28aba --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/docx/block_stream.py @@ -0,0 +1,315 @@ +# pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportGeneralTypeIssues=false +from __future__ import annotations + +import io +import zipfile + +from app.services.document_parser.formats.docx.toc import detect_doc_tocs, detect_sdt_toc +from docx import Document +from docx.oxml.table import CT_Tbl +from docx.oxml.text.paragraph import CT_P +from docx.table import Table +from docx.text.paragraph import Paragraph +from loguru import logger +from lxml import etree + + +def iter_block_items(doc_data): + doc_stream = io.BytesIO(doc_data) + doc = Document(doc_stream) + + # python-docx mapping + p_tbl_map = [] + for child in doc.element.body: + if isinstance(child, CT_P): + p_tbl_map.append(("p", child)) + elif isinstance(child, CT_Tbl): + p_tbl_map.append(("tbl", child)) + + with zipfile.ZipFile(io.BytesIO(doc_data), "r") as docx: + xml = docx.read("word/document.xml") + rels = etree.fromstring(docx.read("word/_rels/document.xml.rels")) + rel_map = { + r.get("Id"): r.get("Target") for r in rels.findall(".//{*}Relationship") + } + ns = { + "a": "http://schemas.openxmlformats.org/drawingml/2006/main", + "r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships", + "w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main", + "v": "urn:schemas-microsoft-com:vml", + "o": "urn:schemas-microsoft-com:office:office", + } + r_ns = "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}" + + root = etree.fromstring(xml) + body = root.find(".//w:body", namespaces=ns) + + ele_num = 1 + map_index = 0 # point to p_tbl_map + toc_field_active = False + + for elem in body.iterchildren(): + if not isinstance(elem.tag, str): + continue + + tag = etree.QName(elem.tag).localname + + # --- SDT (Structured Document Tag) container --- + # TOC generated by MS Word is usually in sdt + if tag == "sdt": + sdt_toc_info = detect_sdt_toc(elem, ns) + is_toc_sdt = sdt_toc_info["is_toc_sdt"] + + sdt_content = elem.find(".//w:sdtContent", namespaces=ns) + if sdt_content is not None: + for p_elem in sdt_content.findall(".//w:p", namespaces=ns): + texts = p_elem.xpath(".//w:t/text()", namespaces=ns) + text = "".join(texts).strip() + + if is_toc_sdt: + label = "TOC-AREA" + toc_info = detect_doc_tocs(p_elem, ns) + else: + toc_info = detect_doc_tocs(p_elem, ns) + if toc_info["is_style"] or toc_info["is_field_start"]: + label = "TOC-AREA" + else: + label = "PTXT" + + if text: + meta = None + if "TOC" in label: + meta = { + "toc_level": toc_info.get("toc_level"), + "toc_outline_level": toc_info.get("outline_level"), + "toc_left_indent": toc_info.get("left_indent"), + "toc_style_name": toc_info.get("style_name"), + "toc_source": "sdt", + } + yield ele_num, text, label, meta + ele_num += 1 + continue + + # --- text paras --- + if tag == "p": + texts = elem.xpath(".//w:t/text()", namespaces=ns) + text = "".join(texts).strip() + + if map_index < len(p_tbl_map) and p_tbl_map[map_index][0] == "p": + p_obj = Paragraph(p_tbl_map[map_index][1], doc) + else: + p_obj = None + + toc_info = detect_doc_tocs(elem, ns) + if toc_info["is_field_start"]: + toc_field_active = True + + if toc_info["is_style"] or toc_field_active: + label = "TOC-AREA" + else: + label = "PTXT" + + if text or p_obj is not None: + meta = None + if "TOC" in label: + meta = { + "toc_level": toc_info.get("toc_level"), + "toc_outline_level": toc_info.get("outline_level"), + "toc_left_indent": toc_info.get("left_indent"), + "toc_style_name": toc_info.get("style_name"), + "toc_source": "paragraph", + } + yield ele_num, p_obj or text, label, meta + ele_num += 1 + + # images (DrawingML: ) + seen_rids = set() + blips = elem.xpath(".//a:blip", namespaces=ns) + for b in blips: + rid = b.get(f"{r_ns}embed") + if not rid or rid in seen_rids: + continue + target = rel_map.get(rid) + if not target or not target.startswith("media/"): + continue + seen_rids.add(rid) + data = docx.read("word/" + target) + yield ( + ele_num, + None, + "IMAGE", + { + "image_name": target.split("/")[-1], + "from": "paragraph", + "size": len(data), + "data": data, + }, + ) + ele_num += 1 + + # TODO: Re-evaluate VML group extraction strategy. + # Complex VML composite images () are currently skipped because extracting + # piece-by-piece loses textual overlay and positioning. + # Future plan: Use LibreOffice headless conversion to render the entire document + # and map the perfectly rendered images back to the layout via text anchors. + # + # Temporary: detect VML-only paragraphs and inject a placeholder so the + # paragraph isn't silently swallowed, leaving its parent section empty. + if not text and not seen_rids: + # No text and no DrawingML images — check for VML content + vml_groups = elem.xpath(".//v:group", namespaces=ns) + vml_images_check = elem.xpath(".//v:imagedata", namespaces=ns) + if vml_groups or vml_images_check: + vml_placeholder = "[VML graphic \u2014 extraction not yet supported]" + yield ele_num, vml_placeholder, "PTXT", None + ele_num += 1 + logger.debug( + f"Injected VML placeholder for paragraph with " + f"{len(vml_groups)} v:group, {len(vml_images_check)} v:imagedata" + ) + """ + # images (VML: ) — convert to PNG + from PIL import Image as PILImage + + vml_images = elem.xpath(".//v:imagedata", namespaces=ns) + for v in vml_images: + rid = v.get(f"{r_ns}id") + if not rid or rid in seen_rids: + continue + target = rel_map.get(rid) + if not target or not target.startswith("media/"): + continue + seen_rids.add(rid) + raw_data = docx.read("word/" + target) + # Convert to PNG for uniform downstream handling + try: + pil_img = PILImage.open(io.BytesIO(raw_data)) + png_buf = io.BytesIO() + pil_img.save(png_buf, format="PNG") + png_data = png_buf.getvalue() + except Exception as e: + logger.warning(f"Failed to convert VML image to PNG: {e}") + continue + orig_name = target.split("/")[-1] + png_name = os.path.splitext(orig_name)[0] + ".png" + yield ( + ele_num, + None, + "IMAGE", + { + "image_name": png_name, + "from": "paragraph_vml", + "size": len(png_data), + "data": png_data, + }, + ) + ele_num += 1 + """ + map_index += 1 + + if toc_info["is_field_end"]: + toc_field_active = False + + # --- tables --- + elif tag == "tbl": + if map_index < len(p_tbl_map) and p_tbl_map[map_index][0] == "tbl": + tbl = Table(p_tbl_map[map_index][1], doc) + else: + tbl = Table(elem, doc) + + # Extract images from each cell, keyed by (row_idx, col_idx) + cell_images = {} # {(row_idx, col_idx): [{'image_name', 'data', 'size'}]} + for row_idx, tr in enumerate(elem.findall(".//w:tr", namespaces=ns)): + for col_idx, tc in enumerate(tr.findall(".//w:tc", namespaces=ns)): + cell_seen_rids = set() + imgs_in_cell = [] + # DrawingML images in cell + blips = tc.xpath(".//a:blip", namespaces=ns) + for b in blips: + rid = b.get(f"{r_ns}embed") + if not rid or rid in cell_seen_rids: + continue + target = rel_map.get(rid) + if not target or not target.startswith("media/"): + continue + cell_seen_rids.add(rid) + data = docx.read("word/" + target) + if ( + len(data) < 10 * 1024 + ): # Skip small images (<10KB, likely icons) + continue + imgs_in_cell.append( + { + "image_name": target.split("/")[-1], + "data": data, + "size": len(data), + } + ) + # TODO: VML in tables is temporarily skipped to avoid extracting + # fragmented textless background images. (Same as paragraph VML logic) + """ + # VML images in cell — convert to PNG + from PIL import Image as PILImage + + vml_in_cell = tc.xpath(".//v:imagedata", namespaces=ns) + for v in vml_in_cell: + rid = v.get(f"{r_ns}id") + if not rid or rid in cell_seen_rids: + continue + target = rel_map.get(rid) + if not target or not target.startswith("media/"): + continue + cell_seen_rids.add(rid) + raw_data = docx.read("word/" + target) + try: + pil_img = PILImage.open(io.BytesIO(raw_data)) + png_buf = io.BytesIO() + pil_img.save(png_buf, format="PNG") + png_data = png_buf.getvalue() + except Exception as e: + logger.warning( + f"Failed to convert VML cell image to PNG: {e}" + ) + continue + if len(png_data) < 10 * 1024: + continue + orig_name = target.split("/")[-1] + png_name = os.path.splitext(orig_name)[0] + ".png" + imgs_in_cell.append( + { + "image_name": png_name, + "data": png_data, + "size": len(png_data), + } + ) + """ + if imgs_in_cell: + cell_images[(row_idx, col_idx)] = imgs_in_cell + + yield ele_num, tbl, "TABLE", cell_images if cell_images else None + ele_num += 1 + map_index += 1 + else: + continue + + # --- handle p_tbl_map at the end --- + while map_index < len(p_tbl_map): + tag, node = p_tbl_map[map_index] + if tag == "p": + toc_info = detect_doc_tocs(node, ns) + label = "TOC-AREA" if toc_info["is_style"] else "PTXT" + meta = None + if "TOC" in label: + meta = { + "toc_level": toc_info.get("toc_level"), + "toc_outline_level": toc_info.get("outline_level"), + "toc_left_indent": toc_info.get("left_indent"), + "toc_style_name": toc_info.get("style_name"), + "toc_source": "tail-map", + } + yield ele_num, Paragraph(node, doc), label, meta + elif tag == "tbl": + yield ele_num, Table(node, doc), "TABLE", None + ele_num += 1 + map_index += 1 + diff --git a/apps/worker/app/services/document_parser/doc_parser.py b/apps/worker/app/services/document_parser/formats/docx/parser.py similarity index 50% rename from apps/worker/app/services/document_parser/doc_parser.py rename to apps/worker/app/services/document_parser/formats/docx/parser.py index a253100c2..59c7b297c 100755 --- a/apps/worker/app/services/document_parser/doc_parser.py +++ b/apps/worker/app/services/document_parser/formats/docx/parser.py @@ -1,47 +1,45 @@ # pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportGeneralTypeIssues=false -import io import json import os -import shutil -import zipfile import pandas as pd -from app.services.common.kb_utils import ( +from app.services.document_parser.tables.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.formats.docx.asset_accumulator import DocxAssetAccumulator +from app.services.document_parser.formats.docx.asset_store import DocxAssetStore +from app.services.document_parser.formats.docx.block_stream import iter_block_items +from app.services.document_parser.formats.docx.table_html import table2html +from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.assets.inline_asset import ( + build_image_asset_row, + build_table_asset_row, +) +from app.services.document_parser.support.parser_rows import ParsedRow, ParsedRowsBuilder +from app.services.document_parser.support.path_helpers import ( find_matches_parsing, - gen_str_codes, - get_str_time, - process_dup_paths_df, process_path_texts, remove_spaces, ) -from app.services.document_parser.html_parser import table2html -from app.services.document_parser.image_parser import ( +from app.services.document_parser.structure.heading_hierarchy import ( + HeadingHierarchyInput, + predict_heading_hierarchy, +) +from app.services.document_parser.formats.image.parser import ( _get_vision_client, ask_image, perceptual_hash, ) -from app.services.document_parser.layout_parser import pred_titles -from app.services.document_parser.table_parser import sanitize_table_name_from_header -from app.services.document_parser.toc_parser import ( - build_docx_toc_hierarchies, - detect_doc_tocs, - detect_sdt_toc, -) -from app.services.document_parser.txt_parser import postprocess_leaf_dics -from docx import Document -from docx.oxml.table import CT_Tbl -from docx.oxml.text.paragraph import CT_P -from docx.table import Table +from app.services.document_parser.tables.table_text_parser import sanitize_table_name_from_header +from app.services.document_parser.formats.docx.toc import build_docx_toc_hierarchies +from app.services.document_parser.formats.text.parser import postprocess_leaf_dics from docx.text.paragraph import Paragraph from loguru import logger -from lxml import etree from shared.core.config import settings from shared.core.exceptions.domain_exceptions import DocxParsingException from shared.core.exceptions.knowhere_exception import KnowhereException from shared.utils.chunk_refs import build_chunk_ref, has_chunk_ref -from shared.utils.CommonHelperSync import load_file_bytes -from shared.utils.file_utils import path_handle +from app.services.common.file_loading import load_file_bytes +from app.services.common.file_utils import path_handle from shared.utils.text_utils import tokenize2stw_remove @@ -74,7 +72,7 @@ def _find_img_context(headings_stack, max_chars=100): Returns: The nearest valid text context, or empty string if none found """ - from app.services.common.kb_utils import truncate_text + from app.services.document_parser.support.text_helpers import truncate_text try: content_list = headings_stack[-1].get("content", []) @@ -100,7 +98,7 @@ def _find_img_context(headings_stack, max_chars=100): def handle_image( df_list, img_file, - img_dir, + asset_store, headings_stack, current_heading, img_count, @@ -116,19 +114,13 @@ def handle_image( cached = seen_images[img_hash] headings_stack[-1]["content"].append(cached["image_ref"]) df_list.append( - [ - cached["image_ref"], - cached["img_path"], - "image", - len(cached["image_ref"]), - "", - cached["img_summary_field"], - cached["temp_uid"], - "", - "", - time_stamp, - "", - ] + build_image_asset_row( + content=cached["image_ref"], + relative_path=cached["img_path"], + summary=cached["img_summary_field"], + know_id=cached["temp_uid"], + addtime=time_stamp, + ).to_list() ) logger.debug(f"Skipped duplicate image (hash={img_hash[:12]}...)") return headings_stack, df_list, False # False = cache hit, don't increment @@ -143,21 +135,21 @@ def handle_image( raw_img_name = process_path_texts( f"image-{str(img_count + 1)} {current_heading} {last_context}", last=30 ) - img_raw_path = os.path.join(img_dir, f"{raw_img_name}{img_ext}") - - with open(img_raw_path, "wb") as image_file: - image_file.write(img_file["data"]) + raw_image_asset = asset_store.write_image(raw_img_name, img_ext, img_file["data"]) # LLM title + summary (optional, with fallback to last_context) llm_title = None llm_summary = None if smart_summary: - from app.services.document_parser.txt_parser import split_title_summary + from app.services.document_parser.formats.text.parser import split_title_summary # TODO: Risk of missing text content if the image is a screenshot of pure text. # Consider adding judge-image-type and OCR fallback as done in image_parser.parse_image. llm_resp = ask_image( - client, img_dir, [f"{raw_img_name}{img_ext}"], title_text=last_context + client, + asset_store.image_dir, + [f"{raw_img_name}{img_ext}"], + title_text=last_context, ) if llm_resp: llm_title, llm_summary = split_title_summary(llm_resp) @@ -176,8 +168,7 @@ def handle_image( img_name = process_path_texts( f"image-{str(img_count + 1)} {current_heading} {img_title or ''}", last=30 ) - img_path = os.path.join(img_dir, f"{img_name}{img_ext}") - os.rename(img_raw_path, img_path) # if summary fails, renaming is not applied + image_asset = asset_store.rename_image(raw_image_asset, img_name) temp_uid = gen_str_codes(img_hash) @@ -187,8 +178,7 @@ def handle_image( else: img_summary_field = image_index - img_path = f"images/{img_name}{img_ext}" - img_ref = build_chunk_ref(img_path) + img_ref = build_chunk_ref(image_asset.relative_path) # Build image_ref for heading_stack: optional summary + image path ref if img_summary: @@ -198,25 +188,19 @@ def handle_image( headings_stack[-1]["content"].append(image_ref) df_list.append( - [ - image_ref, - img_path, - "image", - len(image_ref), - "", - img_summary_field, - temp_uid, - "", - "", - time_stamp, - "", - ] + build_image_asset_row( + content=image_ref, + relative_path=image_asset.relative_path, + summary=img_summary_field, + know_id=temp_uid, + addtime=time_stamp, + ).to_list() ) # Cache result for document-level dedup if seen_images is not None: seen_images[img_hash] = { - "img_path": img_path, + "img_path": image_asset.relative_path, "image_ref": image_ref, "img_summary_field": img_summary_field, "temp_uid": temp_uid, @@ -236,7 +220,7 @@ def _first_cols_rows(table_block, max_items=10, max_chars=20): Returns: Tuple of (first_row_text, first_col_text) with ' | ' as separator """ - from app.services.common.kb_utils import truncate_text + from app.services.document_parser.support.text_helpers import truncate_text first_row_text = "" first_col_text = "" @@ -275,14 +259,13 @@ def _first_cols_rows(table_block, max_items=10, max_chars=20): def handle_table( df_list, block, - tb_dir, + asset_store, headings_stack, current_heading, table_count, summary_table=False, summary_image=False, cell_images=None, - img_dir=None, img_count=0, seen_images=None, ): @@ -329,9 +312,9 @@ def handle_table( img_name = process_path_texts( f"table-{table_count + 1}-{image_index} {current_heading}", last=30 ) - img_save_path = os.path.join(img_dir, f"{img_name}{img_ext}") - with open(img_save_path, "wb") as f: - f.write(img_data["data"]) + image_asset = asset_store.write_image( + img_name, img_ext, img_data["data"] + ) # LLM summary (optional) img_summary = None @@ -340,7 +323,7 @@ def handle_table( client = _get_vision_client() img_summary = ask_image( client, - img_dir, + asset_store.image_dir, [f"{img_name}{img_ext}"], title_text=current_heading, ) @@ -355,32 +338,25 @@ def handle_table( img_summary_field = ( f"{image_index}\n{img_summary}" if img_summary else image_index ) - relative_img_path = f"images/{img_name}{img_ext}" - img_ref = build_chunk_ref(relative_img_path) + img_ref = build_chunk_ref(image_asset.relative_path) if img_summary: image_ref = f"\n{img_summary}\n{img_ref}\n" else: image_ref = f"\n{img_ref}\n" table_img_entries.append( - [ - image_ref, - relative_img_path, - "image", - len(image_ref), - "", - img_summary_field, - temp_uid, - "", - "", - time_stamp, - "", - ] + build_image_asset_row( + content=image_ref, + relative_path=image_asset.relative_path, + summary=img_summary_field, + know_id=temp_uid, + addtime=time_stamp, + ).to_list() ) # Cache result for document-level dedup if seen_images is not None: seen_images[cell_img_hash] = { - "img_path": relative_img_path, + "img_path": image_asset.relative_path, "image_ref": image_ref, "img_summary_field": img_summary_field, "temp_uid": temp_uid, @@ -392,7 +368,6 @@ def handle_table( f"Extracted {sum(len(v) for v in cell_images.values())} images from table-{table_count + 1} cells" ) - # Generate HTML with image descriptions embedded tb_html_str = table2html( block, cell_image_map=cell_image_map if cell_image_map else None @@ -417,7 +392,7 @@ def handle_table( llm_summary = None tb_keywords = "" if summary_table: - from app.services.document_parser.txt_parser import ( + from app.services.document_parser.formats.text.parser import ( extract_title_keywords_summary, ) @@ -438,14 +413,8 @@ def handle_table( tb_name = path_handle( f"table-{str(table_count + 1)} {effective_name}", mode="clean_single" ) - tb_path = os.path.join(tb_dir, f"{tb_name}.html") - - with open(tb_path, "w", encoding="utf-8") as f: - f.write(tb_html_str) - - # Use relative path for tables (avoid absolute path in path column) - tb_path = f"tables/{tb_name}.html" - tb_ref = build_chunk_ref(tb_path) + table_asset = asset_store.write_table(tb_name, tb_html_str) + tb_ref = build_chunk_ref(table_asset.relative_path) # Build table_ref for heading_stack: optional LLM summary + table path ref if llm_summary: table_ref = f"\n{llm_summary}\n{tb_ref}\n" @@ -453,331 +422,24 @@ def handle_table( table_ref = f"\n{tb_ref}\n" headings_stack[-1]["content"].append(table_ref) df_list.append( - [ - tb_html_str, - tb_path, - "table", - len(tb_html_str), - tb_keywords, - tb_summary, - temp_uid, - "", - "", - time_stamp, - "", - ] + build_table_asset_row( + content=tb_html_str, + relative_path=table_asset.relative_path, + summary=tb_summary, + keywords=tb_keywords, + know_id=temp_uid, + addtime=time_stamp, + ).to_list() ) return headings_stack, df_list, img_count -def iter_block_items(doc_data): - doc_stream = io.BytesIO(doc_data) - doc = Document(doc_stream) - - # python-docx mapping - p_tbl_map = [] - for child in doc.element.body: - if isinstance(child, CT_P): - p_tbl_map.append(("p", child)) - elif isinstance(child, CT_Tbl): - p_tbl_map.append(("tbl", child)) - - with zipfile.ZipFile(io.BytesIO(doc_data), "r") as docx: - xml = docx.read("word/document.xml") - rels = etree.fromstring(docx.read("word/_rels/document.xml.rels")) - rel_map = { - r.get("Id"): r.get("Target") for r in rels.findall(".//{*}Relationship") - } - ns = { - "a": "http://schemas.openxmlformats.org/drawingml/2006/main", - "r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships", - "w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main", - "v": "urn:schemas-microsoft-com:vml", - "o": "urn:schemas-microsoft-com:office:office", - } - r_ns = "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}" - - root = etree.fromstring(xml) - body = root.find(".//w:body", namespaces=ns) - - ele_num = 1 - map_index = 0 # point to p_tbl_map - toc_field_active = False - - for elem in body.iterchildren(): - if not isinstance(elem.tag, str): - continue - - tag = etree.QName(elem.tag).localname - - # --- SDT (Structured Document Tag) container --- - # TOC generated by MS Word is usually in sdt - if tag == "sdt": - sdt_toc_info = detect_sdt_toc(elem, ns) - is_toc_sdt = sdt_toc_info["is_toc_sdt"] - - sdt_content = elem.find(".//w:sdtContent", namespaces=ns) - if sdt_content is not None: - for p_elem in sdt_content.findall(".//w:p", namespaces=ns): - texts = p_elem.xpath(".//w:t/text()", namespaces=ns) - text = "".join(texts).strip() - - if is_toc_sdt: - label = "TOC-AREA" - toc_info = detect_doc_tocs(p_elem, ns) - else: - toc_info = detect_doc_tocs(p_elem, ns) - if toc_info["is_style"] or toc_info["is_field_start"]: - label = "TOC-AREA" - else: - label = "PTXT" - - if text: - meta = None - if "TOC" in label: - meta = { - "toc_level": toc_info.get("toc_level"), - "toc_outline_level": toc_info.get("outline_level"), - "toc_left_indent": toc_info.get("left_indent"), - "toc_style_name": toc_info.get("style_name"), - "toc_source": "sdt", - } - yield ele_num, text, label, meta - ele_num += 1 - continue - - # --- text paras --- - if tag == "p": - texts = elem.xpath(".//w:t/text()", namespaces=ns) - text = "".join(texts).strip() - - if map_index < len(p_tbl_map) and p_tbl_map[map_index][0] == "p": - p_obj = Paragraph(p_tbl_map[map_index][1], doc) - else: - p_obj = None - - toc_info = detect_doc_tocs(elem, ns) - if toc_info["is_field_start"]: - toc_field_active = True - - if toc_info["is_style"] or toc_field_active: - label = "TOC-AREA" - else: - label = "PTXT" - - if text or p_obj is not None: - meta = None - if "TOC" in label: - meta = { - "toc_level": toc_info.get("toc_level"), - "toc_outline_level": toc_info.get("outline_level"), - "toc_left_indent": toc_info.get("left_indent"), - "toc_style_name": toc_info.get("style_name"), - "toc_source": "paragraph", - } - yield ele_num, p_obj or text, label, meta - ele_num += 1 - - # images (DrawingML: ) - seen_rids = set() - blips = elem.xpath(".//a:blip", namespaces=ns) - for b in blips: - rid = b.get(f"{r_ns}embed") - if not rid or rid in seen_rids: - continue - target = rel_map.get(rid) - if not target or not target.startswith("media/"): - continue - seen_rids.add(rid) - data = docx.read("word/" + target) - yield ( - ele_num, - None, - "IMAGE", - { - "image_name": target.split("/")[-1], - "from": "paragraph", - "size": len(data), - "data": data, - }, - ) - ele_num += 1 - - # TODO: Re-evaluate VML group extraction strategy. - # Complex VML composite images () are currently skipped because extracting - # piece-by-piece loses textual overlay and positioning. - # Future plan: Use LibreOffice headless conversion to render the entire document - # and map the perfectly rendered images back to the layout via text anchors. - # - # Temporary: detect VML-only paragraphs and inject a placeholder so the - # paragraph isn't silently swallowed, leaving its parent section empty. - if not text and not seen_rids: - # No text and no DrawingML images — check for VML content - vml_groups = elem.xpath(".//v:group", namespaces=ns) - vml_images_check = elem.xpath(".//v:imagedata", namespaces=ns) - if vml_groups or vml_images_check: - vml_placeholder = "[VML graphic \u2014 extraction not yet supported]" - yield ele_num, vml_placeholder, "PTXT", None - ele_num += 1 - logger.debug( - f"Injected VML placeholder for paragraph with " - f"{len(vml_groups)} v:group, {len(vml_images_check)} v:imagedata" - ) - """ - # images (VML: ) — convert to PNG - from PIL import Image as PILImage - - vml_images = elem.xpath(".//v:imagedata", namespaces=ns) - for v in vml_images: - rid = v.get(f"{r_ns}id") - if not rid or rid in seen_rids: - continue - target = rel_map.get(rid) - if not target or not target.startswith("media/"): - continue - seen_rids.add(rid) - raw_data = docx.read("word/" + target) - # Convert to PNG for uniform downstream handling - try: - pil_img = PILImage.open(io.BytesIO(raw_data)) - png_buf = io.BytesIO() - pil_img.save(png_buf, format="PNG") - png_data = png_buf.getvalue() - except Exception as e: - logger.warning(f"Failed to convert VML image to PNG: {e}") - continue - orig_name = target.split("/")[-1] - png_name = os.path.splitext(orig_name)[0] + ".png" - yield ( - ele_num, - None, - "IMAGE", - { - "image_name": png_name, - "from": "paragraph_vml", - "size": len(png_data), - "data": png_data, - }, - ) - ele_num += 1 - """ - map_index += 1 - - if toc_info["is_field_end"]: - toc_field_active = False - - # --- tables --- - elif tag == "tbl": - if map_index < len(p_tbl_map) and p_tbl_map[map_index][0] == "tbl": - tbl = Table(p_tbl_map[map_index][1], doc) - else: - tbl = Table(elem, doc) - - # Extract images from each cell, keyed by (row_idx, col_idx) - cell_images = {} # {(row_idx, col_idx): [{'image_name', 'data', 'size'}]} - for row_idx, tr in enumerate(elem.findall(".//w:tr", namespaces=ns)): - for col_idx, tc in enumerate(tr.findall(".//w:tc", namespaces=ns)): - cell_seen_rids = set() - imgs_in_cell = [] - # DrawingML images in cell - blips = tc.xpath(".//a:blip", namespaces=ns) - for b in blips: - rid = b.get(f"{r_ns}embed") - if not rid or rid in cell_seen_rids: - continue - target = rel_map.get(rid) - if not target or not target.startswith("media/"): - continue - cell_seen_rids.add(rid) - data = docx.read("word/" + target) - if ( - len(data) < 10 * 1024 - ): # Skip small images (<10KB, likely icons) - continue - imgs_in_cell.append( - { - "image_name": target.split("/")[-1], - "data": data, - "size": len(data), - } - ) - # TODO: VML in tables is temporarily skipped to avoid extracting - # fragmented textless background images. (Same as paragraph VML logic) - """ - # VML images in cell — convert to PNG - from PIL import Image as PILImage - - vml_in_cell = tc.xpath(".//v:imagedata", namespaces=ns) - for v in vml_in_cell: - rid = v.get(f"{r_ns}id") - if not rid or rid in cell_seen_rids: - continue - target = rel_map.get(rid) - if not target or not target.startswith("media/"): - continue - cell_seen_rids.add(rid) - raw_data = docx.read("word/" + target) - try: - pil_img = PILImage.open(io.BytesIO(raw_data)) - png_buf = io.BytesIO() - pil_img.save(png_buf, format="PNG") - png_data = png_buf.getvalue() - except Exception as e: - logger.warning( - f"Failed to convert VML cell image to PNG: {e}" - ) - continue - if len(png_data) < 10 * 1024: - continue - orig_name = target.split("/")[-1] - png_name = os.path.splitext(orig_name)[0] + ".png" - imgs_in_cell.append( - { - "image_name": png_name, - "data": png_data, - "size": len(png_data), - } - ) - """ - if imgs_in_cell: - cell_images[(row_idx, col_idx)] = imgs_in_cell - - yield ele_num, tbl, "TABLE", cell_images if cell_images else None - ele_num += 1 - map_index += 1 - else: - continue - - # --- handle p_tbl_map at the end --- - while map_index < len(p_tbl_map): - tag, node = p_tbl_map[map_index] - if tag == "p": - toc_info = detect_doc_tocs(node, ns) - label = "TOC-AREA" if toc_info["is_style"] else "PTXT" - meta = None - if "TOC" in label: - meta = { - "toc_level": toc_info.get("toc_level"), - "toc_outline_level": toc_info.get("outline_level"), - "toc_left_indent": toc_info.get("left_indent"), - "toc_style_name": toc_info.get("style_name"), - "toc_source": "tail-map", - } - yield ele_num, Paragraph(node, doc), label, meta - elif tag == "tbl": - yield ele_num, Table(node, doc), "TABLE", None - ele_num += 1 - map_index += 1 - - def parse_docx( docx_path, llm_paras, output_dir=None, filename="", file_url="", - start_text="", - end_text="", relative_root=None, ): doc_data = load_file_bytes(docx_path, file_url=file_url) @@ -787,16 +449,8 @@ def parse_docx( headings_stack = [{"level": -1, "content": doc_structure}] current_heading = "" - # Clean old artifacts to prevent accumulation across debug runs. - # In production each job uses a fresh workspace so rmtree never triggers. - tb_dir = os.path.join(output_dir, "tables") - if os.path.isdir(tb_dir): - shutil.rmtree(tb_dir) - os.makedirs(tb_dir, exist_ok=True) - img_dir = os.path.join(output_dir, "images") - if os.path.isdir(img_dir): - shutil.rmtree(img_dir) - os.makedirs(img_dir, exist_ok=True) + asset_store = DocxAssetStore(output_dir) + asset_store.reset() block_tuples = list(iter_block_items(doc_data)) # Record first TOC block position before filtering, for pre-TOC exclusion in pred_titles @@ -830,15 +484,17 @@ def parse_docx( if llm_paras else (settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL) ) - heading_candidates = pred_titles( - heading_infos, - doc_type="docx", - toc_hierarchies=toc_hierarchies or None, - enable_regx=True, - smart_parse=smart_title_parse, - model_name=model_name, - output_dir=output_dir, - first_toc_ele_num=first_toc_ele_num, + heading_candidates = predict_heading_hierarchy( + HeadingHierarchyInput( + infos=heading_infos, + doc_type="docx", + toc_hierarchies=toc_hierarchies or None, + enable_regex=True, + smart_parse=smart_title_parse, + model_name=model_name, + output_dir=output_dir, + first_toc_ele_num=first_toc_ele_num, + ) ) if len(heading_candidates) > 0 and not (heading_candidates["level"] == -1).all(): @@ -854,10 +510,13 @@ def parse_docx( headings_stack.append(new_content) logger.debug("⚠️no headings detected, using file name or mine a heading=>", text) - df_list = [] - table_count = 0 - image_count = 0 - _seen_images: dict[str, dict] = {} # sha256_hex -> cached image info for dedup + asset_accumulator = DocxAssetAccumulator( + asset_store=asset_store, + should_summary_image=llm_paras["summary_image"], + should_summary_table=llm_paras["summary_table"], + image_handler=handle_image, + table_handler=handle_table, + ) logger.debug("Parsing docx file... total_blocks={}", len(block_tuples)) for block_tuple in block_tuples: @@ -872,12 +531,8 @@ def parse_docx( outline_level = outline_dic.get(ele_num, -1) if outline_level > 0: # logger.debug('Found a title: ', text, ' current level: ', outline_level) - try: - last_heading = headings_stack[-1]["heading"] - if last_heading == text: - continue - except Exception: - pass + if headings_stack and headings_stack[-1].get("heading") == text: + continue while headings_stack and headings_stack[-1]["level"] >= outline_level: headings_stack.pop() @@ -895,43 +550,27 @@ def parse_docx( if meta and meta.get("size", 0) < 10 * 1024: continue - headings_stack, df_list, is_new = handle_image( - df_list, + headings_stack = asset_accumulator.append_image( meta, - img_dir, headings_stack, current_heading, - image_count, - llm_paras["summary_image"], - seen_images=_seen_images, ) - if is_new: - image_count += 1 current_heading = last_heading_before_block elif label == "TABLE": # TODO: handle cross-page tables - headings_stack, df_list, image_count = handle_table( - df_list, + headings_stack = asset_accumulator.append_table( block, - tb_dir, headings_stack, current_heading, - table_count, - summary_table=llm_paras["summary_table"], - summary_image=llm_paras["summary_image"], cell_images=meta, - img_dir=img_dir, - img_count=image_count, - seen_images=_seen_images, ) - table_count += 1 current_heading = last_heading_before_block else: # TODO: handle latex, etc. pass - return {"content": doc_structure}, df_list + return {"content": doc_structure}, asset_accumulator.rows() def convert_doc2dics( @@ -991,7 +630,7 @@ def convert_doc2dics( ] pure_text = "\n".join(text_items).strip() know_id = gen_str_codes(pure_text) - # Use relative_root for path instead of absolute kb_dir + # Use relative_root for path instead of the absolute output directory. path_suffix = key if key.strip() else "" know_path = ( split_char.join([relative_root, path_suffix]) @@ -999,19 +638,16 @@ def convert_doc2dics( else (relative_root or path_suffix) ) df_list.append( - [ - bottom_content, - know_path, - match_type, - len(bottom_content), - keywords, - summary, - know_id, - bottom_tokens, - "", - time_stamp, - "", - ] + ParsedRow( + content=bottom_content, + path=know_path, + type=match_type, + keywords=keywords, + summary=summary, + know_id=know_id, + tokens=bottom_tokens, + addtime=time_stamp, + ).to_list() ) except KnowhereException: raise @@ -1024,6 +660,23 @@ def convert_doc2dics( original_exception=e, ) - doc_df = pd.DataFrame(df_list, columns=settings.ALL_DF_COLS.split(",")) + rows_builder = ParsedRowsBuilder() + for row_values in df_list: + rows_builder.append( + ParsedRow( + content=str(row_values[0]), + path=str(row_values[1]), + type=str(row_values[2]), + length=int(row_values[3]), + keywords=str(row_values[4]), + summary=str(row_values[5]), + know_id=str(row_values[6]), + tokens=str(row_values[7]), + connectto=str(row_values[8]), + addtime=str(row_values[9]), + page_nums=str(row_values[10]), + ) + ) + doc_df = rows_builder.to_dataframe() doc_df = process_dup_paths_df(doc_df) return doc_df diff --git a/apps/worker/app/services/document_parser/formats/docx/table_html.py b/apps/worker/app/services/document_parser/formats/docx/table_html.py new file mode 100644 index 000000000..6c60f46fa --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/docx/table_html.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +from typing import Any + +from docx.table import Table as DocxTable + + +def table2html(table: DocxTable, cell_image_map: dict | None = None) -> str: + """Convert a DOCX table to HTML with colspan, rowspan, and nested tables.""" + + namespace = {"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main"} + + def get_cell_vmerge(cell: Any) -> str | None: + tc = cell._tc + tc_pr = tc.find(".//w:tcPr", namespaces=namespace) + if tc_pr is None: + return None + + v_merge = tc_pr.find(".//w:vMerge", namespaces=namespace) + if v_merge is None: + return None + + val = v_merge.get( + "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}val" + ) + return val if val else "continue" + + row_count = len(table.rows) + if row_count == 0: + return "
" + + grid = [] + for row in table.rows: + row_data = [] + previous_tc_id = None + for cell in row.cells: + tc_id = id(cell._tc) + is_new = tc_id != previous_tc_id + row_data.append((tc_id, cell, is_new)) + previous_tc_id = tc_id + grid.append(row_data) + + column_count = max(len(row) for row in grid) if grid else 0 + colspan_grid = [[0] * column_count for _ in range(row_count)] + + for row_idx in range(row_count): + row_len = len(grid[row_idx]) + col_idx = 0 + while col_idx < row_len: + tc_id = grid[row_idx][col_idx][0] + span = 1 + while ( + col_idx + span < row_len + and grid[row_idx][col_idx + span][0] == tc_id + ): + span += 1 + colspan_grid[row_idx][col_idx] = span + col_idx += span + + rowspan_grid = [[1] * column_count for _ in range(row_count)] + + for col_idx in range(column_count): + row_idx = 0 + while row_idx < row_count: + if col_idx >= len(grid[row_idx]): + row_idx += 1 + continue + + cell = grid[row_idx][col_idx][1] + vmerge = get_cell_vmerge(cell) + + if vmerge == "restart": + span = 1 + while row_idx + span < row_count: + if col_idx >= len(grid[row_idx + span]): + break + next_cell = grid[row_idx + span][col_idx][1] + next_vmerge = get_cell_vmerge(next_cell) + if next_vmerge == "continue": + span += 1 + else: + break + rowspan_grid[row_idx][col_idx] = span + row_idx += span + elif vmerge == "continue": + rowspan_grid[row_idx][col_idx] = 0 + row_idx += 1 + else: + row_idx += 1 + + html_parts = [""] + + for row_idx in range(row_count): + html_parts.append("") + col_idx = 0 + unique_col_idx = 0 + + while col_idx < len(grid[row_idx]): + _, cell, is_new = grid[row_idx][col_idx] + + if not is_new: + col_idx += 1 + continue + + rowspan = rowspan_grid[row_idx][col_idx] + if rowspan == 0: + unique_col_idx += 1 + col_idx += 1 + continue + + colspan = colspan_grid[row_idx][col_idx] + + if cell.tables: + content = "".join(table2html(nested_table) for nested_table in cell.tables) + else: + content = cell.text.strip().replace("\n", "
") + + if cell_image_map: + image_description = cell_image_map.get((row_idx, unique_col_idx)) + if image_description: + content += f"
{image_description}" + + attrs = [] + if colspan > 1: + attrs.append(f'colspan="{colspan}"') + if rowspan > 1: + attrs.append(f'rowspan="{rowspan}"') + + attr_str = " " + " ".join(attrs) if attrs else "" + html_parts.append(f"{content}") + + unique_col_idx += 1 + col_idx += colspan + + html_parts.append("") + + html_parts.append("
") + return "".join(html_parts) diff --git a/apps/worker/app/services/document_parser/formats/docx/toc.py b/apps/worker/app/services/document_parser/formats/docx/toc.py new file mode 100644 index 000000000..deb77ecfd --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/docx/toc.py @@ -0,0 +1,291 @@ +from __future__ import annotations + +import re + +from app.services.document_parser.structure.heading_candidates import ( + judge_by_conditions, + remove_by_conditions, +) +from app.services.document_parser.structure.toc_hierarchy import build_toc_hierarchy_payload +import lxml.etree as etree + +TOC_TITLE_KEYWORDS = {"目录", "目次", "contents", "table of contents"} + + +def parse_w_int_attr(elem, ns, attr_names): + if elem is None: + return None + + for attr_name in attr_names: + raw_val = elem.get("{%s}%s" % (ns["w"], attr_name)) + if raw_val is None: + continue + try: + return int(raw_val) + except (TypeError, ValueError): + continue + return None + + +def get_docx_toc_layout_hints(elem, ns): + ppr = elem.find("./w:pPr", namespaces=ns) + if ppr is None: + ppr = elem.find(".//w:pPr", namespaces=ns) + + if ppr is None: + return { + "outline_level": None, + "left_indent": None, + } + + outline_elem = ppr.find("./w:outlineLvl", namespaces=ns) + outline_level = None + if outline_elem is not None: + outline_val = parse_w_int_attr(outline_elem, ns, ["val"]) + if outline_val is not None: + outline_level = outline_val + 1 + + indent_elem = ppr.find("./w:ind", namespaces=ns) + left_indent = parse_w_int_attr( + indent_elem, ns, ["left", "start", "leftChars", "startChars"] + ) + + return { + "outline_level": outline_level, + "left_indent": left_indent, + } + + +def infer_toc_level_from_text(text: str): + text_clean = str(text).strip() + if not text_clean: + return None + + normalized = re.sub(r"\s+", " ", text_clean).lower() + if normalized in TOC_TITLE_KEYWORDS: + return None + + raw_pos_code = judge_by_conditions(text_clean) + if not isinstance(raw_pos_code, list): + return None + positive_codes = [ + int(value) + for value in raw_pos_code + if isinstance(value, int) and value > 0 + ] + neg_code = remove_by_conditions(text_clean) + if any(value > 0 for value in neg_code) or not positive_codes: + return None + + return max(positive_codes) + + +def is_toc_title_text(text: str) -> bool: + normalized = re.sub(r"\s+", " ", str(text).strip()).lower() + return normalized in TOC_TITLE_KEYWORDS + + +def infer_toc_levels_from_indentation(entries: list) -> None: + indent_values = sorted( + { + entry["left_indent"] + for entry in entries + if entry.get("level") is None + and entry.get("left_indent") is not None + and not is_toc_title_text(entry.get("heading", "")) + } + ) + + if not indent_values: + return + + indent_to_level = {indent: idx + 1 for idx, indent in enumerate(indent_values)} + for entry in entries: + if entry.get("level") is not None: + continue + if is_toc_title_text(entry.get("heading", "")): + continue + left_indent = entry.get("left_indent") + if left_indent is None: + continue + entry["level"] = indent_to_level.get(left_indent) + + +def get_docx_toc_style_info(elem, ns): + style = elem.find(".//w:pPr/w:pStyle", namespaces=ns) + if style is None: + return { + "is_toc_style": False, + "toc_level": None, + "style_name": None, + } + + val = style.get("{%s}val" % ns["w"]) + if not val: + return { + "is_toc_style": False, + "toc_level": None, + "style_name": None, + } + + val_lower = val.lower().strip() + if "toc" not in val_lower and "目录" not in val: + return { + "is_toc_style": False, + "toc_level": None, + "style_name": val, + } + + level = None + match = re.search(r"(?:toc|目录)\s*[_-]?(\d+)$", val_lower) + if match: + level = int(match.group(1)) + + layout_hints = get_docx_toc_layout_hints(elem, ns) + if level is None: + level = layout_hints["outline_level"] + + return { + "is_toc_style": True, + "toc_level": level, + "style_name": val, + "outline_level": layout_hints["outline_level"], + "left_indent": layout_hints["left_indent"], + } + + +def get_toc_level(elem, ns): + style_info = get_docx_toc_style_info(elem, ns) + if not style_info["is_toc_style"]: + return False + + if style_info["toc_level"] is not None: + return style_info["toc_level"] + return True + + +def detect_sdt_toc(elem, ns): + tag = etree.QName(elem.tag).localname if isinstance(elem.tag, str) else None + + if tag != "sdt": + return {"is_toc_sdt": False, "gallery_type": None} + + is_toc_sdt = False + gallery_type = None + + sdt_pr = elem.find(".//w:sdtPr", namespaces=ns) + if sdt_pr is not None: + doc_part_obj = sdt_pr.find(".//w:docPartObj", namespaces=ns) + if doc_part_obj is not None: + doc_part_gallery = doc_part_obj.find(".//w:docPartGallery", namespaces=ns) + if doc_part_gallery is not None: + gallery_type = doc_part_gallery.get("{%s}val" % ns["w"]) + if gallery_type and "table of contents" in gallery_type.lower(): + is_toc_sdt = True + + return {"is_toc_sdt": is_toc_sdt, "gallery_type": gallery_type} + + +def detect_doc_tocs(elem, ns): + style_info = get_docx_toc_style_info(elem, ns) + is_style = style_info["is_toc_style"] + is_field_start = False + + instrs = elem.findall(".//w:instrText", namespaces=ns) + for instr in instrs: + if instr.text: + instr_text_stripped = instr.text.strip() + instr_text_lower = instr_text_stripped.lower() + if ( + instr_text_lower.startswith("toc") + or "table of contents" in instr_text_lower + or "目录" in instr_text_stripped + ): + is_field_start = True + break + + is_field_end = False + fldchars = elem.findall(".//w:fldChar", namespaces=ns) + for fld in fldchars: + if fld.get("{%s}fldCharType" % ns["w"]) == "end": + is_field_end = True + break + + return { + "is_style": is_style, + "toc_level": style_info["toc_level"], + "style_name": style_info["style_name"], + "outline_level": style_info.get("outline_level"), + "left_indent": style_info.get("left_indent"), + "is_field_start": is_field_start, + "is_field_end": is_field_end, + } + + +def build_docx_toc_hierarchies(block_tuples: list) -> list: + toc_areas = [] + current_area = [] + + for ele_num, block, label, meta in block_tuples: + if "TOC" in label: + current_area.append((ele_num, block, meta or {})) + continue + + if current_area: + toc_areas.append(current_area) + current_area = [] + + if current_area: + toc_areas.append(current_area) + + toc_hierarchies = [] + for area in toc_areas: + toc_entries = [] + for ele_num, block, meta in area: + toc_level = meta.get("toc_level") + try: + toc_level = int(toc_level) if toc_level is not None else None + except (TypeError, ValueError): + toc_level = None + + text = getattr(block, "text", str(block)).strip() + if not text: + continue + + if toc_level is None: + outline_level = meta.get("toc_outline_level") + try: + toc_level = ( + int(outline_level) if outline_level is not None else None + ) + except (TypeError, ValueError): + toc_level = None + + if toc_level is None: + toc_level = infer_toc_level_from_text(text) + + left_indent = meta.get("toc_left_indent") + try: + left_indent = int(left_indent) if left_indent is not None else None + except (TypeError, ValueError): + left_indent = None + + toc_entries.append( + { + "id": ele_num, + "heading": text, + "level": toc_level if toc_level and toc_level > 0 else None, + "left_indent": left_indent, + } + ) + + infer_toc_levels_from_indentation(toc_entries) + payload = build_toc_hierarchy_payload( + toc_entries, + toc_range=(area[0][0], area[-1][0]), + scan_range=(area[0][0], area[-1][0]), + ) + if payload: + toc_hierarchies.append(payload) + + return toc_hierarchies diff --git a/apps/worker/app/services/document_parser/formats/excel/structure_parser.py b/apps/worker/app/services/document_parser/formats/excel/structure_parser.py new file mode 100644 index 000000000..bfe3553de --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/excel/structure_parser.py @@ -0,0 +1,743 @@ +from __future__ import annotations + +import datetime +import io +from typing import List, Optional, Tuple, Union +from typing import cast + +import openpyxl +import pandas as pd +from pandas._typing import Axes +from loguru import logger + +from shared.core.exceptions.domain_exceptions import TableParsingException + + +def parse_excel_structure( + file_source: Union[str, io.BytesIO], + sheet_name: Optional[str] = None, + split_subtables: bool = True, + include_hidden_sheets: bool = False, +) -> dict[str, pd.DataFrame]: + try: + if isinstance(file_source, str): + workbook = openpyxl.load_workbook(file_source, data_only=True) + else: + file_source.seek(0) + workbook = openpyxl.load_workbook(file_source, data_only=True) + + results: dict[str, pd.DataFrame] = {} + sheets_to_parse = [sheet_name] if sheet_name else workbook.sheetnames + + for selected_sheet_name in sheets_to_parse: + if selected_sheet_name not in workbook.sheetnames: + logger.warning( + f"Sheet '{selected_sheet_name}' not found in workbook, skipping" + ) + continue + + worksheet = workbook[selected_sheet_name] + + if not include_hidden_sheets and worksheet.sheet_state != "visible": + logger.info( + f"Sheet '{selected_sheet_name}' is hidden " + f"(state={worksheet.sheet_state}), skipping" + ) + continue + + if worksheet.max_row is None or worksheet.max_row == 0: + logger.debug(f"Sheet '{selected_sheet_name}' is empty, skipping") + continue + + merged_ranges = list(worksheet.merged_cells.ranges) + logger.debug( + f"Sheet '{selected_sheet_name}': found {len(merged_ranges)} merged cell ranges" + ) + + if split_subtables: + subtable_regions = _split_sheet_recursive( + worksheet, + (1, worksheet.max_row), + (1, worksheet.max_column or 1), + merged_ranges, + ) + before_count = len(subtable_regions) + subtable_regions = _merge_small_subtables(worksheet, subtable_regions) + if len(subtable_regions) != before_count: + logger.info( + f"Sheet '{selected_sheet_name}': merged {before_count} subtables → " + f"{len(subtable_regions)} " + f"(absorbed {before_count - len(subtable_regions)} small fragments)" + ) + logger.debug( + f"Sheet '{selected_sheet_name}': {len(subtable_regions)} subtables after merge" + ) + + for index, (row_range, col_range) in enumerate(subtable_regions): + result = _parse_subtable( + worksheet, + row_range, + col_range, + merged_ranges, + ) + dataframe = result["df"] + dataframe.attrs["row_header_cols"] = len(result["header_cols"]) + key = selected_sheet_name if index == 0 else f"{selected_sheet_name}_{index + 1}" + logger.debug( + f"Subtable '{key}': rows={row_range}, cols={col_range}, " + f"header_rows={result['header_rows']}, header_cols={result['header_cols']}" + ) + results[key] = dataframe + else: + row_range = (1, worksheet.max_row) + col_range = (1, worksheet.max_column or 1) + result = _parse_subtable(worksheet, row_range, col_range, merged_ranges) + dataframe = result["df"] + dataframe.attrs["row_header_cols"] = len(result["header_cols"]) + logger.debug( + f"Sheet '{selected_sheet_name}': header_rows={result['header_rows']}, " + f"header_cols={result['header_cols']}, " + f"fallback_col={result['fallback_col_header']}, " + f"fallback_row={result['fallback_row_header']}" + ) + results[selected_sheet_name] = dataframe + + workbook.close() + return results + except Exception as exc: + logger.error(f"Error parsing Excel with precision mode: {exc}") + raise TableParsingException( + user_message="Failed to parse Excel file headers", + reason="EXCEL_PRECISION_PARSE_FAILED", + internal_message=str(exc), + original_exception=exc, + ) from exc + + +DATA_TYPES_TO_EXCLUDE = (int, float, datetime.datetime) + + +def _get_merged_cell_value(ws, row: int, col: int, merged_ranges: list): + for merged_range in merged_ranges: + if ( + merged_range.min_row <= row <= merged_range.max_row + and merged_range.min_col <= col <= merged_range.max_col + ): + return ws.cell(merged_range.min_row, merged_range.min_col).value + return ws.cell(row, col).value + + +def _get_unique_cells_in_row( + ws, + row: int, + col_range: Tuple[int, int], + merged_ranges: list, +) -> List[dict]: + col_start, col_end = col_range + cells = [] + visited_cols = set() + + for col in range(col_start, col_end + 1): + if col in visited_cols: + continue + + in_merge = False + for merged_range in merged_ranges: + if ( + merged_range.min_row <= row <= merged_range.max_row + and merged_range.min_col <= col <= merged_range.max_col + ): + value = ws.cell(merged_range.min_row, merged_range.min_col).value + merge_col_end = min(merged_range.max_col, col_end) + + for merged_col in range(merged_range.min_col, merge_col_end + 1): + visited_cols.add(merged_col) + + cells.append( + { + "col_start": merged_range.min_col, + "col_end": merge_col_end, + "value": value, + "is_merged": True, + } + ) + in_merge = True + break + + if not in_merge: + value = ws.cell(row, col).value + cells.append( + {"col_start": col, "col_end": col, "value": value, "is_merged": False} + ) + visited_cols.add(col) + + return cells + + +def _get_unique_cells_in_col( + ws, + col: int, + row_range: Tuple[int, int], + merged_ranges: list, +) -> List[dict]: + row_start, row_end = row_range + cells = [] + visited_rows = set() + + for row in range(row_start, row_end + 1): + if row in visited_rows: + continue + + in_merge = False + for merged_range in merged_ranges: + if ( + merged_range.min_row <= row <= merged_range.max_row + and merged_range.min_col <= col <= merged_range.max_col + ): + value = ws.cell(merged_range.min_row, merged_range.min_col).value + merge_row_end = min(merged_range.max_row, row_end) + + for merged_row in range(merged_range.min_row, merge_row_end + 1): + visited_rows.add(merged_row) + + cells.append( + { + "row_start": merged_range.min_row, + "row_end": merge_row_end, + "value": value, + "is_merged": True, + } + ) + in_merge = True + break + + if not in_merge: + value = ws.cell(row, col).value + cells.append( + {"row_start": row, "row_end": row, "value": value, "is_merged": False} + ) + visited_rows.add(row) + + return cells + + +def _is_candidate_header_row( + ws, + row: int, + col_range: Tuple[int, int], + merged_ranges: list, + exclude_types: tuple = DATA_TYPES_TO_EXCLUDE, +) -> bool: + cells = _get_unique_cells_in_row(ws, row, col_range, merged_ranges) + + has_any_value = False + for cell in cells: + value = cell["value"] + if value is None: + continue + has_any_value = True + + if isinstance(value, bool): + continue + if isinstance(value, exclude_types): + return False + + return has_any_value + + +def _is_candidate_header_col( + ws, + col: int, + row_range: Tuple[int, int], + merged_ranges: list, + exclude_types: tuple = DATA_TYPES_TO_EXCLUDE, +) -> bool: + cells = _get_unique_cells_in_col(ws, col, row_range, merged_ranges) + + has_any_value = False + for cell in cells: + value = cell["value"] + if value is None: + continue + has_any_value = True + + if isinstance(value, bool): + continue + if isinstance(value, exclude_types): + return False + + return has_any_value + + +def _detect_header_regions( + ws, + row_range: Tuple[int, int], + col_range: Tuple[int, int], + merged_ranges: list, +) -> Tuple[List[int], List[int]]: + row_start, row_end = row_range + col_start, col_end = col_range + + header_rows = [] + for row in range(row_start, row_end + 1): + if _is_candidate_header_row(ws, row, col_range, merged_ranges): + header_rows.append(row) + else: + break + + data_row_start = header_rows[-1] + 1 if header_rows else row_start + if data_row_start > row_end: + return header_rows, [] + + header_cols = [] + data_row_range = (data_row_start, row_end) + for col in range(col_start, col_end + 1): + if _is_candidate_header_col(ws, col, data_row_range, merged_ranges): + header_cols.append(col) + else: + break + + return header_rows, header_cols + + +def _build_column_multiindex( + ws, + header_rows: List[int], + col_range: Tuple[int, int], + merged_ranges: list, +) -> Union[pd.Index, pd.MultiIndex]: + col_start, col_end = col_range + levels = [] + + for row in header_rows: + row_values = [] + for col in range(col_start, col_end + 1): + value = _get_merged_cell_value(ws, row, col, merged_ranges) + row_values.append(str(value).strip() if value else "") + levels.append(row_values) + + for index, level in enumerate(levels): + filled = [] + last = "" + for value in level: + if value: + last = value + filled.append(last if last else value) + levels[index] = filled + + if len(levels) == 1: + return pd.Index(levels[0]) + return pd.MultiIndex.from_arrays(levels) + + +def _build_row_multiindex( + ws, + header_cols: List[int], + row_range: Tuple[int, int], + merged_ranges: list, + header_rows: List[int] | None = None, +) -> Union[pd.Index, pd.MultiIndex]: + row_start, row_end = row_range + levels = [] + names = [] + + for col in header_cols: + col_values = [] + for row in range(row_start, row_end + 1): + value = _get_merged_cell_value(ws, row, col, merged_ranges) + col_values.append(str(value).strip() if value else "") + levels.append(col_values) + + if header_rows: + name_row = header_rows[-1] + name_value = _get_merged_cell_value(ws, name_row, col, merged_ranges) + names.append(str(name_value).strip() if name_value else None) + else: + names.append(None) + + for index, level in enumerate(levels): + filled = [] + last = "" + for value in level: + if value: + last = value + filled.append(last if last else value) + levels[index] = filled + + if len(levels) == 1: + row_index = pd.Index(levels[0]) + row_index.name = names[0] if names else None + return row_index + return pd.MultiIndex.from_arrays(levels, names=names) + + +def _parse_subtable( + ws, + row_range: Tuple[int, int], + col_range: Tuple[int, int], + merged_ranges: list, +) -> dict: + row_start, row_end = row_range + col_start, col_end = col_range + + header_rows, header_cols = _detect_header_regions( + ws, row_range, col_range, merged_ranges + ) + + total_rows = row_end - row_start + 1 + total_cols = col_end - col_start + 1 + fallback_col_header = len(header_rows) == total_rows + fallback_row_header = len(header_cols) == total_cols + + if fallback_col_header: + data_row_start = row_start + columns = None + else: + data_row_start = header_rows[-1] + 1 if header_rows else row_start + columns = ( + _build_column_multiindex(ws, header_rows, col_range, merged_ranges) + if header_rows + else None + ) + + if fallback_row_header: + data_col_start = col_start + row_index = None + else: + data_col_start = header_cols[-1] + 1 if header_cols else col_start + row_index = ( + _build_row_multiindex( + ws, header_cols, (data_row_start, row_end), merged_ranges, header_rows + ) + if header_cols + else None + ) + + data = [] + for row in range(data_row_start, row_end + 1): + row_data = [] + for col in range(data_col_start, col_end + 1): + value = _get_merged_cell_value(ws, row, col, merged_ranges) + row_data.append(value) + data.append(row_data) + + if columns is not None and header_cols and not fallback_row_header: + columns = cast(Axes, columns[len(header_cols) :]) + + dataframe = pd.DataFrame(data, columns=columns, index=row_index) + + excel_row_numbers = list(range(data_row_start, row_end + 1)) + if isinstance(dataframe.columns, pd.MultiIndex): + level_count = dataframe.columns.nlevels + src_row_key = tuple(["_src_row"] + [""] * (level_count - 1)) + dataframe[src_row_key] = excel_row_numbers + else: + dataframe["_src_row"] = excel_row_numbers + + return { + "df": dataframe, + "header_rows": header_rows if not fallback_col_header else [], + "header_cols": header_cols if not fallback_row_header else [], + "fallback_col_header": fallback_col_header, + "fallback_row_header": fallback_row_header, + } + + +def _find_effective_range( + ws, + row_range: Tuple[int, int], + col_range: Tuple[int, int], +) -> Tuple[Tuple[int, int], Tuple[int, int]]: + row_start, row_end = row_range + col_start, col_end = col_range + + effective_row_start = None + effective_row_end = None + effective_col_start = None + effective_col_end = None + + for row in range(row_start, row_end + 1): + for col in range(col_start, col_end + 1): + if ws.cell(row, col).value is not None: + if effective_row_start is None: + effective_row_start = row + effective_row_end = row + if effective_col_start is None or col < effective_col_start: + effective_col_start = col + if effective_col_end is None or col > effective_col_end: + effective_col_end = col + + if effective_row_start is None: + return ((row_start, row_start), (col_start, col_start)) + + if ( + effective_row_end is None + or effective_col_start is None + or effective_col_end is None + ): + return ((row_start, row_start), (col_start, col_start)) + + return ( + (effective_row_start, effective_row_end), + (effective_col_start, effective_col_end), + ) + + +def _is_true_separator_row( + ws, + row: int, + effective_col_range: Tuple[int, int], + merged_ranges: list | None = None, +) -> bool: + col_start, col_end = effective_col_range + merged_ranges = merged_ranges or [] + + for col in range(col_start, col_end + 1): + if ws.cell(row, col).value is not None: + return False + for merged_range in merged_ranges: + if ( + merged_range.min_row <= row <= merged_range.max_row + and merged_range.min_col <= col <= merged_range.max_col + ): + return False + return True + + +def _is_true_separator_col( + ws, + col: int, + effective_row_range: Tuple[int, int], + merged_ranges: list | None = None, +) -> bool: + row_start, row_end = effective_row_range + merged_ranges = merged_ranges or [] + + for row in range(row_start, row_end + 1): + if ws.cell(row, col).value is not None: + return False + for merged_range in merged_ranges: + if ( + merged_range.min_row <= row <= merged_range.max_row + and merged_range.min_col <= col <= merged_range.max_col + ): + return False + return True + + +def _find_separator_groups(items: List[int]) -> List[List[int]]: + if not items: + return [] + + groups = [] + current_group = [items[0]] + + for index in range(1, len(items)): + if items[index] == items[index - 1] + 1: + current_group.append(items[index]) + else: + groups.append(current_group) + current_group = [items[index]] + + groups.append(current_group) + return groups + + +def _split_sheet_recursive( + ws, + row_range: Tuple[int, int], + col_range: Tuple[int, int], + merged_ranges: list | None = None, + min_rows: int = 2, + min_cols: int = 2, +) -> List[Tuple[Tuple[int, int], Tuple[int, int]]]: + merged_ranges = merged_ranges or [] + + (effective_row_start, effective_row_end), ( + effective_col_start, + effective_col_end, + ) = _find_effective_range(ws, row_range, col_range) + + if ( + effective_row_end - effective_row_start + 1 < min_rows + or effective_col_end - effective_col_start + 1 < min_cols + ): + if effective_row_start is not None: + return [ + ( + (effective_row_start, effective_row_end), + (effective_col_start, effective_col_end), + ) + ] + return [] + + separator_rows = [] + for row in range(effective_row_start + 1, effective_row_end): + if _is_true_separator_row( + ws, row, (effective_col_start, effective_col_end), merged_ranges + ): + separator_rows.append(row) + + separator_cols = [] + for col in range(effective_col_start + 1, effective_col_end): + if _is_true_separator_col( + ws, col, (effective_row_start, effective_row_end), merged_ranges + ): + separator_cols.append(col) + + row_groups = _find_separator_groups(separator_rows) + col_groups = _find_separator_groups(separator_cols) + + should_split_rows = len(row_groups) > 0 and ( + len(col_groups) == 0 or len(row_groups) <= len(col_groups) + ) + should_split_cols = len(col_groups) > 0 and not should_split_rows + + if should_split_rows: + subtables = [] + previous_end = effective_row_start + for group in row_groups: + if group[0] > previous_end: + sub_result = _split_sheet_recursive( + ws, + (previous_end, group[0] - 1), + (effective_col_start, effective_col_end), + merged_ranges, + min_rows, + min_cols, + ) + subtables.extend(sub_result) + previous_end = group[-1] + 1 + if previous_end <= effective_row_end: + sub_result = _split_sheet_recursive( + ws, + (previous_end, effective_row_end), + (effective_col_start, effective_col_end), + merged_ranges, + min_rows, + min_cols, + ) + subtables.extend(sub_result) + return subtables + + if should_split_cols: + subtables = [] + previous_end = effective_col_start + for group in col_groups: + if group[0] > previous_end: + sub_result = _split_sheet_recursive( + ws, + (effective_row_start, effective_row_end), + (previous_end, group[0] - 1), + merged_ranges, + min_rows, + min_cols, + ) + subtables.extend(sub_result) + previous_end = group[-1] + 1 + if previous_end <= effective_col_end: + sub_result = _split_sheet_recursive( + ws, + (effective_row_start, effective_row_end), + (previous_end, effective_col_end), + merged_ranges, + min_rows, + min_cols, + ) + subtables.extend(sub_result) + return subtables + + return [((effective_row_start, effective_row_end), (effective_col_start, effective_col_end))] + + +def _count_non_empty_cells( + ws, + row_range: Tuple[int, int], + col_range: Tuple[int, int], +) -> int: + count = 0 + for row in range(row_range[0], row_range[1] + 1): + for col in range(col_range[0], col_range[1] + 1): + if ws.cell(row, col).value is not None: + count += 1 + return count + + +def _merge_small_subtables( + ws, + subtables: List[Tuple[Tuple[int, int], Tuple[int, int]]], + min_cells: int = 4, +) -> List[Tuple[Tuple[int, int], Tuple[int, int]]]: + if len(subtables) <= 1: + return subtables + + items = [] + for row_range, col_range in subtables: + count = _count_non_empty_cells(ws, row_range, col_range) + items.append({"rr": row_range, "cr": col_range, "cells": count}) + + changed = True + while changed and len(items) > 1: + changed = False + + min_index = None + for index, item in enumerate(items): + if item["cells"] < min_cells: + if min_index is None or item["cells"] < items[min_index]["cells"]: + min_index = index + + if min_index is None: + break + + source = items[min_index] + best_index = None + best_distance = float("inf") + for index, target in enumerate(items): + if index == min_index: + continue + row_gap = max( + 0, + target["rr"][0] - source["rr"][1] - 1, + source["rr"][0] - target["rr"][1] - 1, + ) + col_gap = max( + 0, + target["cr"][0] - source["cr"][1] - 1, + source["cr"][0] - target["cr"][1] - 1, + ) + distance = row_gap + col_gap + if distance < best_distance or ( + best_index is not None + and distance == best_distance + and target["cells"] > items[best_index]["cells"] + ): + best_distance = distance + best_index = index + + if best_index is None: + break + + target = items[best_index] + merged_row_range = ( + min(source["rr"][0], target["rr"][0]), + max(source["rr"][1], target["rr"][1]), + ) + merged_col_range = ( + min(source["cr"][0], target["cr"][0]), + max(source["cr"][1], target["cr"][1]), + ) + items[best_index] = { + "rr": merged_row_range, + "cr": merged_col_range, + "cells": source["cells"] + target["cells"], + } + + logger.debug( + f"Merged small fragment (rows={source['rr']}, cols={source['cr']}, " + f"cells={source['cells']}) into neighbor (rows={target['rr']}, cols={target['cr']})" + ) + + del items[min_index] + changed = True + + return [(item["rr"], item["cr"]) for item in items] diff --git a/apps/worker/app/services/document_parser/formats/excel/table_parser.py b/apps/worker/app/services/document_parser/formats/excel/table_parser.py new file mode 100644 index 000000000..3f5499ea3 --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/excel/table_parser.py @@ -0,0 +1,286 @@ +# pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportOptionalSubscript=false, reportReturnType=false +from __future__ import annotations + +import io +import os +from dataclasses import dataclass +from typing import Any + +import pandas as pd +from app.services.document_parser.tables.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.formats.excel.structure_parser import parse_excel_structure +from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.support.parser_rows import ParsedRow, ParsedRowsBuilder +from app.services.document_parser.support.path_helpers import remove_spaces +from app.services.document_parser.tables.table_asset_writer import ( + TableAssetInput, + write_table_asset, +) +from app.services.document_parser.tables.table_frame_parser import ( + parse_headers, + parse_tb_contents, + parse_tb_keywords, + postprocess_tb, +) +from bs4 import BeautifulSoup +from loguru import logger + +from shared.core.exceptions.domain_exceptions import TableParsingException +from shared.core.exceptions.knowhere_exception import KnowhereException +from shared.utils.chunk_refs import build_chunk_ref +from app.services.common.file_loading import load_file_bytes +from app.services.common.file_utils import path_handle +from shared.utils.text_utils import tokenize2stw_remove + + +@dataclass(frozen=True) +class ExcelWorkbookParseRequest: + file_path: str + file_name: str + output_dir: str + baseurl: str + base_llm_paras: dict[str, Any] + window_h: int + relative_root: str | None + use_precision_mode: bool + include_hidden_sheets: bool + + +def parse_xlsx( + file_path: str, + file_name: str, + output_dir: str, + baseurl: str, + base_llm_paras: dict[str, Any] | None = None, + window_h: int = 10, + relative_root: str | None = None, + use_precision_mode: bool = True, + include_hidden_sheets: bool = False, +) -> pd.DataFrame: + request = ExcelWorkbookParseRequest( + file_path=file_path, + file_name=file_name, + output_dir=output_dir, + baseurl=baseurl, + base_llm_paras=_normalise_llm_parameters(base_llm_paras), + window_h=window_h, + relative_root=relative_root, + use_precision_mode=use_precision_mode, + include_hidden_sheets=include_hidden_sheets, + ) + return parse_excel_workbook(request) + + +def parse_excel_workbook(request: ExcelWorkbookParseRequest) -> pd.DataFrame: + time_stamp = get_str_time() + sheets_dict, precision_mode_active = _load_excel_sheets(request) + parsed_rows: list[ParsedRow] = [] + + for sheet_name, sheet_frame in _iter_unique_sheets(sheets_dict): + table_rows = _parse_excel_sheet( + request=request, + sheet_name=sheet_name, + sheet_frame=sheet_frame, + precision_mode_active=precision_mode_active, + time_stamp=time_stamp, + ) + parsed_rows.extend(table_rows) + + return _rows_to_dataframe(parsed_rows) + + +def _normalise_llm_parameters( + base_llm_paras: dict[str, Any] | None, +) -> dict[str, Any]: + llm_parameters = dict(base_llm_paras or {}) + llm_parameters.setdefault("summary_table", False) + llm_parameters.setdefault("stopwords", []) + return llm_parameters + + +def _load_excel_sheets( + request: ExcelWorkbookParseRequest, +) -> tuple[dict[str, pd.DataFrame], bool]: + table_data = load_file_bytes(request.file_path, file_url=request.baseurl) + table_stream = io.BytesIO(table_data) + + os.makedirs(os.path.join(request.output_dir, "tables"), exist_ok=True) + + if not request.use_precision_mode: + return pd.read_excel(table_stream, sheet_name=None), False + + logger.info("Using precision mode for Excel header detection") + try: + return ( + parse_excel_structure( + table_stream, + include_hidden_sheets=request.include_hidden_sheets, + ), + True, + ) + except Exception as exc: + logger.warning(f"Precision mode failed, falling back to legacy mode: {exc}") + table_stream.seek(0) + return pd.read_excel(table_stream, sheet_name=None), False + + +def _iter_unique_sheets( + sheets_dict: dict[str, pd.DataFrame], +) -> list[tuple[str, pd.DataFrame]]: + used_sheet_names: list[str] = [] + unique_sheets: list[tuple[str, pd.DataFrame]] = [] + + for raw_sheet_name, sheet_content in sheets_dict.items(): + sheet_name = raw_sheet_name.strip() + if sheet_name in used_sheet_names: + sheet_name = sheet_name + str(len(used_sheet_names)) + else: + used_sheet_names.append(sheet_name) + unique_sheets.append((sheet_name, sheet_content)) + + return unique_sheets + + +def _parse_excel_sheet( + *, + request: ExcelWorkbookParseRequest, + sheet_name: str, + sheet_frame: pd.DataFrame, + precision_mode_active: bool, + time_stamp: str, +) -> list[ParsedRow]: + parsed_rows: list[ParsedRow] = [] + + try: + table_frame = postprocess_tb(sheet_frame, drop=True) + if len(table_frame) == 0 or table_frame.empty or table_frame.isna().all().all(): + return parsed_rows + + if not precision_mode_active: + table_frame = parse_headers(table_frame, paras=request.base_llm_paras) + + table_frame = _drop_source_row_columns(table_frame) + row_header_cols = int(table_frame.attrs.get("row_header_cols", 0)) + + _table_paths, table_html = parse_tb_contents( + table_frame, + parent_dic={request.file_name: {sheet_name: {}}}, + file_name=request.file_name, + sheet_name=sheet_name, + row_header_cols=row_header_cols, + ) + + parsed_rows.append( + _write_excel_table_asset( + request=request, + sheet_name=sheet_name, + table_frame=table_frame, + table_html=table_html, + time_stamp=time_stamp, + ) + ) + return parsed_rows + except KnowhereException: + raise + except Exception as exc: + logger.error(f"Table parsing failed: {exc}") + raise TableParsingException( + user_message="Failed to parse Excel table content", + reason="TABLE_PROCESSING_FAILED", + internal_message=str(exc), + original_exception=exc, + ) from exc + + +def _drop_source_row_columns(table_frame: pd.DataFrame) -> pd.DataFrame: + source_row_columns = [ + column + for column in table_frame.columns + if (isinstance(column, tuple) and column[0] == "_src_row") + or column == "_src_row" + ] + if not source_row_columns: + return table_frame + return table_frame.drop(columns=source_row_columns) + + +def _write_excel_table_asset( + *, + request: ExcelWorkbookParseRequest, + sheet_name: str, + table_frame: pd.DataFrame, + table_html: str, + time_stamp: str, +) -> ParsedRow: + title, keywords, summary = _summarize_excel_table( + table_frame=table_frame, + table_html=table_html, + sheet_name=sheet_name, + llm_parameters=request.base_llm_paras, + ) + table_index = f"table-{sheet_name}" + table_summary = f"{table_index}\n{summary}" if summary else table_index + effective_name = title or sheet_name + table_stem = path_handle( + remove_spaces("table-" + effective_name), + mode="clean_single", + ) + if not isinstance(table_stem, str) or not table_stem: + raise ValueError(f"Failed to sanitize Excel table name: {effective_name}") + table_name = table_stem + ".html" + table_html_string = BeautifulSoup(table_html, features="html.parser").prettify() + know_id = gen_str_codes(table_html + str(sheet_name)) + table_ref = build_chunk_ref(f"tables/{table_name}") + table_content = ( + f"{table_ref}\nTable summary:\n{table_summary}\nMain columns:\n{keywords}" + ) + table_tokens = tokenize2stw_remove( + [table_content], + request.base_llm_paras["stopwords"], + ) + + return write_table_asset( + TableAssetInput( + html=table_html_string, + output_dir=request.output_dir, + table_name=table_name, + summary=table_summary, + keywords=keywords, + know_id=know_id, + addtime=time_stamp, + content=table_content, + tokens=table_tokens, + length=len(table_html), + ) + ) + + +def _summarize_excel_table( + *, + table_frame: pd.DataFrame, + table_html: str, + sheet_name: str, + llm_parameters: dict[str, Any], +) -> tuple[str | None, str, str | None]: + mechanical_keywords = parse_tb_keywords(table_frame) + + if llm_parameters["summary_table"]: + from app.services.document_parser.formats.text.parser import ( + extract_title_keywords_summary, + ) + + title, keywords, summary = extract_title_keywords_summary( + table_html, + max_keywords=3, + ) + return title, keywords or mechanical_keywords, summary + + return None, mechanical_keywords, None + + +def _rows_to_dataframe(parsed_rows: list[ParsedRow]) -> pd.DataFrame: + rows_builder = ParsedRowsBuilder() + for row in parsed_rows: + rows_builder.append(row) + table_df = rows_builder.to_dataframe() + return process_dup_paths_df(table_df) diff --git a/apps/worker/app/services/document_parser/fragment_parser.py b/apps/worker/app/services/document_parser/formats/fragment/parser.py similarity index 77% rename from apps/worker/app/services/document_parser/fragment_parser.py rename to apps/worker/app/services/document_parser/formats/fragment/parser.py index 5aef9f004..66a20069e 100644 --- a/apps/worker/app/services/document_parser/fragment_parser.py +++ b/apps/worker/app/services/document_parser/formats/fragment/parser.py @@ -2,19 +2,21 @@ Fragment Parser - for user-pasted text content injection This parser handles .fragment files which represent user-pasted text content -that needs to be injected into the knowledge base without requiring a physical file. +that needs to be parsed without requiring a physical file. """ import os from typing import Any, Optional -from app.services.document_parser.md_parser import parse_md +from app.services.document_parser.formats.markdown.parser import parse_md +from app.services.document_parser.orchestration.path_segment import ( + build_parser_path_segment, +) from loguru import logger from openai.types.chat import ChatCompletionMessageParam -from shared.core.config import settings -from shared.utils.file_utils import path_handle -from shared.utils.OpenAICompatibleClientSync import get_openai_client +from app.services.common.file_utils import path_handle +from shared.services.ai.openai_compatible_client_sync import get_openai_client def generate_fragment_title(content: str, max_tokens: int = 30) -> Optional[str]: @@ -54,7 +56,6 @@ def parse_fragment( fragment_content: str, filename: Optional[str] = None, output_dir: Optional[str] = None, - kb_dir: str = "Default_Root", base_llm_paras: Optional[dict[str, Any]] = None, **kwargs, ): @@ -65,13 +66,11 @@ def parse_fragment( fragment_content: The text content to parse filename: Optional filename, if not provided will be auto-generated output_dir: Base output directory - kb_dir: Knowledge base directory name base_llm_paras: LLM parameters for parsing Returns: tuple: (full_output_dir, relative_root, parsed_df) """ - split_char = settings.SPLIT_CHAR or "/" if output_dir is None: raise ValueError("output_dir is required for fragment parsing") @@ -91,13 +90,22 @@ def parse_fragment( logger.debug(f"Fragment filename: {filename}") # Build relative_root and full_output_dir - kb_dir_parts = kb_dir.split(split_char) - relative_root = "/".join(kb_dir_parts + [filename]) - full_output_dir = os.path.join(output_dir, relative_root.replace("/", os.sep)) - sanitized_output_dir = path_handle(full_output_dir, mode="sanitize") - if not isinstance(sanitized_output_dir, str): - raise ValueError("sanitized fragment output path must be a string") - full_output_dir = sanitized_output_dir + filename_segment = build_parser_path_segment( + filename, + default=f"fragment_{os.urandom(4).hex()}.fragment", + ) + relative_root = filename_segment + full_output_dir = os.path.realpath( + os.path.join(output_dir, filename_segment) + ) + resolved_output_dir = os.path.realpath(output_dir) + if ( + os.path.commonpath([resolved_output_dir, full_output_dir]) + != resolved_output_dir + ): + raise ValueError( + f"Fragment output directory escaped workspace: {full_output_dir}" + ) os.makedirs(full_output_dir, exist_ok=True) logger.debug(f"Fragment relative_root: {relative_root}") diff --git a/apps/worker/app/services/document_parser/formats/html/parser.py b/apps/worker/app/services/document_parser/formats/html/parser.py new file mode 100644 index 000000000..bc39731b2 --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/html/parser.py @@ -0,0 +1,448 @@ +# pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportReturnType=false +""" +HTML table parsing and extraction utilities. + +This module provides functions for converting pandas DataFrames to HTML tables +with support for: +- HTML to DataFrame/Markdown conversion +- HTML header expansion +- Nested HTML table parsing +""" + +from typing import Dict, List, Optional + +import pandas as pd +from bs4 import BeautifulSoup + +from shared.core.exceptions.domain_exceptions import TableParsingException +from shared.utils.text_utils import remove_duplicates_orderkept + + +class HTMLHeaderExpander: + """ + Expands multi-level HTML table headers into a flat list of column names. + + Handles: + - rowspan: cells that span multiple rows + - colspan: cells that span multiple columns + - Multi-row headers: combines parent -> child relationships with " > " + + Example: + Row 0: | A(rs=2) | B(cs=2) | + Row 1: | | B1 | B2 | + + Output: ["A", "B > B1", "B > B2"] + + Ported from snap-fill's HeaderMatrixExpander, using BeautifulSoup instead of PyQuery. + """ + + def __init__(self, table_html: str): + """Initialize with table HTML content.""" + soup = BeautifulSoup(table_html, "html.parser") + table = soup.find("table") + if table: + self.rows = table.find_all("tr", recursive=False) + # Also check inside thead/tbody + if not self.rows: + self.rows = table.find_all("tr") + else: + self.rows = [] + + def expand_headers( + self, header_row_count: int = 2, start_row: int = 0 + ) -> List[str]: + """ + Expand multi-row headers into a flat column list. + + Args: + header_row_count: Number of rows that form the header (default 2) + start_row: Starting row index for headers (default 0) + + Returns: + List of column header strings, with nested headers joined by " > " + """ + grid = self._build_grid(header_row_count, start_row) + if not grid or not grid[0]: + return [] + + max_cols = len(grid[0]) + + # Build column headers by combining rows vertically + headers = [] + for col in range(max_cols): + parts = [] + prev_text = None + for row in range(len(grid)): + text = grid[row][col] + # Only add if non-empty and different from previous + if text and text != prev_text: + parts.append(text) + prev_text = text + + if parts: + headers.append(" > ".join(parts)) + else: + headers.append("") # Empty column + + return headers + + def _build_grid( + self, header_row_count: int, start_row: int = 0 + ) -> List[List[Optional[str]]]: + """ + Build expanded grid from table rows. + + Algorithm: + 1. First pass: calculate max columns by expanding all colspan + 2. Create occupied grid to track cells filled by rowspan/colspan + 3. For each cell, place text in all positions it spans + """ + if start_row + header_row_count > len(self.rows): + return [] + + header_rows = self.rows[start_row : start_row + header_row_count] + + # First pass: calculate max columns (sum of colspan for each row) + max_cols = 0 + for tr in header_rows: + col_count = 0 + for td in tr.find_all(["td", "th"], recursive=False): + col_count += int(td.get("colspan", 1)) + max_cols = max(max_cols, col_count) + + if max_cols == 0: + return [] + + # Create grid and occupied tracker + grid = [[None] * max_cols for _ in range(header_row_count)] + occupied = [[False] * max_cols for _ in range(header_row_count)] + + # Process each row + for row_idx, tr in enumerate(header_rows): + col_ptr = 0 # Current column position in this row + + for td in tr.find_all(["td", "th"], recursive=False): + text = td.get_text(strip=True) + rowspan = int(td.get("rowspan", 1)) + colspan = int(td.get("colspan", 1)) + + # Skip past any already-occupied cells (from previous row's rowspan) + while col_ptr < max_cols and occupied[row_idx][col_ptr]: + col_ptr += 1 + + if col_ptr >= max_cols: + break + + # Fill all cells covered by this rowspan/colspan + for r in range(row_idx, min(row_idx + rowspan, header_row_count)): + for c in range(col_ptr, min(col_ptr + colspan, max_cols)): + grid[r][c] = text + occupied[r][c] = True + + # Move column pointer past this cell + col_ptr += colspan + + return grid + + def get_grid_debug( + self, header_row_count: int = 2, start_row: int = 0 + ) -> List[List[Optional[str]]]: + """Return the expanded grid for debugging and testing purposes.""" + return self._build_grid(header_row_count, start_row) + + def get_unique_headers( + self, header_row_count: int = 2, start_row: int = 0 + ) -> List[str]: + """ + Get deduplicated headers for use in keywords/LLM prompts. + Removes duplicate column names (from colspan expansion) while preserving order. + """ + headers = self.expand_headers(header_row_count, start_row) + seen = set() + unique = [] + for h in headers: + if h and h not in seen: + seen.add(h) + unique.append(h) + return unique + + def detect_header_row_count( + self, start_row: int = 0, max_scan_rows: int = 5 + ) -> int: + """ + Automatically detect header row count using the maximum rowspan. + + Logic: + - scan at most ``max_scan_rows`` rows starting at ``start_row`` + - find the largest rowspan value among all cells + - ``rowspan=2`` means a two-row header, ``rowspan=3`` means three rows + + Returns: + int: Detected header row count, at least 1. + """ + if start_row >= len(self.rows): + return 1 + + max_rowspan = 1 + scan_end = min(start_row + max_scan_rows, len(self.rows)) + + for row in self.rows[start_row:scan_end]: + for cell in row.find_all(["td", "th"], recursive=False): + rowspan = int(cell.get("rowspan", 1)) + if rowspan > max_rowspan: + max_rowspan = rowspan + + return max_rowspan + + def detect_row_indices( + self, + header_row_count: int, + start_row: int = 0, + end_row: int = None, + max_scan_cols: int = 3, + ) -> Dict: + """ + Detect row-index columns and return the row-index metadata. + + Logic: + 1. start from data rows after the header + 2. scan left to right; a non-empty, non-numeric column is a row index + 3. stop once a data column is reached + + Returns: + dict: { + 'row_index_col_count': int, + 'row_index_col_name': str or None, + 'row_indices': List[str], + } + """ + result = { + "row_index_col_count": 0, + "row_index_col_name": None, + "row_indices": [], + } + + data_start_row = start_row + header_row_count + if data_start_row >= len(self.rows): + return result + + # Build header grid for column names + grid = self._build_grid(header_row_count, start_row) + if not grid or not grid[0]: + return result + + # Data rows + if end_row is not None: + data_end_row = min(end_row + 1, len(self.rows)) + else: + data_end_row = len(self.rows) + + data_rows = self.rows[data_start_row:data_end_row] + if not data_rows: + return result + + row_index_col_count = 0 + row_index_col_names = [] + + for cell_idx in range(max_scan_cols): + col_values = [] + is_index_col = True + + for row in data_rows: + cells = row.find_all(["td", "th"], recursive=False) + + if cell_idx >= len(cells): + continue + + cell_value = cells[cell_idx].get_text(strip=True) + + if not cell_value: + is_index_col = False + break + + # Check if purely numeric + if ( + cell_value.replace(".", "") + .replace("-", "") + .replace(" ", "") + .isdigit() + ): + is_index_col = False + break + + col_values.append(cell_value) + + if is_index_col and col_values: + row_index_col_count += 1 + row_index_col_names.append(col_values) + else: + break # Stop at first data column + + if row_index_col_count > 0: + result["row_index_col_count"] = row_index_col_count + result["row_index_col_name"] = grid[-1][0] if grid and grid[-1] else None + result["row_indices"] = row_index_col_names[0] + + return result + + +def parse_nested_htmltb(table): + """Parse nested HTML table into a list of rows. + + Handles nested tables recursively. + + Args: + table: BeautifulSoup table element + + Returns: + List of rows, where each row is a list of cell values or nested tables + """ + rows = [] + for tr in table.find_all("tr", recursive=False): + row = [] + for td in tr.find_all(["td", "th"], recursive=False): + inner_table = td.find("table") + if inner_table: + row.append(parse_nested_htmltb(inner_table)) + else: + row.append(td.get_text(strip=True)) + if row: + rows.append(row) + return rows + + +def html_to_md_lines(html: str): + """Convert HTML table to list of markdown-like lines. + + Args: + html: HTML string containing a table + + Returns: + List of strings, each representing a row with cells separated by ' | ' + """ + soup = BeautifulSoup(html, "html.parser") + lines = [] + for row in soup.find_all("tr"): + row_text = [] + for cell in row.find_all("td", recursive=False): + text = cell.get_text(separator=" ", strip=True) + if text: + row_text.append(text) + if row_text: + lines.append(" | ".join(row_text)) + lines = remove_duplicates_orderkept(lines) + return lines + + +def tb_htmlstr_to_df(html_str): + """Convert first table in HTML string to DataFrame""" + soup = BeautifulSoup(html_str, "html.parser") + table = soup.find("table") + if not table: + raise TableParsingException( + user_message="No table structure found in the document", + reason="INVALID_FORMAT", + internal_message="No found in the HTML string", + ) + nested_list = parse_nested_htmltb(table) + try: + df = pd.DataFrame(nested_list[1:], columns=nested_list[0]) + except Exception: + df = pd.DataFrame(nested_list) + return df + + +def merge_html_tables(lines: list) -> list: + """Merge multi-line HTML tables into single lines. + + MinerU and some parsers output HTML tables with line breaks inside. + This function joins all lines from
to
into a single line. + + Args: + lines: List of text lines (already stripped) + + Returns: + List of lines with HTML tables merged into single lines + """ + merged_lines = [] + in_table = False + table_buffer = [] + + for line in lines: + if "" in line: + merged_lines.append(" ".join(table_buffer)) + table_buffer = [] + in_table = False + elif in_table: + table_buffer.append(line) + if "" in line: + merged_lines.append(" ".join(table_buffer)) + table_buffer = [] + in_table = False + else: + merged_lines.append(line) + + # Handle unclosed table at end of file + if table_buffer: + merged_lines.append(" ".join(table_buffer)) + + return merged_lines + + +def first_cols_rows_html(html_str, max_items=10, max_chars=20): + """Extract first row and first column from HTML table string. + + This function mirrors the logic of _first_cols_rows in doc_parser.py + but works with HTML strings instead of DOCX Table objects. + + Args: + html_str: HTML table string + max_items: Maximum number of items to extract (default 10) + max_chars: Maximum characters per item (default 20) + + Returns: + Tuple of (first_row_text, first_col_text) with ' | ' as separator + """ + from app.services.document_parser.support.text_helpers import truncate_text + + soup = BeautifulSoup(html_str, "html.parser") + table = soup.find("table") + if not table: + return "", "" + + rows = table.find_all("tr") + if not rows: + return "", "" + + # First row extraction (deduplicated, order preserved, max items, truncated) + first_row_cells = rows[0].find_all(["td", "th"]) + seen_row = set() + unique_row = [] + for cell in first_row_cells: + if len(unique_row) >= max_items: + break + text = cell.get_text(strip=True) + if text and text not in seen_row: + seen_row.add(text) + unique_row.append(truncate_text(text, max_chars, 0)) + first_row_text = " | ".join(unique_row) if unique_row else "" + + # First column extraction (deduplicated, order preserved, max items, truncated) + seen_col = set() + unique_col = [] + for row in rows: + if len(unique_col) >= max_items: + break + cells = row.find_all(["td", "th"]) + if cells: + text = cells[0].get_text(strip=True) + if text and text not in seen_col: + seen_col.add(text) + unique_col.append(truncate_text(text, max_chars, 0)) + first_col_text = " | ".join(unique_col) if unique_col else "" + + return first_row_text, first_col_text diff --git a/apps/worker/app/services/document_parser/image_parser.py b/apps/worker/app/services/document_parser/formats/image/parser.py similarity index 87% rename from apps/worker/app/services/document_parser/image_parser.py rename to apps/worker/app/services/document_parser/formats/image/parser.py index dc8ffb541..9dd3808ac 100755 --- a/apps/worker/app/services/document_parser/image_parser.py +++ b/apps/worker/app/services/document_parser/formats/image/parser.py @@ -8,11 +8,9 @@ from pathlib import Path import pandas as pd -from app.services.common.kb_utils import ( - gen_str_codes, - get_str_time, - process_dup_paths_df, -) +from app.services.document_parser.tables.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.support.parser_rows import ParsedRow, ParsedRowsBuilder from loguru import logger from PIL import Image @@ -25,9 +23,9 @@ from shared.services.ai.prompt_service import build_prompt from shared.services.ai.response_process_service import eval_response from shared.utils.chunk_refs import build_chunk_ref -from shared.utils.CommonHelperSync import is_remote, load_file_bytes -from shared.utils.file_utils import path_handle -from shared.utils.OpenAICompatibleClientSync import ( +from app.services.common.file_loading import is_remote, load_file_bytes +from app.services.common.file_utils import path_handle +from shared.services.ai.openai_compatible_client_sync import ( OpenAICompatibleClientSync, get_openai_client, ) @@ -99,12 +97,12 @@ def local_image_to_data_url(path, cut=True, min_size=None, max_size=None): return img_data_base64 -def process_img_path4read(paths_, kb_dir, cut): +def process_img_path4read(paths_, image_root_dir, cut): urls = [] for path_ in paths_: if not is_remote(path_): - kb_dir = Path(kb_dir).resolve() - url_ = local_image_to_data_url(kb_dir / path_, cut) + resolved_image_root = Path(image_root_dir).resolve() + url_ = local_image_to_data_url(resolved_image_root / path_, cut) if url_ is not None: urls.append(url_) else: @@ -114,7 +112,7 @@ def process_img_path4read(paths_, kb_dir, cut): def ask_image( client: OpenAICompatibleClientSync, - kb_dir, + image_root_dir, paths_, title_text="", task="summary-images", @@ -140,7 +138,7 @@ def ask_image( if not valid_paths: return None - urls_ = process_img_path4read(valid_paths, kb_dir, size_cut) + urls_ = process_img_path4read(valid_paths, image_root_dir, size_cut) if task in ("summary-images", "atlas-page-info"): image_model = settings.IMAGE_MODEL or "gpt-4-vision-preview" @@ -189,16 +187,16 @@ def ask_image( return None -def detect_summary_img_md(line, last_context, kb_dir, mode=False): +def detect_summary_img_md(line, last_context, image_root_dir, mode=False): client = _get_vision_client() imgs = [] img_paths = re.findall(MD_IMAGE_PATTERN, line, flags=re.IGNORECASE) for i, ip in enumerate(img_paths): if mode: try: - llm_resp = ask_image(client, kb_dir, paths_=[ip]) + llm_resp = ask_image(client, image_root_dir, paths_=[ip]) if llm_resp: - from app.services.document_parser.txt_parser import ( + from app.services.document_parser.formats.text.parser import ( split_title_summary, ) @@ -227,7 +225,6 @@ def parse_image( relative_root=None, ): split_char = settings.SPLIT_CHAR or "/" - df_list = [] time_stamp = get_str_time() os.makedirs(output_dir, exist_ok=True) img_dir = os.path.join(output_dir, "images") @@ -299,7 +296,7 @@ def parse_image( size_cut=False, ) if llm_resp: - from app.services.document_parser.txt_parser import split_title_summary + from app.services.document_parser.formats.text.parser import split_title_summary img_title, image_summary = split_title_summary(llm_resp) else: @@ -308,7 +305,7 @@ def parse_image( else: # For non-text images, split title from summary-images response if base_llm_paras["summary_image"] and image_content != filename: - from app.services.document_parser.txt_parser import split_title_summary + from app.services.document_parser.formats.text.parser import split_title_summary img_title, image_summary = split_title_summary(image_content) else: @@ -322,20 +319,18 @@ def parse_image( if inferred_suffix.lower() in {".png", ".jpg", ".jpeg", ".gif", ".webp"}: img_name = img_stem or img_name if auto_rename: - update_img_path = os.path.join(img_dir, f"{img_name}{img_suffix}") + target_img_path = os.path.join(img_dir, f"{img_name}{img_suffix}") if os.path.exists(img_path): - if img_path != update_img_path: - os.rename(img_path, update_img_path) + if img_path != target_img_path: + os.rename(img_path, target_img_path) # Store the relative filename for path construction final_img_name = f"{img_name}{img_suffix}" else: logger.warning( f"Image file missing before rename, keeping original name: {filename}" ) - update_img_path = img_path final_img_name = filename else: - update_img_path = img_path final_img_name = filename except KnowhereException: raise @@ -358,23 +353,19 @@ def parse_image( ) img_ref = build_chunk_ref(relative_img_path) img_bottom_content = f"{img_ref}\nImage Content:\n{image_content}" - df_list.append( - [ - img_bottom_content, - relative_img_path, - "image", - len(img_bottom_content), - "", - image_summary, - temp_uid, - "", - "", - time_stamp, - "", - ] + rows_builder = ParsedRowsBuilder() + rows_builder.append( + ParsedRow( + content=img_bottom_content, + path=relative_img_path, + type="image", + summary=image_summary, + know_id=temp_uid, + addtime=time_stamp, + ) ) - img_df = pd.DataFrame(df_list, columns=settings.ALL_DF_COLS.split(",")) + img_df = rows_builder.to_dataframe() img_df = process_dup_paths_df(img_df) return img_df diff --git a/apps/worker/app/services/document_parser/formats/markdown/deferred_summary.py b/apps/worker/app/services/document_parser/formats/markdown/deferred_summary.py new file mode 100644 index 000000000..394ed06cc --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/markdown/deferred_summary.py @@ -0,0 +1,308 @@ +from __future__ import annotations + +import os +import re +from dataclasses import dataclass +from typing import Literal, TypeGuard + +import gevent +from app.services.document_parser.formats.markdown.deferred_task import ( + ImageDeferredSummaryTask, + MarkdownDeferredSummaryTask, + TableDeferredSummaryTask, + TextDeferredSummaryTask, +) +from app.services.document_parser.formats.image.parser import _get_vision_client, ask_image +from app.services.document_parser.support.stage_profiler import stage_timer +from app.services.document_parser.tables.table_text_parser import sanitize_table_name_from_header +from app.services.document_parser.formats.text.parser import ( + extract_title_keywords_summary, + split_title_summary, +) +from gevent.pool import Pool as GeventPool +from loguru import logger + +from shared.core.config import settings +from shared.utils.chunk_refs import build_chunk_ref +from app.services.common.file_utils import path_handle + +DeferredResult = ( + tuple[ + int, + Literal["image", "table", "text"], + tuple[str | None, str | None] | tuple[str, str, str] | tuple[str, str], + ] +) +ImageSummaryResult = tuple[str | None, str | None] +TableSummaryResult = tuple[str, str, str] +TextSummaryResult = tuple[str, str] + + +@dataclass(frozen=True) +class MarkdownDeferredSummaryInput: + rows: list[list[str | int]] + tasks: list[MarkdownDeferredSummaryTask] + output_dir: str + summary_len: int = 1500 + + +def apply_markdown_deferred_summaries( + deferred_input: MarkdownDeferredSummaryInput, +) -> None: + if not deferred_input.tasks: + return + + image_task_count = sum( + 1 for task in deferred_input.tasks if isinstance(task, ImageDeferredSummaryTask) + ) + table_task_count = sum( + 1 for task in deferred_input.tasks if isinstance(task, TableDeferredSummaryTask) + ) + text_task_count = sum( + 1 for task in deferred_input.tasks if isinstance(task, TextDeferredSummaryTask) + ) + logger.info( + f"Running {len(deferred_input.tasks)} deferred summary LLM calls in parallel" + ) + max_concurrent = getattr(settings, "SUMMARY_LLM_MAX_CONCURRENT", 8) + + with stage_timer( + "md.deferred_summaries", + total_tasks=len(deferred_input.tasks), + image_tasks=image_task_count, + table_tasks=table_task_count, + text_tasks=text_task_count, + max_concurrent=min(max_concurrent, len(deferred_input.tasks)), + ): + results = _run_deferred_summary_tasks(deferred_input, max_concurrent) + _apply_deferred_summary_results(deferred_input, results) + + logger.info(f"Completed {len(deferred_input.tasks)} deferred summary LLM calls") + + +def replace_chunk_ref_in_rows( + rows: list[list[str | int]], old_path: str, new_path: str +) -> None: + old_ref = build_chunk_ref(old_path) + new_ref = build_chunk_ref(new_path) + if not old_ref or old_ref == new_ref: + return + + for row in rows: + if len(row) > 0 and isinstance(row[0], str): + row[0] = row[0].replace(old_ref, new_ref) + if len(row) > 1 and row[1] == old_path: + row[1] = new_path + if len(row) > 2 and isinstance(row[2], str): + row[2] = row[2].replace(old_ref, new_ref) + if len(row) > 8 and isinstance(row[8], str): + row[8] = row[8].replace(old_ref, new_ref) + + +def _run_deferred_summary_tasks( + deferred_input: MarkdownDeferredSummaryInput, + max_concurrent: int, +) -> list[DeferredResult | None]: + pool = GeventPool(size=min(max_concurrent, len(deferred_input.tasks))) + greenlets = [ + pool.spawn(_run_deferred_summary_task, task, deferred_input) + for task in deferred_input.tasks + ] + gevent.joinall(greenlets) + return [greenlet.value for greenlet in greenlets] + + +def _run_deferred_summary_task( + task: MarkdownDeferredSummaryTask, + deferred_input: MarkdownDeferredSummaryInput, +) -> DeferredResult | None: + try: + if isinstance(task, ImageDeferredSummaryTask): + client = _get_vision_client() + # TODO: Risk of missing text content if MinerU outputted a pure text image. + # Consider adding judge-image-type and OCR fallback as done in image_parser.parse_image. + llm_resp = ask_image( + client, deferred_input.output_dir, paths_=[task.relative_path] + ) + if llm_resp: + img_title, img_summary = split_title_summary(llm_resp) + else: + img_title, img_summary = None, None + return task.row_index, "image", (img_title, img_summary) + + if isinstance(task, TableDeferredSummaryTask): + title, keywords, summary = extract_title_keywords_summary( + task.table_html, max_keywords=3 + ) + return task.row_index, "table", (title, keywords, summary) + + if isinstance(task, TextDeferredSummaryTask): + _, keywords, summary = extract_title_keywords_summary( + task.content, + max_keywords=3, + summary_len=deferred_input.summary_len, + ) + return task.row_index, "text", (keywords, summary) + except Exception as exc: + logger.warning( + f"Deferred summary LLM call failed for idx={task.row_index}: {exc}" + ) + return None + + logger.warning(f"Unknown deferred markdown summary task type: {type(task).__name__}") + return None + + +def _apply_deferred_summary_results( + deferred_input: MarkdownDeferredSummaryInput, + results: list[DeferredResult | None], +) -> None: + deferred_by_index = {task.row_index: task for task in deferred_input.tasks} + + for result in results: + if result is None: + continue + + row_index, task_type, task_result = result + if task_type == "image": + if not _is_image_summary_result(task_result): + logger.warning(f"Invalid image deferred result for idx={row_index}") + continue + _apply_image_summary_result( + deferred_input.rows, + _get_image_task(deferred_by_index[row_index]), + row_index, + task_result, + ) + elif task_type == "table": + if not _is_table_summary_result(task_result): + logger.warning(f"Invalid table deferred result for idx={row_index}") + continue + _apply_table_summary_result( + deferred_input.rows, + _get_table_task(deferred_by_index[row_index]), + row_index, + task_result, + ) + elif task_type == "text": + if not _is_text_summary_result(task_result): + logger.warning(f"Invalid text deferred result for idx={row_index}") + continue + _apply_text_summary_result(deferred_input.rows, row_index, task_result) + + +def _is_image_summary_result(result: object) -> TypeGuard[ImageSummaryResult]: + return ( + isinstance(result, tuple) + and len(result) == 2 + and all(isinstance(value, (str, type(None))) for value in result) + ) + + +def _is_table_summary_result(result: object) -> TypeGuard[TableSummaryResult]: + return ( + isinstance(result, tuple) + and len(result) == 3 + and all(isinstance(value, str) for value in result) + ) + + +def _is_text_summary_result(result: object) -> TypeGuard[TextSummaryResult]: + return ( + isinstance(result, tuple) + and len(result) == 2 + and all(isinstance(value, str) for value in result) + ) + + +def _get_image_task(task: MarkdownDeferredSummaryTask) -> ImageDeferredSummaryTask: + if isinstance(task, ImageDeferredSummaryTask): + return task + raise TypeError(f"Expected image deferred task, got {type(task).__name__}") + + +def _get_table_task(task: MarkdownDeferredSummaryTask) -> TableDeferredSummaryTask: + if isinstance(task, TableDeferredSummaryTask): + return task + raise TypeError(f"Expected table deferred task, got {type(task).__name__}") + + +def _apply_image_summary_result( + rows: list[list[str | int]], + original_task: ImageDeferredSummaryTask, + row_index: int, + result: ImageSummaryResult, +) -> None: + img_title, img_summary = result + row = rows[row_index] + if img_summary: + image_index = str(row[5]).split("\n")[0] if row[5] else "image" + row[5] = f"{image_index}\n{img_summary}" + + if not img_title: + return + + image_dir = original_task.image_dir + old_img_name = original_task.image_name + image_suffix = original_task.image_suffix + safe_title = path_handle(str(img_title), mode="clean_single") + img_num_match = re.match(r"image-(\d+)", str(old_img_name)) + img_num = ( + img_num_match.group(1) + if img_num_match + else str(old_img_name).split("-")[1] + if "-" in str(old_img_name) + else "0" + ) + new_img_name = path_handle(f"image-{img_num}-{safe_title}", mode="clean_single") + old_path = os.path.join(image_dir, f"{old_img_name}{image_suffix}") + new_path = os.path.join(image_dir, f"{new_img_name}{image_suffix}") + if old_path == new_path or not os.path.exists(old_path): + return + + os.rename(old_path, new_path) + new_relative_path = f"images/{new_img_name}{image_suffix}" + replace_chunk_ref_in_rows(rows, str(row[1]), new_relative_path) + row[1] = new_relative_path + + +def _apply_table_summary_result( + rows: list[list[str | int]], + original_task: TableDeferredSummaryTask, + row_index: int, + result: TableSummaryResult, +) -> None: + title, keywords, summary = result + row = rows[row_index] + row[4] = keywords if isinstance(keywords, str) else "" + if summary: + table_index = str(row[5]) if "\n" not in str(row[5]) else str(row[5]).split("\n")[0] + row[5] = f"{table_index}\n{summary}" + + if not title: + return + + table_dir = original_task.table_dir + old_table_name = original_task.table_name + table_count = original_task.table_count + safe_title = sanitize_table_name_from_header(str(title)) + new_table_name = path_handle( + f"table-{table_count} {safe_title}", mode="clean_single" + ) + old_path = os.path.join(table_dir, f"{old_table_name}.html") + new_path = os.path.join(table_dir, f"{new_table_name}.html") + if old_path == new_path or not os.path.exists(old_path): + return + + os.rename(old_path, new_path) + new_relative_path = f"tables/{new_table_name}.html" + replace_chunk_ref_in_rows(rows, str(row[1]), new_relative_path) + row[1] = new_relative_path + + +def _apply_text_summary_result( + rows: list[list[str | int]], row_index: int, result: TextSummaryResult +) -> None: + keywords, summary = result + rows[row_index][4] = keywords if isinstance(keywords, str) else "" + rows[row_index][5] = summary if isinstance(summary, str) else "" diff --git a/apps/worker/app/services/document_parser/formats/markdown/deferred_task.py b/apps/worker/app/services/document_parser/formats/markdown/deferred_task.py new file mode 100644 index 000000000..f2cb82409 --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/markdown/deferred_task.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TypeAlias + + +@dataclass(frozen=True) +class ImageDeferredSummaryTask: + row_index: int + relative_path: str + image_dir: str + image_name: str + image_suffix: str + + +@dataclass(frozen=True) +class TableDeferredSummaryTask: + row_index: int + table_html: str + table_dir: str + table_name: str + table_count: int + + +@dataclass(frozen=True) +class TextDeferredSummaryTask: + row_index: int + content: str + + +MarkdownDeferredSummaryTask: TypeAlias = ( + ImageDeferredSummaryTask | TableDeferredSummaryTask | TextDeferredSummaryTask +) + +__all__ = [ + "ImageDeferredSummaryTask", + "MarkdownDeferredSummaryTask", + "TableDeferredSummaryTask", + "TextDeferredSummaryTask", +] diff --git a/apps/worker/app/services/document_parser/formats/markdown/image_asset.py b/apps/worker/app/services/document_parser/formats/markdown/image_asset.py new file mode 100644 index 000000000..b71ace46a --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/markdown/image_asset.py @@ -0,0 +1,229 @@ +from __future__ import annotations + +import os +from dataclasses import dataclass +from pathlib import Path +from typing import cast + +from app.services.document_parser.support.identifiers import gen_str_codes +from app.services.document_parser.formats.image.parser import perceptual_hash +from app.services.document_parser.assets.inline_asset import build_image_asset_row +from app.services.document_parser.formats.markdown.deferred_task import ( + ImageDeferredSummaryTask, + MarkdownDeferredSummaryTask, +) +from app.services.document_parser.formats.markdown.parse_state import ParserRowValues +from loguru import logger + +from shared.utils.chunk_refs import build_chunk_ref +from app.services.common.file_utils import path_handle + + +@dataclass(frozen=True) +class MarkdownImageAsset: + content_item: str | None + row_values: ParserRowValues | None + cache_key: str | None + cache_entry: dict[str, str] | None + deferred_task: MarkdownDeferredSummaryTask | None + should_advance_image_count: bool + + +@dataclass(frozen=True) +class MarkdownImageAssetRequest: + output_dir: str + image_dir: str + image_path: str + image_name: str + image_count: int + last_context: str + image_summary: str | None + timestamp: str + current_page_number: int + seen_images: dict[str, dict[str, str]] + summary_image: bool + row_index: int + + +def build_markdown_image_asset( + request: MarkdownImageAssetRequest, +) -> MarkdownImageAsset: + image_suffix = os.path.splitext(request.image_path)[-1] + source_path = resolve_markdown_image_source_path( + request.output_dir, + request.image_path, + ) + if source_path is None or not source_path.exists(): + logger.warning(f"Image file not found, skipping rename: {request.image_path}") + return _empty_asset(should_advance_image_count=True) + + with open(source_path, "rb") as image_file: + image_binary_hash = perceptual_hash(image_file.read()) + + if image_binary_hash in request.seen_images: + return _build_duplicate_image_asset( + source_path=source_path, + cache_entry=request.seen_images[image_binary_hash], + timestamp=request.timestamp, + current_page_number=request.current_page_number, + ) + + relative_image_path = f"images/{request.image_name}{image_suffix}" + target_image_path = os.path.join( + request.image_dir, + f"{request.image_name}{image_suffix}", + ) + os.rename(source_path, target_image_path) + + image_index = f"image-{request.image_count}" + effective_summary = request.image_summary or request.last_context or None + image_summary_field = ( + f"{image_index}\n{effective_summary}" if effective_summary else image_index + ) + image_content = _build_image_content( + relative_image_path=relative_image_path, + summary=effective_summary, + ) + image_know_id = gen_str_codes(image_binary_hash) + row_values = _build_image_row_values( + content=image_content, + relative_path=relative_image_path, + summary=image_summary_field, + know_id=image_know_id, + timestamp=request.timestamp, + current_page_number=request.current_page_number, + ) + cache_entry = { + "relative_img_path": relative_image_path, + "img_content": image_content, + "img_summary_field": image_summary_field, + "temp_uid": image_know_id, + } + + deferred_task = None + if request.summary_image: + deferred_task = ImageDeferredSummaryTask( + row_index=request.row_index, + relative_path=relative_image_path, + image_dir=request.image_dir, + image_name=request.image_name, + image_suffix=image_suffix, + ) + + return MarkdownImageAsset( + content_item=image_content, + row_values=row_values, + cache_key=image_binary_hash, + cache_entry=cache_entry, + deferred_task=deferred_task, + should_advance_image_count=True, + ) + + +def build_markdown_image_name(*, image_count: int, last_context: str) -> str: + image_name_context = path_handle(last_context[:10], mode="clean_single") + return f"image-{str(image_count)}-{image_name_context}" + + +def resolve_workspace_image_path( + candidate_path: Path, workspace_path: Path, +) -> Path | None: + """Return the candidate only when it exists inside the current job workspace.""" + resolved_path = candidate_path.resolve(strict=False) + try: + resolved_path.relative_to(workspace_path) + except ValueError: + return None + return resolved_path if resolved_path.exists() else None + + +def resolve_markdown_image_source_path(output_dir: str, image_path: str) -> Path | None: + """Handle local absolute refs and container cwd-relative refs safely.""" + if not image_path: + return None + + workspace_path = Path(output_dir).resolve() + raw_path = Path(image_path).expanduser() + candidate_paths = ( + [raw_path] + if raw_path.is_absolute() + else [ + workspace_path / raw_path, + Path.cwd() / raw_path, + ] + ) + + for candidate_path in candidate_paths: + resolved_path = resolve_workspace_image_path(candidate_path, workspace_path) + if resolved_path is not None: + return resolved_path + + return None + + +def _build_duplicate_image_asset( + *, + source_path: Path, + cache_entry: dict[str, str], + timestamp: str, + current_page_number: int, +) -> MarkdownImageAsset: + row_values = _build_image_row_values( + content=cache_entry["img_content"], + relative_path=cache_entry["relative_img_path"], + summary=cache_entry["img_summary_field"], + know_id=cache_entry["temp_uid"], + timestamp=timestamp, + current_page_number=current_page_number, + ) + try: + source_path.unlink() + except OSError as exc: + logger.debug(f"Failed to remove duplicate image source {source_path}: {exc}") + logger.debug("Skipped duplicate image") + return MarkdownImageAsset( + content_item=cache_entry["img_content"], + row_values=row_values, + cache_key=None, + cache_entry=None, + deferred_task=None, + should_advance_image_count=False, + ) + + +def _build_image_content(*, relative_image_path: str, summary: str | None) -> str: + image_reference = build_chunk_ref(relative_image_path) + if summary: + return f"\n{summary}\n{image_reference}\n" + return f"\n{image_reference}\n" + + +def _build_image_row_values( + *, + content: str, + relative_path: str, + summary: str, + know_id: str, + timestamp: str, + current_page_number: int, +) -> ParserRowValues: + image_row = build_image_asset_row( + content=content, + relative_path=relative_path, + summary=summary, + know_id=know_id, + addtime=timestamp, + page_nums=str(current_page_number) if current_page_number > 0 else "", + ) + return cast(ParserRowValues, image_row.to_list()) + + +def _empty_asset(*, should_advance_image_count: bool) -> MarkdownImageAsset: + return MarkdownImageAsset( + content_item=None, + row_values=None, + cache_key=None, + cache_entry=None, + deferred_task=None, + should_advance_image_count=should_advance_image_count, + ) diff --git a/apps/worker/app/services/document_parser/formats/markdown/parse_state.py b/apps/worker/app/services/document_parser/formats/markdown/parse_state.py new file mode 100644 index 000000000..b7c6da508 --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/markdown/parse_state.py @@ -0,0 +1,185 @@ +from __future__ import annotations + +import re +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +import pandas as pd + +from app.services.document_parser.tables.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.formats.markdown.deferred_task import ( + MarkdownDeferredSummaryTask, + TextDeferredSummaryTask, +) +from app.services.document_parser.support.parser_rows import ParsedRow, ParsedRowsBuilder + +ParserRowValues = list[str | int] + +RowUpdater = Callable[ + [list[ParserRowValues], list[str], str, dict[str, Any], str, str, int, bool], + list[ParserRowValues], +] + + +@dataclass +class MarkdownParseState: + relative_root: str + split_char: str + llm_parameters: dict[str, Any] + timestamp: str + row_updater: RowUpdater + rows: list[ParserRowValues] = field(default_factory=list) + content_items: list[str] = field(default_factory=list) + path_stack: list[tuple[str, int]] = field(default_factory=list) + inner_paths: list[str] = field(default_factory=list) + error_line_numbers: list[int] = field(default_factory=list) + table_lines: list[str] = field(default_factory=list) + current_page_number: int = 0 + chunk_pages: set[int] = field(default_factory=set) + base_level: int | None = None + path: str = "" + path_counter: dict[str, int] = field(default_factory=dict) + deferred_llm_tasks: list[MarkdownDeferredSummaryTask] = field(default_factory=list) + seen_images: dict[str, dict[str, str]] = field(default_factory=dict) + image_count: int = 1 + table_count: int = 1 + + def __post_init__(self) -> None: + if not self.path: + self.path = self.relative_root + + def record_page_marker(self, line: str) -> bool: + if "" not in line: + return False + if "page" not in line and "Slide number" not in line: + return False + + page_match = re.search(r"page\s+(\d+)", line) + if page_match: + self.current_page_number = int(page_match.group(1)) + else: + self.current_page_number += 1 + self.chunk_pages.add(self.current_page_number) + return True + + def flush_current_content(self) -> None: + page_numbers = self._format_chunk_pages() + self.rows = self.row_updater( + self.rows, + self.content_items, + self.path, + self.llm_parameters, + self.timestamp, + page_numbers, + 1500, + True, + ) + self.content_items = [] + self.chunk_pages = set() + if self.current_page_number > 0: + self.chunk_pages.add(self.current_page_number) + + def flush_placeholder_chunk(self) -> None: + page_numbers = self._format_chunk_pages() + self.rows = self.row_updater( + self.rows, + [], + self.path, + self.llm_parameters, + self.timestamp, + page_numbers, + 1500, + True, + ) + + def enter_heading(self, heading: str, level: int) -> None: + if self.base_level is None: + self.base_level = level + elif level < self.base_level: + self.base_level = level + + adjusted_level = level - self.base_level + 1 + self.path_stack = [ + (item_heading, item_level) + for item_heading, item_level in self.path_stack + if item_level < adjusted_level + ] + + current_heading = ( + heading.replace(self.split_char, "∕") + if self.split_char in heading + else heading + ) + tentative_names = [item_heading for item_heading, _ in self.path_stack] + tentative_names.append(current_heading) + tentative_path_parts = [self.relative_root] if self.relative_root else [] + tentative_path_parts.extend(tentative_names) + tentative_path = self.split_char.join(tentative_path_parts) + + if tentative_path in self.path_counter: + self.path_counter[tentative_path] += 1 + current_heading = f"{current_heading}_{self.path_counter[tentative_path]}" + else: + self.path_counter[tentative_path] = 1 + + self.path_stack.append((current_heading, adjusted_level)) + heading_names = [item_heading for item_heading, _ in self.path_stack] + path_parts = [self.relative_root] if self.relative_root else [] + path_parts.extend(heading_names) + self.inner_paths.append(self.split_char.join(heading_names)) + self.path = self.split_char.join(path_parts) + + def append_content_item(self, item: str) -> None: + self.content_items.append(item) + + def append_plain_text(self, text: str) -> None: + self.content_items.append(text.strip()) + if self.current_page_number > 0: + self.chunk_pages.add(self.current_page_number) + + def append_row(self, row: ParserRowValues) -> None: + self.rows.append(row) + + def schedule_deferred_task(self, task: MarkdownDeferredSummaryTask) -> None: + self.deferred_llm_tasks.append(task) + + def collect_text_summary_tasks(self, summary_len: int) -> None: + if not self.llm_parameters.get("summary_txt"): + return + + for index, entry in enumerate(self.rows): + marker = entry[2] + if isinstance(marker, str) and marker.strip().split("\n", 1)[0].lower() in { + "image", + "table", + }: + continue + content = str(entry[0]) + if len(content) > summary_len and not entry[4] and not entry[5]: + self.deferred_llm_tasks.append( + TextDeferredSummaryTask(row_index=index, content=content) + ) + + def to_dataframe(self) -> pd.DataFrame: + rows_builder = ParsedRowsBuilder() + for row_values in self.rows: + rows_builder.append( + ParsedRow( + content=str(row_values[0]), + path=str(row_values[1]), + type=str(row_values[2]), + length=int(row_values[3]), + keywords=str(row_values[4]), + summary=str(row_values[5]), + know_id=str(row_values[6]), + tokens=str(row_values[7]), + connectto=str(row_values[8]), + addtime=str(row_values[9]), + page_nums=str(row_values[10]), + ) + ) + return process_dup_paths_df(rows_builder.to_dataframe()) + + def _format_chunk_pages(self) -> str: + return ",".join(str(page) for page in sorted(self.chunk_pages)) diff --git a/apps/worker/app/services/document_parser/formats/markdown/parser.py b/apps/worker/app/services/document_parser/formats/markdown/parser.py new file mode 100755 index 000000000..a33be14e9 --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/markdown/parser.py @@ -0,0 +1,444 @@ +# pyright: reportArgumentType=false, reportAssignmentType=false, reportOptionalIterable=false, reportOptionalMemberAccess=false, reportOptionalOperand=false, reportOptionalSubscript=false +import json +import os +import re +import shutil + +from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.formats.markdown.deferred_summary import ( + MarkdownDeferredSummaryInput, + apply_markdown_deferred_summaries, +) +from app.services.document_parser.formats.markdown.image_asset import ( + MarkdownImageAssetRequest, + build_markdown_image_name, + build_markdown_image_asset, +) +from app.services.document_parser.formats.markdown.parse_state import MarkdownParseState +from app.services.document_parser.formats.markdown.table_asset import ( + MarkdownTableAssetRequest, + build_markdown_table_asset, +) +from app.services.document_parser.support.parser_rows import ParsedRow +from app.services.document_parser.support.path_helpers import find_matches_parsing +from app.services.document_parser.formats.html.parser import ( + merge_html_tables, +) +from app.services.document_parser.formats.image.parser import ( + MD_IMAGE_PATTERN, + detect_summary_img_md, +) +from app.services.document_parser.structure.heading_hierarchy import ( + HeadingHierarchyInput, + predict_heading_hierarchy, +) +from app.services.document_parser.structure.heading_candidates import md_heading_match +from app.services.document_parser.support.stage_profiler import stage_timer +from app.services.document_parser.tables.table_text_parser import ( + extract_tables_by_forms, + identify_tables, +) +from app.services.document_parser.structure.toc_parser import detect_tocs_in_texts +from app.services.document_parser.formats.text.parser import extract_title_keywords_summary +from loguru import logger + +from shared.core.config import settings +from shared.utils.chunk_refs import has_chunk_ref +from shared.utils.text_utils import tokenize2stw_remove + + +def find_surround_context(md_lines, lid): + def is_skip(line): + s = line.strip() + is_image = re.findall(MD_IMAGE_PATTERN, line, flags=re.IGNORECASE) + is_table, _, _ = identify_tables(line) + return not s or is_image or is_table + + n = len(md_lines) + prev_text = "" + for i in range(max(lid - 5, 0), lid): + if not is_skip(md_lines[i]): + prev_text = md_lines[i].strip() + break + + next_text = "" + for i in range(lid + 1, min(lid + 6, n)): + if not is_skip(md_lines[i]): + next_text = md_lines[i].strip() + break + return f"{prev_text} {next_text}".strip() + + +def heading_md_relocate(md_lines, heading_preds): + """Relocate markdown headings based on predicted levels (sxjg simplified logic)""" + + def remove_hash(txt): + return re.sub(r"^\s*(#+)\s*", "", txt) + + for lid, line_txt in enumerate(md_lines): + pred_level_df = heading_preds[heading_preds["id"] == lid] + + if pred_level_df.empty: # if the line does not enter predicting + line_txt = remove_hash(line_txt) + else: + pred_level = pred_level_df["level"].iloc[0] + if pred_level < 0: + line_txt = remove_hash(line_txt) + else: + # sxjg simplified: remove all #, then add correct number of # + clean_text = line_txt.lstrip("#").lstrip() + line_txt = f"{'#' * int(pred_level)} {clean_text}" + # update lines + md_lines[lid] = line_txt.strip() + + md_lines = [line for line in md_lines if line.strip() != ""] + return md_lines # note the length=original md_lines but contents/level are updated + + +def eval_md_headings( + md_lines, + source_type, + toc_hierarchies=None, + smart_parse=False, + model_name=None, + output_dir=None, + layout_json_path=None, +): + """Evaluate markdown headings with optional TOC hierarchies context""" + heading_preds = predict_heading_hierarchy( + HeadingHierarchyInput( + infos=md_lines, + doc_type=source_type, + toc_hierarchies=toc_hierarchies, + enable_regex=True, + smart_parse=smart_parse, + model_name=model_name, + output_dir=output_dir, + layout_json_path=layout_json_path, + ) + ) + + if len(heading_preds) == 0: + lines_with_heading = md_lines + else: + lines_with_heading = heading_md_relocate(md_lines, heading_preds) + return lines_with_heading + + +def clean_md_table_lines(table_lines, start_line_num): + expected_columns = table_lines[0].count("|") - 1 + cleaned_lines = [] + error_lines = [] # To record line numbers that need cleaning + + for i, line in enumerate(table_lines): + line_columns = line.count("|") - 1 + current_line_num = ( + start_line_num + i + ) # Calculate the current line number in the original file + if line_columns == expected_columns: + cleaned_lines.append(line) + else: + error_lines.append(current_line_num) + if line_columns > expected_columns: + parts = line.split("|") + cleaned_line = "|".join( + parts[: expected_columns + 1] + ) # If there are more columns, combine them (or drop extra columns) + cleaned_lines.append(cleaned_line) + elif line_columns < expected_columns: + # If there are fewer columns, pad the line (or you could skip it) + cleaned_line = line + "|" * (expected_columns - line_columns) + cleaned_lines.append(cleaned_line) + return cleaned_lines, error_lines + + +def update_df_list( + df_list, + content_items, + path, + llm_paras, + time_stamp, + page_nums="", + summary_len=1500, + skip_llm=False, +): + """Flush accumulated content_items into a chunk row in df_list. + + Args: + content_items: list of content strings. Each item is either pure text + or an IMAGE/TABLE ref block. know_id is generated from pure text + items only (deterministic), while full content includes all items. + skip_llm: if True, skip inline LLM calls (deferred to parallel batch). + """ + # Separate pure text from IMAGE/TABLE ref blocks for deterministic know_id + text_items = [item for item in content_items if not has_chunk_ref(str(item))] + pure_text = "\n".join(text_items).strip() + bottom_content = "\n".join(content_items).strip() + + match_type = find_matches_parsing(bottom_content, path) + know_id_source = pure_text if pure_text else f"{path or ''}::{page_nums or ''}" + know_id = gen_str_codes(know_id_source) + bottom_tokens = tokenize2stw_remove([bottom_content], llm_paras["stopwords"]) + + keywords = "" + summary = "" + needs_llm = ( + not skip_llm and len(bottom_content) > summary_len and llm_paras["summary_txt"] + ) + if needs_llm: + _title, keywords, summary = extract_title_keywords_summary( + bottom_content, max_keywords=3, summary_len=summary_len + ) + + df_list.append( + ParsedRow( + content=bottom_content, + path=path, + type=match_type, + keywords=keywords, + summary=summary, + know_id=know_id, + tokens=bottom_tokens, + addtime=time_stamp, + page_nums=page_nums, + ).to_list() + ) + return df_list + + +def parse_md( + output_dir, + source_type, + file_path=None, + md_lines=None, + base_llm_paras=None, + relative_root=None, +): + if md_lines is None and file_path is not None: + from app.services.common.file_loading import is_remote, load_file_bytes + + if is_remote(file_path): + file_bytes = load_file_bytes(file_path) + md_content = file_bytes.decode("utf-8") + md_lines = md_content.splitlines() + else: + with open(file_path, "r", encoding="utf-8") as file: + md_lines = file.readlines() + + md_lines = [line.strip() for line in md_lines if line.strip() != ""] + + # Preprocess: merge multi-line HTML tables into single lines + md_lines = merge_html_tables(md_lines) + + # Detect TOC using async LLM-based detection + toc_model_name = ( + base_llm_paras.get("model_name", settings.NORMOL_MODEL) + if base_llm_paras + else settings.NORMOL_MODEL + ) + hierarchy_model_name = ( + (base_llm_paras.get("hierarchy_model_name") or toc_model_name) + if base_llm_paras + else (settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL) + ) + + with stage_timer( + "md.detect_toc", line_count=len(md_lines), model_name=toc_model_name + ): + toc_hierarchies, md_lines = detect_tocs_in_texts( + md_lines, + model_name=toc_model_name, + hierarchy_model_name=hierarchy_model_name, + ) + + # Save toc_hierarchies.json to output_dir (will be included in final zip package) + if toc_hierarchies: + toc_json_path = os.path.join(output_dir, "toc_hierarchies.json") + with open(toc_json_path, "w", encoding="utf-8") as f: + json.dump(toc_hierarchies, f, ensure_ascii=False, indent=2) + logger.info(f"Saved TOC hierarchies to {toc_json_path}") + + # Clean old artifacts to prevent accumulation across debug runs. + # In production each job uses a fresh workspace so rmtree never triggers. + tb_dir = os.path.join(output_dir, "tables") + if os.path.isdir(tb_dir): + shutil.rmtree(tb_dir) + os.makedirs(tb_dir, exist_ok=True) + img_dir = os.path.join(output_dir, "images") + if os.path.isdir(img_dir): + # Only remove parse_md's own output (image-N-*) from previous runs + for fname in os.listdir(img_dir): + if re.match(r"^image-\d+", fname): + os.remove(os.path.join(img_dir, fname)) + os.makedirs(img_dir, exist_ok=True) + + # initialize vars + split_char = settings.SPLIT_CHAR or "/" + parser_state = MarkdownParseState( + relative_root=relative_root or "", + split_char=split_char, + llm_parameters=base_llm_paras, + timestamp=get_str_time(), + row_updater=update_df_list, + ) + + # Find layout.json path + layout_json_path = os.path.join(output_dir, "layout.json") + if not os.path.exists(layout_json_path): + layout_json_path = None + logger.debug("layout.json not found, META features will not be added") + + # estimate hierarchies with toc_hierarchies context + with stage_timer( + "md.predict_headings", + line_count=len(md_lines), + smart_parse=base_llm_paras["smart_title_parse"], + model_name=hierarchy_model_name, + ): + lines_with_heading = eval_md_headings( + md_lines, + source_type, + toc_hierarchies=toc_hierarchies, + smart_parse=base_llm_paras["smart_title_parse"], + model_name=hierarchy_model_name, + output_dir=output_dir, + layout_json_path=layout_json_path, + ) + + logger.debug("Parsing md data... total_lines={}", len(lines_with_heading)) + for i, line in enumerate(lines_with_heading): + if parser_state.record_page_marker(line): + continue + + last_context = find_surround_context( + lines_with_heading, i + ) # record the previous and next line which is not table/image + current_heading, current_heading_level = md_heading_match(line, as_is=False) + + if ( + not current_heading_level == -1 + ): # indicate a new path should be evaluated or added + if parser_state.content_items: + parser_state.flush_current_content() + elif parser_state.path and parser_state.path != (relative_root or ""): + # Consecutive headings with no body text between them: + # Create a placeholder chunk so the previous heading's path + parser_state.flush_placeholder_chunk() + + parser_state.enter_heading(current_heading, current_heading_level) + + else: # no path change, remain in the same hierarchy + # a. handle lines containing images (LLM deferred to post-loop parallel batch) + # Always skip inline LLM — vision calls are deferred to parallel batch + imgs = detect_summary_img_md(line, last_context, output_dir, mode=False) + image_name = build_markdown_image_name( + image_count=parser_state.image_count, + last_context=last_context, + ) + + for img_path, _img_title, img_summary in imgs: + image_asset = build_markdown_image_asset( + MarkdownImageAssetRequest( + output_dir=output_dir, + image_dir=img_dir, + image_path=img_path, + image_name=image_name, + image_count=parser_state.image_count, + last_context=last_context, + image_summary=img_summary, + timestamp=parser_state.timestamp, + current_page_number=parser_state.current_page_number, + seen_images=parser_state.seen_images, + summary_image=bool(base_llm_paras["summary_image"]), + row_index=len(parser_state.rows), + ) + ) + if ( + image_asset.content_item is not None + and image_asset.row_values is not None + ): + parser_state.append_content_item(image_asset.content_item) + parser_state.append_row(image_asset.row_values) + if ( + image_asset.cache_key is not None + and image_asset.cache_entry is not None + ): + parser_state.seen_images[image_asset.cache_key] = ( + image_asset.cache_entry + ) + if image_asset.deferred_task is not None: + parser_state.schedule_deferred_task(image_asset.deferred_task) + if image_asset.should_advance_image_count: + parser_state.image_count += 1 + + # TODO for large and dense tables, such as "Epstein flight logs", + # integrate tabula-py as an independent extraction path to solve VLM hallucinations and misplacement + # b. handle lines containing tables + tb_bool, form, _ = identify_tables(line) + if tb_bool: + if form == "html": + # each line is a complete table - process immediately + tb_str = line + elif form == "md": + # For MD tables, accumulate lines until table ends + parser_state.table_lines.append(line) + if i + 1 >= len(lines_with_heading): + tb_bool_next = False + else: + tb_bool_next, _, _ = identify_tables( + lines_with_heading[i + 1].strip() + ) + + if not tb_bool_next or i == len(lines_with_heading) - 1: + cleaned_table_lines, error_lines = clean_md_table_lines( + parser_state.table_lines, start_line_num=i + ) + tb_str = "\n".join(cleaned_table_lines) + parser_state.error_line_numbers.extend(error_lines) + tb_str = extract_tables_by_forms(tb_str, form="md") + else: + continue # Keep accumulating MD table lines + else: + continue # Unknown form, skip + + table_asset = build_markdown_table_asset( + MarkdownTableAssetRequest( + table_html=tb_str, + table_dir=tb_dir, + table_count=parser_state.table_count, + timestamp=parser_state.timestamp, + current_page_number=parser_state.current_page_number, + summary_table=bool(base_llm_paras["summary_table"]), + row_index=len(parser_state.rows), + ) + ) + parser_state.append_content_item(table_asset.content_item) + parser_state.append_row(table_asset.row_values) + if table_asset.deferred_task is not None: + parser_state.schedule_deferred_task(table_asset.deferred_task) + parser_state.table_lines = [] + parser_state.table_count += 1 + + # c. handle plain texts + if len(imgs) == 0 and not tb_bool: + parser_state.append_plain_text(line) + + if parser_state.content_items: + parser_state.flush_current_content() + + # Collect text chunk deferred tasks (entries needing summary/keywords) + summary_len = 1500 + parser_state.collect_text_summary_tasks(summary_len) + apply_markdown_deferred_summaries( + MarkdownDeferredSummaryInput( + rows=parser_state.rows, + tasks=parser_state.deferred_llm_tasks, + output_dir=output_dir, + summary_len=summary_len, + ) + ) + + with stage_timer("md.build_dataframe", row_count=len(parser_state.rows)): + doc_df = parser_state.to_dataframe() + + return doc_df diff --git a/apps/worker/app/services/document_parser/formats/markdown/table_asset.py b/apps/worker/app/services/document_parser/formats/markdown/table_asset.py new file mode 100644 index 000000000..19f795f4f --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/markdown/table_asset.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import cast + +from app.services.document_parser.formats.html.parser import first_cols_rows_html +from app.services.document_parser.support.identifiers import gen_str_codes +from app.services.document_parser.assets.inline_asset import build_table_asset_row +from app.services.document_parser.formats.markdown.deferred_task import ( + MarkdownDeferredSummaryTask, + TableDeferredSummaryTask, +) +from app.services.document_parser.formats.markdown.parse_state import ParserRowValues +from app.services.document_parser.tables.table_text_parser import sanitize_table_name_from_header + +from shared.utils.chunk_refs import build_chunk_ref +from app.services.common.file_utils import path_handle + + +@dataclass(frozen=True) +class MarkdownTableAsset: + content_item: str + row_values: ParserRowValues + deferred_task: MarkdownDeferredSummaryTask | None + relative_path: str + + +@dataclass(frozen=True) +class MarkdownTableAssetRequest: + table_html: str + table_dir: str + table_count: int + timestamp: str + current_page_number: int + summary_table: bool + row_index: int + + +def build_markdown_table_asset( + request: MarkdownTableAssetRequest, +) -> MarkdownTableAsset: + first_row_text, _first_col_text = first_cols_rows_html(request.table_html) + table_index = f"table-{request.table_count}" + + raw_table_name = ( + sanitize_table_name_from_header(first_row_text) if first_row_text else "" + ) + table_name = _sanitize_table_file_stem( + f"table-{str(request.table_count)} {raw_table_name}" + ) + relative_table_path = f"tables/{table_name}.html" + table_ref = build_chunk_ref(relative_table_path) + table_content_item = f"\n{table_ref}\n" + table_path = os.path.join(request.table_dir, f"{table_name}.html") + _write_table_html(table_path=table_path, table_html=request.table_html) + + table_row = build_table_asset_row( + content=request.table_html, + relative_path=relative_table_path, + summary=table_index, + keywords="", + know_id=gen_str_codes((request.table_html + str(request.table_count))), + addtime=request.timestamp, + page_nums=str(request.current_page_number) + if request.current_page_number > 0 + else "", + ) + + deferred_task = None + if request.summary_table: + deferred_task = TableDeferredSummaryTask( + row_index=request.row_index, + table_html=request.table_html, + table_dir=request.table_dir, + table_name=table_name, + table_count=request.table_count - 1, + ) + + return MarkdownTableAsset( + content_item=table_content_item, + row_values=cast(ParserRowValues, table_row.to_list()), + deferred_task=deferred_task, + relative_path=relative_table_path, + ) + + +def _sanitize_table_file_stem(raw_name: str) -> str: + table_name = path_handle(raw_name, mode="clean_single") + if not isinstance(table_name, str) or not table_name: + raise ValueError(f"Failed to sanitize Markdown table name: {raw_name}") + return table_name + + +def _write_table_html(*, table_path: str, table_html: str) -> None: + table_html_with_border = table_html.replace("", "
").replace( + "
None: - """Compatibility wrapper for the extracted MinerU workflow module.""" - parse_via_full(pdf_url, filename, output_dir, s3_key=s3_key) - - def parse_pdfs( pdf_path, filename, @@ -230,7 +223,7 @@ def parse_pdfs( # ── Atlas routing: bypass MinerU entirely, use PyMuPDF for per-page chunking ── if profile and profile.doc_category == "atlas": logger.info(f"📐 Atlas detected, bypassing MinerU for {filename}") - from app.services.document_parser.atlas_parser import parse_atlas + from app.services.document_parser.formats.atlas.parser import parse_atlas return parse_atlas( pdf_path, output_dir, base_llm_paras, relative_root, profile=profile @@ -259,14 +252,14 @@ def parse_pdfs( # ) # else: # with stage_timer("pdf.extract.standard", filename=filename): - # upload_and_parse(pdf_path, filename, output_dir, s3_key=s3_key) + # parse_via_full(pdf_path, filename, output_dir, s3_key=s3_key) # _inject_page_markers(output_dir) logger.info( f"🛡️ Conservative mode: forcing MinerU (standard) for {filename} [route={route}]" ) with stage_timer("pdf.extract.standard", filename=filename): - upload_and_parse(pdf_path, filename, output_dir, s3_key=s3_key) + parse_via_full(pdf_path, filename, output_dir, s3_key=s3_key) # Inject page markers from MinerU layout.json _inject_page_markers(output_dir) diff --git a/apps/worker/app/services/document_parser/pptx_pdf_rendering.py b/apps/worker/app/services/document_parser/formats/pdf/pptx_rendering.py similarity index 95% rename from apps/worker/app/services/document_parser/pptx_pdf_rendering.py rename to apps/worker/app/services/document_parser/formats/pdf/pptx_rendering.py index 6e0a881a9..de5373ead 100644 --- a/apps/worker/app/services/document_parser/pptx_pdf_rendering.py +++ b/apps/worker/app/services/document_parser/formats/pdf/pptx_rendering.py @@ -1,7 +1,7 @@ import os import tempfile -from app.services.document_parser.pymupdf_subprocess import run_in_child_process, worker +from app.services.document_parser.formats.pdf.pymupdf_subprocess import run_in_child_process, worker from loguru import logger diff --git a/apps/worker/app/services/document_parser/pymupdf_subprocess.py b/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py similarity index 96% rename from apps/worker/app/services/document_parser/pymupdf_subprocess.py rename to apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py index c98b4678c..fdcf64d35 100644 --- a/apps/worker/app/services/document_parser/pymupdf_subprocess.py +++ b/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py @@ -24,6 +24,7 @@ from multiprocessing.process import BaseProcess from multiprocessing.queues import Queue as MultiprocessingQueue from threading import RLock +from typing import TYPE_CHECKING from app.core.runtime_limits import read_pymupdf_max_concurrent from loguru import logger @@ -33,7 +34,8 @@ TimeoutException, ) -from gevent.threadpool import ThreadPool as GeventThreadPool +if TYPE_CHECKING: + from gevent.threadpool import ThreadPool as GeventThreadPool # Default timeout for child processes (seconds) DEFAULT_TIMEOUT = 3000 @@ -119,13 +121,13 @@ def _close_result_queue(result_queue: MultiprocessingQueue) -> None: """Release parent-side queue resources once the child result is no longer needed.""" try: result_queue.close() - except Exception: - pass + except Exception as exc: + logger.debug(f"Failed to close PyMuPDF result queue: {exc}") try: result_queue.join_thread() - except Exception: - pass + except Exception as exc: + logger.debug(f"Failed to join PyMuPDF result queue thread: {exc}") def _run_worker_in_spawned_process( diff --git a/apps/worker/app/services/document_parser/formats/pdf/rendered_transform.py b/apps/worker/app/services/document_parser/formats/pdf/rendered_transform.py new file mode 100644 index 000000000..aa5f9adda --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/pdf/rendered_transform.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import os + +import pandas as pd + +from app.services.document_parser.providers.mineru.pdf_service import ( + get_existing_mineru_source_s3_key, +) +from app.services.document_parser.formats.pdf.parser import parse_pdfs +from app.services.document_parser.formats.pdf.pptx_rendering import render_pdf_to_image_pdf +from loguru import logger + +from shared.core.config import settings +from shared.services.storage.job_file_storage import JobFileStorage + +RENDERED_PDF_TEMP_FILENAME = "_pptx_tmp.pdf" + + +def build_rendered_pdf_s3_key(job_id: str | None) -> str | None: + """Store rendered parser artifacts under a stable transform/ prefix.""" + if settings.ENVIRONMENT == "development" or not job_id: + return None + return f"transform/{job_id}.rendered.pdf" + + +def parse_cached_rendered_pdf( + *, + rendered_pdf_s3_key: str | None, + filename: str, + output_dir: str, + base_llm_paras: dict[str, object], + relative_root: str | None, +) -> pd.DataFrame | None: + """Parse a previously rendered PDF from S3 without re-reading the source deck.""" + if rendered_pdf_s3_key is None: + return None + + cached_rendered_pdf_s3_key = get_existing_mineru_source_s3_key(rendered_pdf_s3_key) + if cached_rendered_pdf_s3_key is None: + return None + + logger.info( + f"[rendered_pdf_transform] Reusing rendered PDF for MinerU URL mode: {rendered_pdf_s3_key}" + ) + cached_rendered_pdf_path = JobFileStorage().download_upload_to_temp( + cached_rendered_pdf_s3_key, + suffix=".pdf", + temp_dir=output_dir, + ) + try: + return parse_pdfs( + cached_rendered_pdf_path, + filename, + output_dir, + base_llm_paras, + relative_root=relative_root, + s3_key=cached_rendered_pdf_s3_key, + ) + finally: + if os.path.exists(cached_rendered_pdf_path): + os.remove(cached_rendered_pdf_path) + + +def parse_rendered_pdf_bytes( + *, + pdf_bytes: bytes, + filename: str, + output_dir: str, + base_llm_paras: dict[str, object], + relative_root: str | None, + rendered_pdf_s3_key: str | None = None, +) -> pd.DataFrame: + image_only_pdf_bytes = render_pdf_to_image_pdf(pdf_bytes) + temporary_pdf_path = os.path.join(output_dir, RENDERED_PDF_TEMP_FILENAME) + with open(temporary_pdf_path, "wb") as temporary_pdf_file: + temporary_pdf_file.write(image_only_pdf_bytes) + + try: + return parse_pdfs( + temporary_pdf_path, + filename=filename, + output_dir=output_dir, + base_llm_paras=base_llm_paras, + relative_root=relative_root, + s3_key=rendered_pdf_s3_key, + ) + finally: + if os.path.exists(temporary_pdf_path): + os.remove(temporary_pdf_path) diff --git a/apps/worker/app/services/document_parser/pptx_parser.py b/apps/worker/app/services/document_parser/formats/pptx/parser.py similarity index 84% rename from apps/worker/app/services/document_parser/pptx_parser.py rename to apps/worker/app/services/document_parser/formats/pptx/parser.py index dde943110..fb8a93acc 100755 --- a/apps/worker/app/services/document_parser/pptx_parser.py +++ b/apps/worker/app/services/document_parser/formats/pptx/parser.py @@ -6,20 +6,17 @@ import jwt import requests -from app.services.common.kb_utils import find_images -from app.services.document_parser.legacy_converter import ( +from app.services.document_parser.support.path_helpers import find_images +from app.services.document_parser.conversion.legacy_converter import ( _convert_with_libreoffice, ) -from app.services.document_parser.md_parser import parse_md -from app.services.document_parser.mineru_pdf_service import ( - get_existing_mineru_source_s3_key, +from app.services.document_parser.formats.markdown.parser import parse_md +from app.services.document_parser.support.parser_log_utils import truncate_log_value +from app.services.document_parser.formats.pdf.rendered_transform import ( + build_rendered_pdf_s3_key, + parse_cached_rendered_pdf, + parse_rendered_pdf_bytes, ) -from app.services.document_parser.parser_log_utils import truncate_log_value -from app.services.document_parser.pdf_parser import parse_pdfs -from app.services.document_parser.pptx_pdf_rendering import ( - render_pdf_to_image_pdf as _render_pdf_to_image_pdf, -) -from app.services.storage.sync_storage_service import download_s3_object_to_temp from loguru import logger from markitdown import MarkItDown from pptx2md import ConversionConfig, convert @@ -29,8 +26,8 @@ FileSystemException, ) from shared.core.logging import LogEvent -from shared.utils.CommonHelperSync import load_file_bytes -from shared.utils.file_utils import path_handle +from app.services.common.file_loading import load_file_bytes +from app.services.common.file_utils import path_handle # ==================== LibreOffice conversion ==================== @@ -59,7 +56,7 @@ def pptx_to_pdf_libreoffice(pptx_path, outdir="."): def _get_iloveapi_token_lease(): """acquire iLoveAPI token lease from the quotas pool and generate a JWT token""" - from shared.utils.iloveapi_quota_manager import get_iloveapi_quota_manager + from shared.services.ai.iloveapi_quota_manager import get_iloveapi_quota_manager quota_manager = get_iloveapi_quota_manager() @@ -114,7 +111,7 @@ def _pptx_bytes_to_pdf_bytes(pptx_bytes: bytes, filename: str) -> bytes: Acquires an in-flight slot before starting. Slots are released only if the Redis-backed reservation actually succeeded. """ - from shared.utils.iloveapi_quota_manager import get_iloveapi_quota_manager + from shared.services.ai.iloveapi_quota_manager import get_iloveapi_quota_manager quota_manager = get_iloveapi_quota_manager() @@ -306,50 +303,6 @@ class _ILoveApiConcurrencyExceeded(Exception): pass -def _build_rendered_pdf_s3_key(job_id: str | None) -> str | None: - """Store rendered parser artifacts under a stable transform/ prefix.""" - if settings.ENVIRONMENT == "development" or not job_id: - return None - return f"transform/{job_id}.rendered.pdf" - - -def _parse_cached_rendered_pdf( - rendered_pdf_s3_key: str | None, - filename: str, - output_dir: str, - base_llm_paras, - relative_root, -): - """Parse a previously rendered PPTX PDF from S3 without re-reading the source deck.""" - if rendered_pdf_s3_key is None: - return None - - cached_rendered_pdf_s3_key = get_existing_mineru_source_s3_key(rendered_pdf_s3_key) - if cached_rendered_pdf_s3_key is None: - return None - - logger.info( - f"[parse_pptx] Reusing rendered PDF for MinerU URL mode: {rendered_pdf_s3_key}" - ) - cached_rendered_pdf_path = download_s3_object_to_temp( - cached_rendered_pdf_s3_key, - suffix=".pdf", - temp_dir=output_dir, - ) - try: - return parse_pdfs( - cached_rendered_pdf_path, - filename, - output_dir, - base_llm_paras, - relative_root=relative_root, - s3_key=cached_rendered_pdf_s3_key, - ) - finally: - if os.path.exists(cached_rendered_pdf_path): - os.remove(cached_rendered_pdf_path) - - # ==================== main parsing entrance ==================== @@ -372,12 +325,12 @@ def parse_pptx( - "to_pdf_api": use iLoveAPI to convert to PDF, then parse via MinerU (recommended) """ rendered_pdf_s3_key = ( - _build_rendered_pdf_s3_key(job_id) + build_rendered_pdf_s3_key(job_id) if strategy in {"to_pdf_api", "to_pdf"} else None ) if strategy in {"to_pdf_api", "to_pdf"}: - cached_result = _parse_cached_rendered_pdf( + cached_result = parse_cached_rendered_pdf( rendered_pdf_s3_key=rendered_pdf_s3_key, filename=filename, output_dir=output_dir, @@ -491,27 +444,14 @@ def _parse_pptx_via_api( # Step 1: PPTX → PDF (in memory) pdf_bytes = _pptx_bytes_to_pdf_bytes(pptx_data, filename) - # Step 2: PDF → image-only PDF (in memory) - img_pdf_bytes = _render_pdf_to_image_pdf(pdf_bytes) - - # Step 3: Write to output_dir for MinerU upload, then clean up - tmp_path = os.path.join(output_dir, "_pptx_tmp.pdf") - with open(tmp_path, "wb") as f: - f.write(img_pdf_bytes) - - try: - parsed_df = parse_pdfs( - tmp_path, - filename=filename, - output_dir=output_dir, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - s3_key=rendered_pdf_s3_key, - ) - return parsed_df - finally: - if os.path.exists(tmp_path): - os.remove(tmp_path) + return parse_rendered_pdf_bytes( + pdf_bytes=pdf_bytes, + filename=filename, + output_dir=output_dir, + base_llm_paras=base_llm_paras, + relative_root=relative_root, + rendered_pdf_s3_key=rendered_pdf_s3_key, + ) def _parse_pptx_via_libreoffice( @@ -552,25 +492,14 @@ def _parse_pptx_via_libreoffice( finally: shutil.rmtree(tmp_dir, ignore_errors=True) - img_pdf_bytes = _render_pdf_to_image_pdf(pdf_bytes) - - tmp_path = os.path.join(output_dir, "_pptx_tmp.pdf") - with open(tmp_path, "wb") as f: - f.write(img_pdf_bytes) - - try: - parsed_df = parse_pdfs( - tmp_path, - filename=filename, - output_dir=output_dir, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - s3_key=rendered_pdf_s3_key, - ) - return parsed_df - finally: - if os.path.exists(tmp_path): - os.remove(tmp_path) + return parse_rendered_pdf_bytes( + pdf_bytes=pdf_bytes, + filename=filename, + output_dir=output_dir, + base_llm_paras=base_llm_paras, + relative_root=relative_root, + rendered_pdf_s3_key=rendered_pdf_s3_key, + ) def _parse_pptx_to_md(pptx_data, filename, output_dir, base_llm_paras, relative_root): diff --git a/apps/worker/app/services/document_parser/txt_parser.py b/apps/worker/app/services/document_parser/formats/text/parser.py similarity index 98% rename from apps/worker/app/services/document_parser/txt_parser.py rename to apps/worker/app/services/document_parser/formats/text/parser.py index 3851817b5..161e09864 100755 --- a/apps/worker/app/services/document_parser/txt_parser.py +++ b/apps/worker/app/services/document_parser/formats/text/parser.py @@ -12,8 +12,8 @@ from shared.services.ai.prompt_service import build_prompt from shared.services.ai.response_process_service import eval_response from shared.utils.chunk_refs import CHUNK_REF_PATTERN -from shared.utils.CommonHelperSync import load_file_bytes -from shared.utils.OpenAICompatibleClientSync import get_openai_client +from app.services.common.file_loading import load_file_bytes +from shared.services.ai.openai_compatible_client_sync import get_openai_client def clean_texts_by_form(text, form="html"): diff --git a/apps/worker/app/services/document_parser/html_parser.py b/apps/worker/app/services/document_parser/html_parser.py deleted file mode 100644 index c94fbfce8..000000000 --- a/apps/worker/app/services/document_parser/html_parser.py +++ /dev/null @@ -1,1005 +0,0 @@ -# pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportReturnType=false -""" -HTML rendering utilities for DataFrame to HTML conversion. - -This module provides functions for converting pandas DataFrames to HTML tables -with support for: -- MultiIndex columns with colspan/rowspan merging -- Row headers (semantic
elements) -- Proper HTML escaping and formatting -- DOCX table to HTML conversion -- HTML to DataFrame/Markdown conversion -""" - -from typing import Dict, List, Optional, Union - -import pandas as pd -from bs4 import BeautifulSoup -from docx.table import Table as DocxTable - -from shared.core.exceptions.domain_exceptions import TableParsingException -from shared.utils.text_utils import remove_duplicates_orderkept - - -class HTMLHeaderExpander: - """ - Expands multi-level HTML table headers into a flat list of column names. - - Handles: - - rowspan: cells that span multiple rows - - colspan: cells that span multiple columns - - Multi-row headers: combines parent -> child relationships with " > " - - Example: - Row 0: | A(rs=2) | B(cs=2) | - Row 1: | | B1 | B2 | - - Output: ["A", "B > B1", "B > B2"] - - Ported from snap-fill's HeaderMatrixExpander, using BeautifulSoup instead of PyQuery. - """ - - def __init__(self, table_html: str): - """Initialize with table HTML content.""" - soup = BeautifulSoup(table_html, "html.parser") - table = soup.find("table") - if table: - self.rows = table.find_all("tr", recursive=False) - # Also check inside thead/tbody - if not self.rows: - self.rows = table.find_all("tr") - else: - self.rows = [] - - def expand_headers( - self, header_row_count: int = 2, start_row: int = 0 - ) -> List[str]: - """ - Expand multi-row headers into a flat column list. - - Args: - header_row_count: Number of rows that form the header (default 2) - start_row: Starting row index for headers (default 0) - - Returns: - List of column header strings, with nested headers joined by " > " - """ - grid = self._build_grid(header_row_count, start_row) - if not grid or not grid[0]: - return [] - - max_cols = len(grid[0]) - - # Build column headers by combining rows vertically - headers = [] - for col in range(max_cols): - parts = [] - prev_text = None - for row in range(len(grid)): - text = grid[row][col] - # Only add if non-empty and different from previous - if text and text != prev_text: - parts.append(text) - prev_text = text - - if parts: - headers.append(" > ".join(parts)) - else: - headers.append("") # Empty column - - return headers - - def _build_grid( - self, header_row_count: int, start_row: int = 0 - ) -> List[List[Optional[str]]]: - """ - Build expanded grid from table rows. - - Algorithm: - 1. First pass: calculate max columns by expanding all colspan - 2. Create occupied grid to track cells filled by rowspan/colspan - 3. For each cell, place text in all positions it spans - """ - if start_row + header_row_count > len(self.rows): - return [] - - header_rows = self.rows[start_row : start_row + header_row_count] - - # First pass: calculate max columns (sum of colspan for each row) - max_cols = 0 - for tr in header_rows: - col_count = 0 - for td in tr.find_all(["td", "th"], recursive=False): - col_count += int(td.get("colspan", 1)) - max_cols = max(max_cols, col_count) - - if max_cols == 0: - return [] - - # Create grid and occupied tracker - grid = [[None] * max_cols for _ in range(header_row_count)] - occupied = [[False] * max_cols for _ in range(header_row_count)] - - # Process each row - for row_idx, tr in enumerate(header_rows): - col_ptr = 0 # Current column position in this row - - for td in tr.find_all(["td", "th"], recursive=False): - text = td.get_text(strip=True) - rowspan = int(td.get("rowspan", 1)) - colspan = int(td.get("colspan", 1)) - - # Skip past any already-occupied cells (from previous row's rowspan) - while col_ptr < max_cols and occupied[row_idx][col_ptr]: - col_ptr += 1 - - if col_ptr >= max_cols: - break - - # Fill all cells covered by this rowspan/colspan - for r in range(row_idx, min(row_idx + rowspan, header_row_count)): - for c in range(col_ptr, min(col_ptr + colspan, max_cols)): - grid[r][c] = text - occupied[r][c] = True - - # Move column pointer past this cell - col_ptr += colspan - - return grid - - def get_grid_debug( - self, header_row_count: int = 2, start_row: int = 0 - ) -> List[List[Optional[str]]]: - """Return the expanded grid for debugging and testing purposes.""" - return self._build_grid(header_row_count, start_row) - - def get_unique_headers( - self, header_row_count: int = 2, start_row: int = 0 - ) -> List[str]: - """ - Get deduplicated headers for use in keywords/LLM prompts. - Removes duplicate column names (from colspan expansion) while preserving order. - """ - headers = self.expand_headers(header_row_count, start_row) - seen = set() - unique = [] - for h in headers: - if h and h not in seen: - seen.add(h) - unique.append(h) - return unique - - def detect_header_row_count( - self, start_row: int = 0, max_scan_rows: int = 5 - ) -> int: - """ - Automatically detect header row count using the maximum rowspan. - - Logic: - - scan at most ``max_scan_rows`` rows starting at ``start_row`` - - find the largest rowspan value among all cells - - ``rowspan=2`` means a two-row header, ``rowspan=3`` means three rows - - Returns: - int: Detected header row count, at least 1. - """ - if start_row >= len(self.rows): - return 1 - - max_rowspan = 1 - scan_end = min(start_row + max_scan_rows, len(self.rows)) - - for row in self.rows[start_row:scan_end]: - for cell in row.find_all(["td", "th"], recursive=False): - rowspan = int(cell.get("rowspan", 1)) - if rowspan > max_rowspan: - max_rowspan = rowspan - - return max_rowspan - - def detect_row_indices( - self, - header_row_count: int, - start_row: int = 0, - end_row: int = None, - max_scan_cols: int = 3, - ) -> Dict: - """ - Detect row-index columns and return the row-index metadata. - - Logic: - 1. start from data rows after the header - 2. scan left to right; a non-empty, non-numeric column is a row index - 3. stop once a data column is reached - - Returns: - dict: { - 'row_index_col_count': int, - 'row_index_col_name': str or None, - 'row_indices': List[str], - } - """ - result = { - "row_index_col_count": 0, - "row_index_col_name": None, - "row_indices": [], - } - - data_start_row = start_row + header_row_count - if data_start_row >= len(self.rows): - return result - - # Build header grid for column names - grid = self._build_grid(header_row_count, start_row) - if not grid or not grid[0]: - return result - - # Data rows - if end_row is not None: - data_end_row = min(end_row + 1, len(self.rows)) - else: - data_end_row = len(self.rows) - - data_rows = self.rows[data_start_row:data_end_row] - if not data_rows: - return result - - row_index_col_count = 0 - row_index_col_names = [] - - for cell_idx in range(max_scan_cols): - col_values = [] - is_index_col = True - - for row in data_rows: - cells = row.find_all(["td", "th"], recursive=False) - - if cell_idx >= len(cells): - continue - - cell_value = cells[cell_idx].get_text(strip=True) - - if not cell_value: - is_index_col = False - break - - # Check if purely numeric - if ( - cell_value.replace(".", "") - .replace("-", "") - .replace(" ", "") - .isdigit() - ): - is_index_col = False - break - - col_values.append(cell_value) - - if is_index_col and col_values: - row_index_col_count += 1 - row_index_col_names.append(col_values) - else: - break # Stop at first data column - - if row_index_col_count > 0: - result["row_index_col_count"] = row_index_col_count - result["row_index_col_name"] = grid[-1][0] if grid and grid[-1] else None - result["row_indices"] = row_index_col_names[0] - - return result - - -def parse_nested_htmltb(table): - """Parse nested HTML table into a list of rows. - - Handles nested tables recursively. - - Args: - table: BeautifulSoup table element - - Returns: - List of rows, where each row is a list of cell values or nested tables - """ - rows = [] - for tr in table.find_all("tr", recursive=False): - row = [] - for td in tr.find_all(["td", "th"], recursive=False): - inner_table = td.find("table") - if inner_table: - row.append(parse_nested_htmltb(inner_table)) - else: - row.append(td.get_text(strip=True)) - if row: - rows.append(row) - return rows - - -def html_to_md_lines(html: str): - """Convert HTML table to list of markdown-like lines. - - Args: - html: HTML string containing a table - - Returns: - List of strings, each representing a row with cells separated by ' | ' - """ - soup = BeautifulSoup(html, "html.parser") - lines = [] - for row in soup.find_all("tr"): - row_text = [] - for cell in row.find_all("td", recursive=False): - text = cell.get_text(separator=" ", strip=True) - if text: - row_text.append(text) - if row_text: - lines.append(" | ".join(row_text)) - lines = remove_duplicates_orderkept(lines) - return lines - - -def tb_htmlstr_to_df(html_str): - """Convert first table in HTML string to DataFrame""" - soup = BeautifulSoup(html_str, "html.parser") - table = soup.find("table") - if not table: - raise TableParsingException( - user_message="No table structure found in the document", - reason="INVALID_FORMAT", - internal_message="No found in the HTML string", - ) - nested_list = parse_nested_htmltb(table) - try: - df = pd.DataFrame(nested_list[1:], columns=nested_list[0]) - except Exception: - df = pd.DataFrame(nested_list) - return df - - -def merge_html_tables(lines: list) -> list: - """Merge multi-line HTML tables into single lines. - - MinerU and some parsers output HTML tables with line breaks inside. - This function joins all lines from
to
into a single line. - - Args: - lines: List of text lines (already stripped) - - Returns: - List of lines with HTML tables merged into single lines - """ - merged_lines = [] - in_table = False - table_buffer = [] - - for line in lines: - if "" in line: - merged_lines.append(" ".join(table_buffer)) - table_buffer = [] - in_table = False - elif in_table: - table_buffer.append(line) - if "
" in line: - merged_lines.append(" ".join(table_buffer)) - table_buffer = [] - in_table = False - else: - merged_lines.append(line) - - # Handle unclosed table at end of file - if table_buffer: - merged_lines.append(" ".join(table_buffer)) - - return merged_lines - - -def first_cols_rows_html(html_str, max_items=10, max_chars=20): - """Extract first row and first column from HTML table string. - - This function mirrors the logic of _first_cols_rows in doc_parser.py - but works with HTML strings instead of DOCX Table objects. - - Args: - html_str: HTML table string - max_items: Maximum number of items to extract (default 10) - max_chars: Maximum characters per item (default 20) - - Returns: - Tuple of (first_row_text, first_col_text) with ' | ' as separator - """ - from app.services.common.kb_utils import truncate_text - - soup = BeautifulSoup(html_str, "html.parser") - table = soup.find("table") - if not table: - return "", "" - - rows = table.find_all("tr") - if not rows: - return "", "" - - # First row extraction (deduplicated, order preserved, max items, truncated) - first_row_cells = rows[0].find_all(["td", "th"]) - seen_row = set() - unique_row = [] - for cell in first_row_cells: - if len(unique_row) >= max_items: - break - text = cell.get_text(strip=True) - if text and text not in seen_row: - seen_row.add(text) - unique_row.append(truncate_text(text, max_chars, 0)) - first_row_text = " | ".join(unique_row) if unique_row else "" - - # First column extraction (deduplicated, order preserved, max items, truncated) - seen_col = set() - unique_col = [] - for row in rows: - if len(unique_col) >= max_items: - break - cells = row.find_all(["td", "th"]) - if cells: - text = cells[0].get_text(strip=True) - if text and text not in seen_col: - seen_col.add(text) - unique_col.append(truncate_text(text, max_chars, 0)) - first_col_text = " | ".join(unique_col) if unique_col else "" - - return first_row_text, first_col_text - - -def table2html(table: DocxTable, cell_image_map: dict = None) -> str: - """Convert a DOCX table to HTML string with proper colspan/rowspan handling. - - Handles merged cells by: - - Detecting horizontal merges via comparing cell._tc objects - - Detecting vertical merges via vMerge XML attribute - - Generating proper colspan and rowspan attributes - - Args: - table: python-docx Table object - cell_image_map: Optional dict mapping (row_idx, col_idx) to image description - strings. col_idx corresponds to the unique tc index in each row - (matching XML ordering, not expanded python-docx cells). - - Returns: - HTML string representation of the table with merged cells - """ - - NS = {"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main"} - - def get_cell_vmerge(cell): - """Get vMerge status: 'restart', 'continue', or None""" - tc = cell._tc - tcPr = tc.find(".//w:tcPr", namespaces=NS) - if tcPr is not None: - vMerge = tcPr.find(".//w:vMerge", namespaces=NS) - if vMerge is not None: - val = vMerge.get( - "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}val" - ) - return ( - val if val else "continue" - ) # If no val attribute, it's a continuation - return None - - n_rows = len(table.rows) - if n_rows == 0: - return "
" - - # Build grid: track unique cells and their positions - # grid[row][col] = (cell_tc_id, cell, is_new_cell) - # We use id(cell._tc) as unique identifier for cells - - grid = [] - for row_idx, row in enumerate(table.rows): - row_data = [] - prev_tc_id = None - for cell in row.cells: - tc_id = id(cell._tc) - is_new = tc_id != prev_tc_id - row_data.append((tc_id, cell, is_new)) - prev_tc_id = tc_id - grid.append(row_data) - - # Rows may have different cell counts due to complex merges; - # use the maximum for grid allocation, per-row length for access. - n_cols = max(len(r) for r in grid) if grid else 0 - - # Calculate colspan for each cell (count consecutive cells with same _tc) - colspan_grid = [[0] * n_cols for _ in range(n_rows)] - - for row_idx in range(n_rows): - row_len = len(grid[row_idx]) - col_idx = 0 - while col_idx < row_len: - tc_id = grid[row_idx][col_idx][0] - span = 1 - while ( - col_idx + span < row_len and grid[row_idx][col_idx + span][0] == tc_id - ): - span += 1 - colspan_grid[row_idx][col_idx] = span - col_idx += span - - # Calculate rowspan for cells with vMerge='restart' - rowspan_grid = [[1] * n_cols for _ in range(n_rows)] - - for col_idx in range(n_cols): - row_idx = 0 - while row_idx < n_rows: - if col_idx >= len(grid[row_idx]): - row_idx += 1 - continue - cell = grid[row_idx][col_idx][1] - vmerge = get_cell_vmerge(cell) - - if vmerge == "restart": - # Count how many 'continue' cells follow - span = 1 - while row_idx + span < n_rows: - if col_idx >= len(grid[row_idx + span]): - break - next_cell = grid[row_idx + span][col_idx][1] - next_vmerge = get_cell_vmerge(next_cell) - if next_vmerge == "continue": - span += 1 - else: - break - rowspan_grid[row_idx][col_idx] = span - row_idx += span - elif vmerge == "continue": - # This cell is part of a vertical merge, mark as 0 (skip) - rowspan_grid[row_idx][col_idx] = 0 - row_idx += 1 - else: - row_idx += 1 - - # Build HTML - html_parts = [""] - - for row_idx in range(n_rows): - html_parts.append("") - col_idx = 0 - unique_col_idx = 0 # Tracks unique tc index per row (matches XML order) - - while col_idx < len(grid[row_idx]): - tc_id, cell, is_new = grid[row_idx][col_idx] - - # Skip if this cell is a horizontal continuation - if not is_new: - col_idx += 1 - continue - - # Skip if this cell is a vertical continuation - rowspan = rowspan_grid[row_idx][col_idx] - if rowspan == 0: - unique_col_idx += 1 - col_idx += 1 - continue - - colspan = colspan_grid[row_idx][col_idx] - - # Build cell content - if cell.tables: - # Nested table - content = "".join( - table2html(nested_table) for nested_table in cell.tables - ) - else: - content = cell.text.strip().replace("\n", "
") - - # Append image descriptions if available - if cell_image_map: - img_desc = cell_image_map.get((row_idx, unique_col_idx)) - if img_desc: - content += f"
{img_desc}" - - # Build attributes - attrs = [] - if colspan > 1: - attrs.append(f'colspan="{colspan}"') - if rowspan > 1: - attrs.append(f'rowspan="{rowspan}"') - - attr_str = " " + " ".join(attrs) if attrs else "" - html_parts.append(f"{content}") - - unique_col_idx += 1 - col_idx += colspan - - html_parts.append("") - - html_parts.append("
") - return "".join(html_parts) - - -def render_multiindex_thead(columns: pd.MultiIndex, escape: bool = False) -> str: - """ - Convert MultiIndex columns to HTML thead with colspan/rowspan. - - This function generates a proper multi-row structure where: - - Horizontally adjacent identical values are merged with colspan - - Vertically repeated values are merged with rowspan - - Args: - columns: pandas MultiIndex representing the column headers - escape: Whether to HTML-escape the cell values - - Returns: - HTML string for the element - """ - import html as html_lib - - n_levels = columns.nlevels - n_cols = len(columns) - - # Build a 2D grid of values [level][col] - grid = [] - for level in range(n_levels): - row = [columns.get_level_values(level)[col] for col in range(n_cols)] - grid.append(row) - - # Calculate colspan for each cell (horizontal merging) - # colspan[level][col] = number of columns this cell spans - colspan = [[1] * n_cols for _ in range(n_levels)] - - for level in range(n_levels): - col = 0 - while col < n_cols: - span = 1 - while col + span < n_cols and grid[level][col] == grid[level][col + span]: - # Check if the parent cells also match (for correct hierarchical merging) - parent_match = True - for parent_level in range(level): - if grid[parent_level][col] != grid[parent_level][col + span]: - parent_match = False - break - if parent_match: - span += 1 - else: - break - colspan[level][col] = span - col += span - - # Calculate rowspan for each cell (vertical merging) - # A cell has rowspan > 1 if all cells in the same column below have the same value - # AND if they would have the same colspan - rowspan = [[1] * n_cols for _ in range(n_levels)] - - for col in range(n_cols): - level = 0 - while level < n_levels: - span = 1 - # Check if cells below have the same value AND same colspan - while level + span < n_levels: - if ( - grid[level][col] == grid[level + span][col] - and colspan[level][col] == colspan[level + span][col] - ): - span += 1 - else: - break - rowspan[level][col] = span - level += span - - # Build HTML rows - # Track which cells are "covered" by rowspan from above - covered = [[False] * n_cols for _ in range(n_levels)] - - html_parts = [""] - - for level in range(n_levels): - html_parts.append('') - col = 0 - while col < n_cols: - if covered[level][col]: - # This cell is covered by a rowspan from above, skip it - col += 1 - continue - - # Get cell value - val = grid[level][col] - val_str = str(val) if val is not None else "" - if escape: - val_str = html_lib.escape(val_str) - - # Get spans - cs = colspan[level][col] - rs = rowspan[level][col] - - # Mark covered cells - for r_offset in range(rs): - for c_offset in range(cs): - if r_offset > 0 or c_offset > 0: - if level + r_offset < n_levels and col + c_offset < n_cols: - covered[level + r_offset][col + c_offset] = True - - # Build th element with attributes - attrs = [] - if cs > 1: - attrs.append(f'colspan="{cs}"') - if rs > 1: - attrs.append(f'rowspan="{rs}"') - - attr_str = " " + " ".join(attrs) if attrs else "" - html_parts.append(f"{val_str}") - - col += cs - - html_parts.append("") - - html_parts.append("") - return "".join(html_parts) - - -def render_tbody_with_row_headers( - tb_df: pd.DataFrame, - row_header_cols: int = 0, - na_rep: str = "—", - escape: bool = False, -) -> str: - """ - Render DataFrame body with support for row headers and cell merging. - - This function generates a proper structure where: - - Row header columns use instead of - - Horizontally adjacent identical values in row headers are merged with colspan - - Vertically adjacent identical values in row headers are merged with rowspan - - Merging respects hierarchical structure - - Args: - tb_df: DataFrame to render - row_header_cols: Number of leftmost columns to render as - na_rep: String representation for NaN values - escape: Whether to HTML-escape values - - Returns: - HTML string for the element - """ - import html as html_lib - - if row_header_cols <= 0: - # No row headers - simple rendering without merging - html_parts = [""] - for _, row in tb_df.iterrows(): - html_parts.append("") - for val in row: - if pd.isna(val): - val_str = na_rep - else: - val_str = str(val) - if escape: - val_str = html_lib.escape(val_str) - html_parts.append(f"{val_str}") - html_parts.append("") - html_parts.append("") - return "".join(html_parts) - - n_rows = len(tb_df) - n_cols = len(tb_df.columns) - - if n_rows == 0: - return "" - - # Build 2D grid of values for row header columns - # grid[row_idx][col_idx] = value - grid = [] - for row_idx in range(n_rows): - row_values = [] - for col_idx in range(row_header_cols): - val = tb_df.iloc[row_idx, col_idx] - if pd.isna(val): - val = na_rep - else: - val = str(val) - row_values.append(val) - grid.append(row_values) - - # Calculate colspan for each cell (horizontal merging within same row) - # colspan[row_idx][col_idx] = number of columns this cell spans - colspan = [[1] * row_header_cols for _ in range(n_rows)] - - for row_idx in range(n_rows): - col_idx = 0 - while col_idx < row_header_cols: - span = 1 - while ( - col_idx + span < row_header_cols - and grid[row_idx][col_idx] == grid[row_idx][col_idx + span] - ): - span += 1 - colspan[row_idx][col_idx] = span - col_idx += span - - # Calculate rowspan for each cell (vertical merging) - # Only calculate rowspan for cells that start a colspan group - # rowspan[row_idx][col_idx] = number of rows this cell spans - rowspan = [[1] * row_header_cols for _ in range(n_rows)] - - col_idx = 0 - while col_idx < row_header_cols: - row_idx = 0 - while row_idx < n_rows: - # Only process cells that start a colspan group (not covered by colspan from left) - if col_idx > 0 and grid[row_idx][col_idx] == grid[row_idx][col_idx - 1]: - row_idx += 1 - continue - - current_colspan = colspan[row_idx][col_idx] - span = 1 - - while row_idx + span < n_rows: - # Check if the value matches - if grid[row_idx][col_idx] != grid[row_idx + span][col_idx]: - break - # Check if colspan in the next row also matches - if colspan[row_idx + span][col_idx] != current_colspan: - break - # Check if all parent columns (to the left) also have same rowspan behavior - parent_match = True - for parent_col in range(col_idx): - if grid[row_idx][parent_col] != grid[row_idx + span][parent_col]: - parent_match = False - break - if parent_match: - span += 1 - else: - break - - rowspan[row_idx][col_idx] = span - row_idx += span - col_idx += 1 - - # Track which cells are covered by rowspan from above or colspan from left - covered = [[False] * row_header_cols for _ in range(n_rows)] - - # Mark cells covered by colspan (horizontal) - for row_idx in range(n_rows): - col_idx = 0 - while col_idx < row_header_cols: - cs = colspan[row_idx][col_idx] - for offset in range(1, cs): - if col_idx + offset < row_header_cols: - covered[row_idx][col_idx + offset] = True - col_idx += cs - - # Mark cells covered by rowspan (vertical) - for row_idx in range(n_rows): - for col_idx in range(row_header_cols): - if covered[row_idx][col_idx]: - continue # Skip cells already covered by colspan - rs = rowspan[row_idx][col_idx] - for offset in range(1, rs): - if row_idx + offset < n_rows: - # Mark all cells in the rowspan as covered - cs = colspan[row_idx][col_idx] - for c_offset in range(cs): - if col_idx + c_offset < row_header_cols: - covered[row_idx + offset][col_idx + c_offset] = True - - # Build HTML - html_parts = [""] - - for row_idx in range(n_rows): - html_parts.append("") - - # Render row header columns with rowspan/colspan - for col_idx in range(row_header_cols): - if covered[row_idx][col_idx]: - # This cell is covered by a rowspan/colspan, skip it - continue - - val_str = grid[row_idx][col_idx] - if escape: - val_str = html_lib.escape(val_str) - - rs = rowspan[row_idx][col_idx] - cs = colspan[row_idx][col_idx] - - attrs = [] - if rs > 1: - attrs.append(f'rowspan="{rs}"') - if cs > 1: - attrs.append(f'colspan="{cs}"') - - attr_str = " " + " ".join(attrs) if attrs else "" - html_parts.append(f'{val_str}') - - # Render data columns - for col_idx in range(row_header_cols, n_cols): - val = tb_df.iloc[row_idx, col_idx] - if pd.isna(val): - val_str = na_rep - else: - val_str = str(val) - if escape: - val_str = html_lib.escape(val_str) - html_parts.append(f"{val_str}") - - html_parts.append("") - - html_parts.append("") - return "".join(html_parts) - - -def df2html( - tb_df: pd.DataFrame, - *, - index: bool = False, - classes: Union[str, List[str], None] = "table table-striped", - na_rep: str = "—", - escape: bool = False, - row_header_cols: int = 0, -) -> str: - """Convert DataFrame to HTML table. - - Supports: - - MultiIndex columns with proper colspan/rowspan merging - - Row headers (leftmost columns rendered as ) - - Custom CSS classes and NA representation - - Args: - tb_df: DataFrame to convert - index: Whether to include the DataFrame index (not commonly used) - classes: CSS classes to add to the table - na_rep: String representation for NaN values - escape: Whether to HTML-escape values - row_header_cols: Number of leftmost columns to render as row headers (). - These columns will use instead of . - - Returns: - HTML table string - """ - class_str = ( - classes if isinstance(classes, str) else " ".join(classes) if classes else "" - ) - - # Check if columns are MultiIndex - use advanced rendering - if isinstance(tb_df.columns, pd.MultiIndex): - # Use specialized rendering for MultiIndex columns - thead_html = render_multiindex_thead(tb_df.columns, escape=escape) - tbody_html = render_tbody_with_row_headers( - tb_df, row_header_cols, na_rep, escape - ) - return f'{thead_html}{tbody_html}
' - - # Simple columns case - if row_header_cols <= 0: - # Use default pandas to_html for simple case without row headers - table_html = tb_df.to_html( - index=index, - na_rep=na_rep, - classes=classes, - escape=escape, - border=0, - justify="center", - ) - return table_html.replace("\n", "") - - # Simple columns with row headers - custom rendering - import html as html_lib - - html_parts = [f''] - - # Build thead - html_parts.append("") - html_parts.append('') - for col in tb_df.columns: - col_str = str(col) if col is not None else "" - if escape: - col_str = html_lib.escape(col_str) - html_parts.append(f"") - html_parts.append("") - html_parts.append("") - - # Build tbody with row headers - tbody_html = render_tbody_with_row_headers(tb_df, row_header_cols, na_rep, escape) - html_parts.append(tbody_html) - - html_parts.append("
{col_str}
") - - return "".join(html_parts) diff --git a/apps/worker/app/services/document_parser/layout_parser.py b/apps/worker/app/services/document_parser/layout_parser.py deleted file mode 100755 index f308e7f5a..000000000 --- a/apps/worker/app/services/document_parser/layout_parser.py +++ /dev/null @@ -1,2114 +0,0 @@ -# pyright: reportArgumentType=false, reportAssignmentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportGeneralTypeIssues=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalSubscript=false -import os -import re -import unicodedata -from collections import Counter, defaultdict - -import gevent -import pandas as pd -from app.services.common.kb_utils import ( - count_cn_en, - truncate_text_by_tokens, -) -from app.services.document_parser.stage_profiler import stage_timer -from app.services.document_parser.table_parser import df2md -from docx.oxml.ns import qn -from gevent.pool import Pool as GeventPool - -try: - from markitdown import MarkItDown -except ImportError: - # Fall back to a pass-through shim when markitdown is unavailable. - class MarkItDown: - def convert(self, content): - return content - - -from loguru import logger - -from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import WorkerHandlingException - -# TaskRedis dependency is removed, use Redis directly to track -from shared.services.ai.prompt_service import build_prompt -from shared.services.ai.response_process_service import eval_response - -# ARQ dependency is removed, use Celery instead -from shared.utils.OpenAICompatibleClientSync import get_openai_client - -# ==================== Helper Functions ==================== - - -def _resolve_hierarchy_model_name(model_name=None): - """Resolve the dedicated hierarchy LLM model with backward-compatible fallback.""" - return model_name or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL - - -def save_intermediate_csv(df: pd.DataFrame, output_dir: str, filename: str): - """ - save intermediate result to csv file, use utf-8-sig encoding to support Chinese and English - Only saves when LOCAL_DEBUG environment variable is set to 'true'. - - Args: - df: DataFrame to save - output_dir: output directory path - filename: filename (without extension) - """ - if os.environ.get("LOCAL_DEBUG", "").lower() not in ("true", "1"): - return - if output_dir is None or df is None or df.empty: - return - - try: - csv_path = os.path.join(output_dir, f"{filename}.csv") - df.to_csv(csv_path, index=False, encoding="utf-8-sig") - logger.debug(f"📊 Saved intermediate result to {csv_path}, rows={len(df)}") - except Exception as e: - logger.warning(f"Failed to save intermediate CSV {filename}: {e}") - - -# ==================== Tree Structure Functions (from sxjg) ==================== - - -def build_tree_from_dataframe(df): - """ - develop json tree from dataframe - - Args: - df: DataFrame, including id, heading, level columns - - Returns: - tree: pure nested dict structure - node_to_id: map from tree node to id (use unique node key) - id_to_row: map from id to original row data - """ - headings = df[df["level"] > -1].copy() - - node_to_id = {} # {(tree_node_key, parent_path): id} - id_to_node_info = {} # {id: (tree_node_key, parent_path)} - id_to_row = {} - root = {} - stack = [(0, root, "ROOT", "")] - - for _, row in headings.iterrows(): - heading_txt = row["heading"] - row_id = int(row["id"]) - level = int(row["level"]) - - # record id to row mapping - id_to_row[row_id] = row.to_dict() - - # find suitable parent node - while len(stack) > 1 and stack[-1][0] >= level: - stack.pop() - - # get parent node info - parent_level, parent_dict, parent_heading, parent_path = stack[-1] - - # create unique key for tree node: if there are duplicate headings under the same parent, add ID suffix - tree_node_key = heading_txt - - if tree_node_key in parent_dict: - tree_node_key = f"{heading_txt}#{row_id}" - - # build mapping: use (tree_node_key, parent_path) as key - node_key = (tree_node_key, parent_path) - node_to_id[node_key] = row_id - id_to_node_info[row_id] = node_key - - parent_dict[tree_node_key] = {} - current_path = ( - f"{parent_path}/{tree_node_key}" if parent_path else tree_node_key - ) - stack.append((level, parent_dict[tree_node_key], tree_node_key, current_path)) - return root, node_to_id, id_to_row - - -def tree_to_dataframe(tree, node_to_id, original_df): - """ - convert processed tree structure back to dataframe - - Args: - tree: processed pure nested dict structure - node_to_id: map from node to id {(tree_node_key, parent_path): id} - original_df: original dataframe - - Returns: - updated_df: updated dataframe - """ - - # extract all retained headings from tree - def extract_headings(node_dict, current_level=1, parent_path=""): - """recursively extract all headings and their new levels""" - results = [] - for tree_node_key, children in node_dict.items(): - # use (tree_node_key, parent_path) as key to find ID - node_key = (tree_node_key, parent_path) - row_id = node_to_id.get(node_key, -1) - - if row_id >= 0: - # extract original heading from tree_node_key (remove possible ID suffix) - original_heading = ( - tree_node_key.split("#")[0] - if "#" in tree_node_key - else tree_node_key - ) - - results.append( - { - "id": row_id, - "heading": original_heading, - "level": current_level, - "tree_key": tree_node_key, - "parent_path": parent_path, - } - ) - # recursively process child nodes - if isinstance(children, dict) and children: - current_path = ( - f"{parent_path}/{tree_node_key}" - if parent_path - else tree_node_key - ) - results.extend( - extract_headings(children, current_level + 1, current_path) - ) - return results - - preserved_headings = extract_headings(tree) - preserved_ids = set([h["id"] for h in preserved_headings]) - - updated_df = original_df.copy() - removed_count = 0 - level_changed_count = 0 - - for idx, row in original_df.iterrows(): - row_id = int(row["id"]) - old_level = int(row["level"]) if row["level"] not in [-2, "nan", -1] else -1 - - if old_level > -1: - if row_id in preserved_ids: - new_level = next( - (h["level"] for h in preserved_headings if h["id"] == row_id), - old_level, - ) - updated_df.at[idx, "level"] = new_level - if new_level != old_level: - level_changed_count += 1 - else: - updated_df.at[idx, "level"] = -1 - removed_count += 1 - - logger.debug( - f"Tree changed: removed headings={removed_count}, level changed={level_changed_count}, preserved headings={len(preserved_ids)}" - ) - return updated_df - - -def remove_isolated_nodes(tree): - """ - rules: if a heading has only one child heading, and the child heading has no further child headings, - then delete this isolated child heading - - Args: - tree: pure nested dict structure, format as {heading: {child_heading: {...}}} - - Returns: - processed_tree: processed tree structure - """ - - def recursive_check_and_remove(node_dict, parent_path=""): - if not isinstance(node_dict, dict): - return node_dict - - result_dict = {} - - for heading, children in node_dict.items(): - if isinstance(children, dict) and len(children) == 1: - child_heading = list(children.keys())[0] - grandchildren = children[child_heading] - - if not grandchildren or ( - isinstance(grandchildren, dict) and len(grandchildren) == 0 - ): - result_dict[heading] = {} - logger.debug( - f"remove isolated heading: {parent_path}/{heading}/{child_heading}" - ) - else: - processed_children = recursive_check_and_remove( - children, f"{parent_path}/{heading}" if parent_path else heading - ) - result_dict[heading] = processed_children - elif isinstance(children, dict) and children: - processed_children = recursive_check_and_remove( - children, f"{parent_path}/{heading}" if parent_path else heading - ) - result_dict[heading] = processed_children - else: - result_dict[heading] = children - - return result_dict - - processed_tree = recursive_check_and_remove(tree) - return processed_tree - - -# def if_no_pos_code(reason_str: str) -> bool: -# """ -# Check whether all pos_code values are zero. -# reason format: "POS [0, 0, ...] NEG [...]" -# """ -# if not reason_str or not isinstance(reason_str, str): -# return True - -# pos_match = re.search(r'POS\s*\[([^\]]*)\]', reason_str) -# if not pos_match: -# return True -# pos_content = pos_match.group(1) -# try: -# nums = [int(x.strip()) for x in pos_content.split(',') if x.strip()] -# return all(x == 0 for x in nums) -# except: -# return True - - -# ==================== Level Mapping Functions ==================== - - -def build_level_mapping(df, origin_lvls, mode="max"): - df = df.copy() - df["origin_level"] = origin_lvls - - mapping = df.groupby("reason")["level"].apply(list).to_dict() - - processed_mapping = {} - for reason, lvls in mapping.items(): - positive_lvls = [lvl for lvl in lvls if lvl > -1] - counts = Counter(lvls) - - if not positive_lvls: - mapped_lvl = -1 - elif mode == "max": - mapped_lvl = max(positive_lvls) - elif mode == "freq": - mapped_lvl = counts.most_common(1)[0][0] - else: - raise WorkerHandlingException( - internal_message=f"wrong input mode: {mode}. Must be 'max' or 'freq'" - ) - - processed_mapping[reason] = { - "lvls": lvls, - "positive_lvls": positive_lvls, - "freqs": dict(counts), - "mapped_lvl": mapped_lvl, - } - return df, processed_mapping - - -def execute_level_mapping(df: pd.DataFrame, mapping: dict) -> pd.DataFrame: - def map_row(row): - reason = row["reason"] - if reason in mapping: - return mapping[reason]["mapped_lvl"] - return row["level"] - - df = df.copy() - origin_est_lvls = df["level"].tolist() - df["level"] = df.apply(map_row, axis=1) - df["origin_level"] = origin_est_lvls - return df - - -def extract_non_neg_code(reason_str: str) -> str: - """ - Extract the non-NEG code from reason_str (strip only NEG part, preserve META) - - Example: "POS [1, 0, 0] NEG [0, 0, 0] META [1, 2, 1]" -> "POS [1, 0, 0] META [1, 2, 1]" - Example: "3# AND POS [1, 0] NEG [0, 0]" -> "3# AND POS [1, 0]" - Example: "3# AND POS [1, 0] NEG [0, 0] META [1, 1, 0]" -> "3# AND POS [1, 0] META [1, 1, 0]" - """ - if not reason_str or not isinstance(reason_str, str): - return "" - neg_match = re.search(r"\s*NEG\s*\[[^\]]*\]", reason_str) - if neg_match: - # Remove only the NEG [...] part, keep everything before and after - before_neg = reason_str[: neg_match.start()] - after_neg = reason_str[neg_match.end() :] - return (before_neg + after_neg).strip() - return reason_str.strip() - - -def build_non_neg_mapping(lvl_mapping: dict) -> dict: - """ - Build non-NEG code mapping from complete lvl_mapping, select by highest frequency - - Args: - lvl_mapping: reason -> level mapping - - Returns: - non_neg_mapping: {non_neg_code: mapped_lvl} - """ - # collect all levels for each non_neg_code - non_neg_levels = {} - for reason, info in lvl_mapping.items(): - non_neg_code = extract_non_neg_code(reason) - mapped_lvl = info.get("mapped_lvl", -1) - if non_neg_code: - if non_neg_code not in non_neg_levels: - non_neg_levels[non_neg_code] = [] - non_neg_levels[non_neg_code].append(mapped_lvl) - - # select by highest frequency - non_neg_mapping = {} - for non_neg_code, levels in non_neg_levels.items(): - positive_levels = [lvl for lvl in levels if lvl > -1] - if positive_levels: - level_counts = Counter(positive_levels) - most_common_level = level_counts.most_common(1)[0][0] - non_neg_mapping[non_neg_code] = most_common_level - else: - non_neg_mapping[non_neg_code] = -1 - - return non_neg_mapping - - -def handle_unseen_codes( - df: pd.DataFrame, - level_dfs: list, - lvl_mapping: dict, - output_dir: str = None, - window_half_size: int = 10, - strategy: str = "double_mapping", -) -> dict: - """ - Handle unseen codes with configurable strategy - - Args: - df: original complete DataFrame - level_dfs: segment DataFrames - lvl_mapping: existing level mapping - output_dir: output directory (optional, only used for window_llm strategy) - window_half_size: window half size (how many rows above and below) - strategy: "double_mapping" or "window_llm" - - double_mapping: use non-neg code fallback (fast, no LLM call) - - window_llm: create windows for LLM to judge (slower, more accurate) - - Returns: - updated lvl_mapping - """ - - def extract_reason_signature(reason: str) -> str: - """Extract reason signature""" - return reason.strip() if reason else "" - - def has_neg_signal(reason_str: str) -> bool: - """Check if NEG signal exists (any value >= 1)""" - if not reason_str or not isinstance(reason_str, str): - return False - neg_match = re.search(r"NEG\s*\[([^\]]*)\]", reason_str) - if not neg_match: - return False - neg_content = neg_match.group(1) - try: - nums = [int(x.strip()) for x in neg_content.split(",") if x.strip()] - return any(x >= 1 for x in nums) - except Exception: - return False - - def build_context_window( - target_idx: int, known_codes_set: set, total_rows: int, half_size: int = 10 - ) -> dict: - """ - Build context window for unseen codes - 1. window size: half_size - 2. window should contain at least one known code - """ - min_start = max(0, target_idx - half_size) - min_end = min(total_rows - 1, target_idx + half_size) - - start_idx = min_start - end_idx = min_end - - found_known_above = False - found_known_below = False - known_positions = [] - - # check above - for i in range(start_idx, target_idx): - reason = df.iloc[i].get("reason", "") - sig = extract_reason_signature(reason) - if sig in known_codes_set: - found_known_above = True - known_positions.append(i) - - # check below - for i in range(target_idx + 1, end_idx + 1): - reason = df.iloc[i].get("reason", "") - sig = extract_reason_signature(reason) - if sig in known_codes_set: - found_known_below = True - known_positions.append(i) - - # expand above if needed - if not found_known_above and min_start > 0: - search_idx = min_start - 1 - while search_idx >= 0: - reason = df.iloc[search_idx].get("reason", "") - sig = extract_reason_signature(reason) - if sig in known_codes_set: - found_known_above = True - known_positions.append(search_idx) - start_idx = search_idx - break - search_idx -= 1 - - # expand below if needed - if not found_known_below and min_end < total_rows - 1: - search_idx = min_end + 1 - while search_idx < total_rows: - reason = df.iloc[search_idx].get("reason", "") - sig = extract_reason_signature(reason) - if sig in known_codes_set: - found_known_below = True - known_positions.append(search_idx) - end_idx = search_idx - break - search_idx += 1 - - return { - "start": start_idx, - "end": end_idx, - "found_known": found_known_above or found_known_below, - "known_positions": known_positions, - } - - # build non-neg mapping - non_neg_mapping = build_non_neg_mapping(lvl_mapping) - - # get known codes - known_codes = set(lvl_mapping.keys()) - - # record all codes from all segments. Placeholder rows (reason == - # PLACEHOLDER_REASON) are injected by _compact_for_llm and are never real - # heading candidates, so they must be skipped here — otherwise they would - # show up as an "unseen code" and fall through to NO_MATCH_FALLBACK, adding - # harmless but noisy warnings to the log. - all_codes_in_full = {} - for seg_idx, seg_df in enumerate(level_dfs): - for _, row in seg_df.iterrows(): - reason = row.get("reason", "") - sig = extract_reason_signature(reason) - if not sig or sig == PLACEHOLDER_REASON: - continue - if sig not in all_codes_in_full: - all_codes_in_full[sig] = { - "first_seg": seg_idx, - "first_id": row.get("id", 0), - "reason": reason, - } - - # find unseen codes - unseen_codes = {} - unseen_neg_filtered = {} - for sig, info in all_codes_in_full.items(): - if sig in known_codes: - continue - if has_neg_signal(info["reason"]): - unseen_neg_filtered[sig] = info - else: - unseen_codes[sig] = info - - logger.info( - f"Unseen codes total: {len(unseen_codes) + len(unseen_neg_filtered)}, NEG filtered: {len(unseen_neg_filtered)}, to process: {len(unseen_codes)}" - ) - - # if neg signal, map to -1 - for sig in unseen_neg_filtered: - lvl_mapping[sig] = {"mapped_lvl": -1, "note": "NEG_FILTERED"} - - # handle remaining unseen_codes based on strategy - if unseen_codes: - if strategy == "double_mapping": - # Strategy 1: use non-neg code fallback - fallback_success = 0 - fallback_failed = 0 - failed_codes = [] - for sig, info in unseen_codes.items(): - non_neg_code = extract_non_neg_code(sig) - if non_neg_code in non_neg_mapping: - mapped_level = non_neg_mapping[non_neg_code] - lvl_mapping[sig] = { - "mapped_lvl": mapped_level, - "note": f"NON_NEG_FALLBACK from '{non_neg_code}'", - } - fallback_success += 1 - else: - lvl_mapping[sig] = {"mapped_lvl": -1, "note": "NO_MATCH_FALLBACK"} - fallback_failed += 1 - failed_codes.append( - f"'{non_neg_code}' (from '{sig[:60]}...')" - if len(sig) > 60 - else f"'{non_neg_code}' (from '{sig}')" - ) - - logger.info( - f"Double mapping result: success={fallback_success}, failed={fallback_failed}" - ) - if failed_codes: - logger.warning( - f"Failed codes (non_neg not in mapping): {failed_codes[:5]}{'...' if len(failed_codes) > 5 else ''}" - ) - - elif strategy == "window_llm" and output_dir: - # Strategy 2: create windows for LLM to judge - total_rows = len(df) - windows = [] - for sig, info in unseen_codes.items(): - first_id = info["first_id"] - first_seg = info["first_seg"] - df_indices = df.index[df["id"] == first_id].tolist() - if df_indices: - first_df_idx = df_indices[0] - window_info = build_context_window( - first_df_idx, known_codes, total_rows, window_half_size - ) - windows.append( - { - "code": sig, - "first_id": first_id, - "first_seg": first_seg, - "start": window_info["start"], - "end": window_info["end"], - "found_known": window_info["found_known"], - } - ) - - # merge windows - sorted_windows = sorted(windows, key=lambda x: x["start"]) - merged_windows = [] - current_window = None - - for w in sorted_windows: - if current_window is None: - current_window = { - "start": w["start"], - "end": w["end"], - "codes": [w["code"]], - "segments": [w["first_seg"]], - } - elif w["start"] <= current_window["end"]: - current_window["end"] = max(current_window["end"], w["end"]) - current_window["codes"].append(w["code"]) - current_window["segments"].append(w["first_seg"]) - else: - merged_windows.append(current_window) - current_window = { - "start": w["start"], - "end": w["end"], - "codes": [w["code"]], - "segments": [w["first_seg"]], - } - - if current_window: - merged_windows.append(current_window) - - # save windows - windows_dir = os.path.join(output_dir, "merged_windows") - os.makedirs(windows_dir, exist_ok=True) - - unseen_codes_set = set(unseen_codes.keys()) - unseen_neg_set = set(unseen_neg_filtered.keys()) - - for i, mw in enumerate(merged_windows): - window_df = df.iloc[mw["start"] : mw["end"] + 1].copy() - - def get_code_status(row): - reason = row.get("reason", "") - sig = extract_reason_signature(reason) - if not sig: - return "" - if sig in unseen_codes_set: - return "★ UNSEEN_TARGET" - elif sig in unseen_neg_set: - return "NEG→-1" - elif sig in known_codes: - return "KNOWN" - else: - return "" - - window_df["code_status"] = window_df.apply(get_code_status, axis=1) - window_path = os.path.join( - windows_dir, - f"window_{i + 1:02d}_rows_{mw['start']}-{mw['end']}.csv", - ) - window_df.to_csv(window_path, index=False, encoding="utf-8-sig") - - logger.debug( - f"Window LLM: {len(merged_windows)} windows created in {windows_dir}" - ) - # TODO: use llm to assign level based on window data - - return lvl_mapping - - -def detect_outlines_md(line): - pos_code = judge_by_conditions(line) - any(x > 0 for x in pos_code) - - -def get_max_lvl(code_str: str): - match = re.search(r"\[([^]]+)]", code_str) - if not match: - return "Sure" - - nums = [int(x.strip()) for x in match.group(1).split(",")] - max_val = int(max(nums)) - return max_val if max_val > 1 else -2 # -2 = "Not Sure" sentinel (int-safe) - - -PLACEHOLDER_REASON = "__PLACEHOLDER__" - - -def _compact_for_llm(df: pd.DataFrame) -> pd.DataFrame: - """Collapse consecutive ``level == -1`` rows into a single placeholder row. - - Rows with ``level >= 1`` (heading candidates) and ``level == -2`` ("Not Sure") - are preserved verbatim so the LLM can still judge them. Each run of - consecutive ``-1`` rows becomes one placeholder row whose: - - id = "start-end" (always a range; "N-N" when the run is one row) - heading = "[N BODY LINES]" where N is the run length - level = "-" - reason = ``PLACEHOLDER_REASON`` - - The id is ALWAYS a hyphenated string, even for single-row runs, so that - ``int(id)`` fails for every placeholder. This lets downstream code identify - placeholders structurally (non-integer id) without depending on ``reason`` - or length heuristics. - """ - if df is None or len(df) == 0: - return pd.DataFrame(columns=["id", "heading", "level", "reason"]) - - rows = [] - i = 0 - n = len(df) - while i < n: - lvl_raw = df.iloc[i]["level"] - try: - lvl_int = int(lvl_raw) - except (TypeError, ValueError): - lvl_int = None - - if lvl_int == -1: - j = i - while j < n: - try: - nxt_lvl = int(df.iloc[j]["level"]) - except (TypeError, ValueError): - break - if nxt_lvl != -1: - break - j += 1 - start_id = int(df.iloc[i]["id"]) - end_id = int(df.iloc[j - 1]["id"]) - run = j - i - rows.append( - { - "id": f"{start_id}-{end_id}", - "heading": f"[{run} BODY LINES]", - "level": "-", - "reason": PLACEHOLDER_REASON, - } - ) - i = j - else: - r = df.iloc[i] - rows.append( - { - "id": int(r["id"]), - "heading": str(r["heading"]), - "level": ( - int(lvl_int) - if lvl_int is not None and lvl_int != -2 - else "Not Sure" - ), - "reason": str(r.get("reason", "") or ""), - } - ) - i += 1 - - return pd.DataFrame(rows, columns=["id", "heading", "level", "reason"]) - - -def heading_tb_transfer(df, threshold=3000, max_start=50, max_end=10): - raw_headings = df["heading"].tolist() - df["heading"] = df["heading"].apply( - lambda x: truncate_text_by_tokens(x, max_start, max_end) - ) - - sub_dfs = [] - current_rows = [] - current_len = 0 - for _, row in df.iterrows(): - row_filtered = row.drop(labels=["reason"], errors="ignore") - row_len = sum(count_cn_en(str(v)) for v in row_filtered.values) - - if current_len + row_len > threshold and current_rows: - sub_dfs.append(pd.DataFrame(current_rows, columns=df.columns)) - current_rows = [row.tolist()] - current_len = row_len - else: - current_rows.append(row.tolist()) - current_len += row_len - - if current_rows: - sub_dfs.append(pd.DataFrame(current_rows, columns=df.columns)) - return sub_dfs, raw_headings - - -def judge_by_conditions(text, scope=20, return_detail=False, CN_SPECIAL_IDX=12): - """ - judge level features as one-hot embeddings for texts - - Args: - text: input text - scope: text scope for judging - return_detail: whether to return detailed information (including unit type) - CN_SPECIAL_IDX: index of special Chinese number - - Returns: - if return_detail=False: return pos_triggered_code list - if return_detail=True: return (pos_triggered_code, detail_info) tuple - where detail_info is a dictionary containing additional information, such as Chinese unit type - """ - text = text.replace("\u3000", " ") - text = unicodedata.normalize("NFKC", text)[:scope] - - # ========== English Numbering ========== - regex_en_num_dots = r"^\d+(?:\s*\.\s*\d+)+(?![、，。！？；：])(?=\s|$|\w|[一-龥])" - regex_en_num_dun = r"^\d、\s{0,4}(?=\S|$)" # 1、xxx - regex_en_num_dots_dun = r"^\d+(?:\.\d+)*、\s*(?=[A-Za-z一-龥])" - regex_en_num_single_dot = r"^\d+\.(?!\d)\s{0,4}(?=\S)" # 1.xxx - regex_en_num_space = r"^[0-9]{1,2}\s{1,8}(?=\S)" # 1 xxx - # ========== Chinese Numbering ========== - regex_cn_num_dun = r"^[一二三四五六七八九十百千万]+、\s{0,4}(?=\S|$)" - regex_cn_num_mix = ( - r"^[一二三四五六七八九十百千万]+(?:\s*\.[一二三四五六七八九十百千万\d]+)+" - ) - regex_cn_num_plain = r"^[一二三四五六七八九十百千万]+(?=\s|$)" - # ========== English Bracketing ========== - regex_en_brac_paren = r"^[\(\（]\s*\d+(?:\.\d+)*(?!\.0)\s*[\)\）]" - regex_en_brac_right = r"^\d+(?:\.\d+)*(?!\.0)\s*[\)\）]" - # ========== Chinese Bracketing ========== - regex_cn_brac_paren = r"^[\(\（]\s*[一二三四五六七八九十百千万]+(?:\.[一二三四五六七八九十百千万\d]+)*\s*[\)\）]" - regex_cn_brac_right = r"^[一二三四五六七八九十百千万]+(?:\.[一二三四五六七八九十百千万\d]+)*\s*[\)\）]" - # ========== Chinese Special ========== - regex_cn_special = r"^第[一二三四五六七八九十百千万\d]+(?:\.[一二三四五六七八九十百千万\d]+)*(章|节|条|部分|款|目|项|编|篇|卷|辑)?(?=$|\s|[A-Za-z0-9\u4e00-\u9fa5])" - # ========== English Letter Numbering ========== - regex_letter_dot = r"^[A-Za-z](?:\.\d+)*[\.、](?=\s*\S)" - regex_letter_brac_paren = r"^[\(\（]\s*[A-Za-z](?:\.\d+)*(?!\.0)\s*[\)\）]" - regex_letter_brac_right = r"^[A-Za-z](?:\.\d+)*(?!\.0)\s*[\)\）]" - # ========== Appendix ========== - regex_appendix = r"^((附件|附录|附表|附图)|(?i:appendix))[\s_\-—]{0,4}(?:\[)?[一二三四五六七八九十A-Za-z\d]" - - pos_regex_conditions = [ - # English Numbering - regex_en_num_dots, - regex_en_num_dun, - regex_en_num_single_dot, - regex_en_num_space, - regex_en_num_dots_dun, - # Chinese Numbering - regex_cn_num_dun, - regex_cn_num_mix, - regex_cn_num_plain, - # English Bracketing - regex_en_brac_paren, - regex_en_brac_right, - # Chinese Bracketing - regex_cn_brac_paren, - regex_cn_brac_right, - # Chinese Special - regex_cn_special, - # English Letter Numbering - regex_letter_dot, - regex_letter_brac_paren, - regex_letter_brac_right, - # Appendix - regex_appendix, - ] - - pos_triggered_code = [] - reason_suffix_parts = [] - - for idx, regex in enumerate(pos_regex_conditions): - match = re.match(regex, text) - if match: - matched_text = match.group(0) - symbols = ".-" - count_ = sum(matched_text.count(s) for s in symbols) + 1 - - # Special handling for Chinese chapter/section/item markers. - if idx == CN_SPECIAL_IDX and return_detail: - unit_match = re.search( - r"(章|节|条|部分|款|目|项|编|篇|卷|辑)", matched_text - ) - if unit_match: - unit = unit_match.group(1) - reason_suffix_parts.append(f"[CN:{unit}]") - pos_triggered_code.append(count_) - else: - pos_triggered_code.append(0) - - if return_detail: - detail_info = { - "reason_suffix": ( - " ".join(reason_suffix_parts) if reason_suffix_parts else "" - ) - } - if detail_info["reason_suffix"]: - detail_info["reason_suffix"] = " " + detail_info["reason_suffix"] - return pos_triggered_code, detail_info - return pos_triggered_code - - -def remove_by_conditions(text, include_punc=False): - neg_condition_num = r"^\d{3,}" - neg_condition_zero = r"^0\.\d+[\u4e00-\u9fa5A-Za-z\S]*" # 0.2xxx - neg_decimal_only = r"^\d*\.\d+$" # 0.2 .23 - neg_condition_http = ( - r"(?i)(^https?://\S+|^www\.\S+|^P\.S|^\b\d{0,2}\s*(?:a\.m|p\.m)\b)" - ) - # LaTeX: wrapped ($..\cmd..$) OR bare commands (\times, \mathrm, etc.) - neg_condition_latex = ( - r"(?:" - r"\$[^$]*\\[A-Za-z]+(?:\s*\{[^{}]*\})?[^$]*\$" # wrapped: $...\cmd...$ - r"|" - r"\\(?:times|div|cdot|pm|mp|leq|geq|neq|approx|equiv|sim|infty" - r"|sum|prod|int|sqrt|frac|mathrm|mathbf|mathit|mathcal" - r"|text(?:bf|it|rm)?|alpha|beta|gamma|delta|epsilon|theta" - r"|lambda|mu|sigma|pi|omega|partial|nabla" - r"|left|right|begin|end|overline|underline|hat|vec|tilde)\b" - r")" - ) - # Number immediately followed by measurement unit (e.g. 25.40mm, 100kPa) - neg_condition_unit = ( - r"^\d+\.?\d*\s{0,2}" - r"(?:mm|cm|km|nm|μm|inch(?:es)?|ft|yd|mi" - r"|kg|mg|μg|lb|oz" - r"|kPa|MPa|GPa|Pa|psi|bar" - r"|°[CFK]" - r"|Hz|kHz|MHz|GHz" - r"|mol|mL|dL|dB|Nm|kN|MN|kW|MW|GW|hp|rpm|cc|cal|kcal)\b" - ) - neg_condition_punc_mid = r"[。！；].+" - neg_condition_punc_end = r"[.,;，。；]$" - - neg_conditions = [ - neg_condition_num, - neg_condition_http, - neg_condition_latex, - neg_condition_zero, - neg_decimal_only, - neg_condition_punc_mid, - neg_condition_unit, - ] - - neg_triggered_code = [] - for regex in neg_conditions: - match = re.search(regex, text) - neg_triggered_code.append(1 if match else 0) - - if include_punc: - match = re.search(neg_condition_punc_end, text) - neg_triggered_code.append(1 if match else 0) - else: - neg_triggered_code.append(0) - - return neg_triggered_code - - -def md_heading_match(line, as_is=True): - """handle markdown headings, considering # < ! [....""" - match = re.match(r"^\s*(#+)\s*(.*)$", line) - if match: - level = len(match.group(1)) # count the number of '#' - if as_is: # determine if remove the '#' - return line, level - else: - return line.lstrip("#").strip(), level - else: - return line, -1 - - -def filter_md_headings(md_lines, num_pos=17, num_neg=7, layout_json_path=None): - """filter candidate headings for .md - - Args: - md_lines: list of markdown lines - num_pos: number of positive conditions - num_neg: number of negative conditions - layout_json_path: optional path to layout.json for META features (size ranking) - """ - # Create MetadataContext if layout_json_path is provided - meta_ctx = None - if layout_json_path: - try: - from .metadata_extractor import MetadataContext - - meta_ctx = MetadataContext(md_lines, layout_json_path) - except Exception as e: - logger.warning(f"Failed to create MetadataContext: {e}") - - raw_candidates = [] - for i, line in enumerate(md_lines): - line = line.strip() - if not line: - continue - - if ( - ("" in line) # annotation line - or line.startswith("|") # table line - or line.startswith("") - or "![" in line - and "](" in line # image line - ): - est_lvl = -1 - zero_pos_code = [0] * num_pos - zero_neg_code = [0] * num_neg - str_lvl = f"POS {zero_pos_code} NEG {zero_neg_code}" - if meta_ctx: - str_lvl += " META [0, 0, 0]" - line = "Figure/Image" - else: - line_clean, hash_lvl = md_heading_match( - line, as_is=False - ) # detect "#" in .md lines - - # NEW: detect and strip full-line bold markers (e.g. **3.4 Title** -> 3.4 Title) - from .metadata_extractor import detect_and_strip_md_bold - - line_clean_stripped, is_full_bold = detect_and_strip_md_bold(line_clean) - - # Use stripped text for POS/NEG analysis (fixes '**3.4' -> '3.4' issue) - pos_code, detail_info = judge_by_conditions( - line_clean_stripped, return_detail=True - ) - neg_code = remove_by_conditions(line_clean_stripped) - - if any(x > 0 for x in neg_code): - code_lvl = -1 - code_str = f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}" - - elif any(x > 0 for x in pos_code) and all(x == 0 for x in neg_code): - code_lvl = get_max_lvl(str(pos_code)) - code_str = f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}" - - else: - code_lvl = -1 - code_str = f"POS {pos_code} NEG {neg_code}" - - # Add META suffix with bold dimension - if meta_ctx: - size_rank, occurrence = meta_ctx.get_meta_for_line(line_clean) - is_bold_int = 1 if is_full_bold else 0 - code_str += meta_ctx.format_meta_suffix( - size_rank, occurrence, is_bold_int - ) - else: - # Even without layout.json, output bold info in META - if is_full_bold: - code_str += " META [0, 0, 1]" - - if hash_lvl <= 0: - est_lvl = code_lvl - str_lvl = code_str - else: - if isinstance(code_lvl, int): - est_lvl = max( - hash_lvl, code_lvl - ) # current miner tend to produce fewer #s - else: - est_lvl = code_lvl # code_lvl could be not sure - str_lvl = f"{hash_lvl}# AND {code_str}" - raw_candidates.append((i, line, est_lvl, str_lvl)) - - preds_df = pd.DataFrame( - raw_candidates, columns=["id", "heading", "level", "reason"], index=None - ) - return preds_df - - -def filter_doc_headings(titles_material, enable_regx=True, enable_style_check=False): - """filter candidate headings for docx""" - - def find_docstyle(para_): - try: - style_name = para_.style.name - except Exception: - style_name = "normal" - if style_name.startswith("Heading") or style_name.startswith("标题"): - try: - outline_level = int(style_name.split(" ")[1]) - except Exception: - outline_level = -2 # "Not Sure" sentinel - return outline_level - else: - return None - - def find_otsetting(para_): - ppr = para_._element.find(qn("w:pPr")) - if ppr is not None: - plvl = ppr.find(qn("w:outlineLvl")) - else: - return None - - if plvl is not None: - outline_level = int(plvl.get(qn("w:val"))) + 1 - return outline_level - else: - return None - - def find_bold(para_): - if para_.runs and all(run.bold for run in para_.runs if run.text.strip()): - return True - else: - return None - - raw_candidates = [] - logger.debug( - "Filtering docx heading candidates... total_items={}", len(titles_material) - ) - for ele_id, para, text in titles_material: - str_lvl = "" - est_lvl = None - style_lvl = find_docstyle(para) - setting_lvl = find_otsetting(para) - - # 1. check .docx style settings - if style_lvl is not None: - est_lvl = style_lvl - str_lvl = f"style-{style_lvl}" - - # 2. check .docx paragraph numbering settings - elif setting_lvl is not None: - est_lvl = setting_lvl - str_lvl = f"outline-{setting_lvl}" - - # 3. detect bold (unconditionally, encode as META dimension) - is_bold = 1 if find_bold(para) else 0 - - # 4. proceed condition judge - if enable_regx: - pos_code, detail_info = judge_by_conditions(text, return_detail=True) - neg_code = remove_by_conditions(text) - - if any(x > 0 for x in neg_code): - code_lvl = -1 - code_str = f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}" - elif any(x > 0 for x in pos_code) and all(x == 0 for x in neg_code): - code_lvl = get_max_lvl(str(pos_code)) - code_str = f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}" - else: - code_lvl = -1 - code_str = f"POS {pos_code} NEG {neg_code}" - - # Append bold as META dimension (DOCX has no layout.json, so only bold) - if is_bold: - code_str += f" META [0, 0, {is_bold}]" - - if est_lvl is None: - est_lvl = code_lvl - str_lvl = code_str - else: - str_lvl = f"{str_lvl} AND {code_str}" - raw_candidates.append((ele_id, text, est_lvl, str_lvl)) - - preds_df = pd.DataFrame( - raw_candidates, columns=["id", "heading", "level", "reason"], index=None - ) - - # initial merge isolated and short texts - preds_df = postprocess_headings(preds_df, task="merge_continuous") - preds_df = postprocess_headings(preds_df, task="merge_short") - return preds_df - - -def format_toc_context_for_llm(toc_context) -> str: - """Convert TOC hierarchy or structured payloads into compact LLM-friendly plain text.""" - if not toc_context: - return "" - - if isinstance(toc_context, str): - return toc_context - - toc_items = toc_context if isinstance(toc_context, list) else [toc_context] - formatted_blocks = [] - - for toc_idx, toc_item in enumerate(toc_items, start=1): - if not isinstance(toc_item, dict): - formatted_blocks.append(str(toc_item)) - continue - - toc_range = toc_item.get("toc_range") - toc_entries = toc_item.get("toc_with_level") or [] - - if toc_range and len(toc_range) == 2: - formatted_blocks.append( - f"TOC {toc_idx} (source rows {toc_range[0]}-{toc_range[1]}):" - ) - else: - formatted_blocks.append(f"TOC {toc_idx}:") - - if not toc_entries: - formatted_blocks.append("- No TOC entries available") - continue - - if isinstance(toc_entries, str): - toc_entries = toc_entries.strip() - if toc_entries: - formatted_blocks.append(toc_entries) - else: - formatted_blocks.append("- No TOC entries available") - continue - - for entry in toc_entries: - if not isinstance(entry, dict): - continue - - heading = str(entry.get("heading", "")).strip().replace("\n", " ") - if not heading: - continue - - level = entry.get("level") - line_id = entry.get("id") - if isinstance(level, int): - formatted_blocks.append(f"- level {level} | id {line_id} | {heading}") - else: - formatted_blocks.append(f"- id {line_id} | {heading}") - - return "\n".join(formatted_blocks) - - -def hiearchy_llm( - df, - model_name=None, - max_depth=6, - toc_context=None, - max_len=8192, - task="eval-headings", -): - """Apply LLM to analyze the hierarchy of headings - - Args: - df: DataFrame with id, heading columns - model_name: LLM model name (optional, uses default if None) - max_depth: Maximum hierarchy depth - max_len: Hard cap for LLM completion max_tokens (default 2048). - Actual value is derived from the number of heading candidates. - task: Prompt task type - "eval-headings" for general document, "eval-toc-headings" for TOC - toc_context: Optional formatted TOC context string for guiding level assignment - - Returns: - List of dicts with id and level, one per row in ``df`` (missing IDs -> level=-1). - """ - - model_name = _resolve_hierarchy_model_name(model_name) - level_md = df2md(df) - - # Completion budget is driven by the number of heading candidates, not the - # markdown input length. Each JSON entry is `{"id":X,"level":Y}` ≈ 25 tokens; - # add 200 tokens overhead for brackets/whitespace and leave a 512 floor for - # tiny inputs. Non-int ids (placeholders like "10-12" or "-") are excluded. - def _is_candidate_id(val): - if isinstance(val, bool): - return False - if isinstance(val, int): - return True - try: - int(val) - return True - except (TypeError, ValueError): - return False - - n_candidates = int(df["id"].apply(_is_candidate_id).sum()) if len(df) > 0 else 0 - ot_limit = max(512, n_candidates * 25 + 200) - ot_limit = min(ot_limit, max_len) - formatted_toc_context = format_toc_context_for_llm(toc_context) - - paras = { - "max_tokens": ot_limit, - "max_depth": max_depth, - "toc_context": formatted_toc_context, - } - prompt, temperature, top_p, max_tokens = build_prompt( - task=task, texts=level_md, query="", paras=paras - ) - messages = [ - {"role": "system", "content": "you are a document auditing expert"}, - {"role": "user", "content": prompt}, - ] - - try: - with stage_timer( - "heading.hierarchy_llm_call", - model_name=model_name, - row_count=len(df), - task=task, - candidate_count=n_candidates, - max_tokens=max_tokens, - ): - answer = get_openai_client(model=model_name).chat_completion( - messages=messages, - model=model_name, - max_tokens=max_tokens, - temperature=temperature, - ) - layout_res = eval_response(answer) - - # Validate eval_response result — it can return a raw string when JSON parsing fails - if not isinstance(layout_res, list): - raise ValueError( - f"LLM returned non-list response (type={type(layout_res).__name__}), " - f"raw content: {str(layout_res)[:200]}" - ) - - # Validate each item is a dict with required keys - for i, item in enumerate(layout_res): - if not isinstance(item, dict) or "id" not in item or "level" not in item: - raise ValueError(f"LLM response item[{i}] is malformed: {item!r}") - - # Drop items whose id is not a clean integer. This includes placeholder - # rows ("10-12", "-") that the LLM may echo back despite the prompt telling - # it not to. - clean_res = [] - dropped = 0 - for item in layout_res: - raw_id = item["id"] - if isinstance(raw_id, bool): - dropped += 1 - continue - if isinstance(raw_id, int): - clean_res.append({"id": raw_id, "level": item["level"]}) - continue - try: - clean_res.append({"id": int(raw_id), "level": item["level"]}) - except (TypeError, ValueError): - dropped += 1 - if dropped: - logger.debug(f"filtered {dropped} non-integer-id entries from LLM response") - - # LLM only returns heading rows (level >= 1). Reconstruct full result so the - # returned list has one entry per row in ``df``, with missing ids defaulting - # to level=-1. Rows whose ``id`` is itself non-integer (placeholders) keep - # their id as-is and level=-1 so the caller can filter them out. - llm_levels = {item["id"]: item["level"] for item in clean_res} - full_result = [] - for row_id in df["id"].tolist(): - if _is_candidate_id(row_id): - try: - int_id = int(row_id) - except (TypeError, ValueError): - int_id = row_id - full_result.append({"id": int_id, "level": llm_levels.get(int_id, -1)}) - else: - full_result.append({"id": row_id, "level": -1}) - logger.debug( - f"LLM returned {len(clean_res)} heading levels out of {n_candidates} candidates " - f"({len(df)} total rows)" - ) - return full_result - except Exception as e: - logger.error(f"detect hierarchy by LLM failed: {e}") - raise - - -def _compute_zone_boundaries(toc_hierarchies, coordinate_mode="post_removal"): - """Compute content zone boundaries for documents with multiple TOC areas. - - When multiple TOCs exist, they divide the document into zones. Each zone - starts right after a TOC area and extends to just before the next TOC area - (or end of document). - - coordinate_mode: - - "post_removal": TOC ranges are in original coordinates, but heading IDs - are measured after TOC rows were removed (MD/PDF path). - - "original": heading IDs stay in original document coordinates, so zones - can be computed directly from TOC boundaries (DOCX path). - - Args: - toc_hierarchies: List of toc hierarchy dicts (sorted by toc_range start) - - Returns: - List of (zone_start_post, zone_end_post_or_None, toc_hierarchy_dict) - zone_end_post is None for the last zone (extends to end of document) - """ - if coordinate_mode not in {"post_removal", "original"}: - raise ValueError(f"Unsupported coordinate_mode: {coordinate_mode}") - - sorted_tocs = sorted(toc_hierarchies, key=lambda t: t["toc_range"][0]) - - zones = [] - cumulative_removed = 0 - - for i, toc in enumerate(sorted_tocs): - toc_start, toc_end = toc["toc_range"] - zone_start = toc_end + 1 - - if coordinate_mode == "post_removal": - toc_size = toc_end - toc_start + 1 - cumulative_removed += toc_size - zone_start -= cumulative_removed - - if i + 1 < len(sorted_tocs): - next_toc_start = sorted_tocs[i + 1]["toc_range"][0] - zone_end = next_toc_start - 1 - if coordinate_mode == "post_removal": - zone_end -= cumulative_removed - else: - zone_end = None # to end of document - - if zone_end is not None and zone_end < zone_start: - continue - zones.append((zone_start, zone_end, toc)) - - return zones - - -def _resolve_first_toc_boundary(toc_hierarchies=None, first_toc_ele_num=None): - """Resolve the earliest available first-TOC boundary across coordinate sources.""" - toc_range_start = None - if toc_hierarchies: - first_range = toc_hierarchies[0].get("toc_range") - if first_range and len(first_range) == 2: - toc_range_start = first_range[0] - - candidates = [ - value for value in (toc_range_start, first_toc_ele_num) if value is not None - ] - if not candidates: - return None - - resolved_start = min(candidates) - if ( - toc_range_start is not None - and first_toc_ele_num is not None - and toc_range_start != first_toc_ele_num - ): - logger.info( - f"📌 TOC boundary mismatch detected: toc_range start={toc_range_start}, " - f"first_toc_ele_num={first_toc_ele_num}, using earliest={resolved_start}" - ) - return resolved_start - - -def pred_titles( - infos, - doc_type, - toc_hierarchies=None, - prompt_limt=4000, - enable_regx=True, - smart_parse=False, - model_name=None, - output_dir=None, - layout_json_path=None, - first_toc_ele_num=None, -): - """ - predict title hierarchy - - Args: - infos: document information - doc_type: document type (pptx, md, docx) - toc_hierarchies: TOC hierarchy information (if any) - prompt_limt: prompt character limit - enable_regx: whether to enable regex matching - smart_parse: whether to use LLM intelligent parsing - model_name: LLM model name - output_dir: output directory for saving intermediate CSV results - layout_json_path: path to layout.json for META features (optional) - first_toc_ele_num: ele_num of the first TOC block in DOCX (for pre-TOC exclusion) - """ - model_name = _resolve_hierarchy_model_name(model_name) - logger.info( - f"Start to predict title hierarchy: doc_type={doc_type}, smart_parse={smart_parse}, candidate titles={len(infos)}" - ) - - if doc_type == "pptx": - raw_preds = filter_md_headings(infos) - elif doc_type == "md": - raw_preds = filter_md_headings(infos, layout_json_path=layout_json_path) - elif doc_type == "docx": - raw_preds = filter_doc_headings(infos, enable_regx) - else: - raw_preds = pd.DataFrame(columns=["id", "heading", "level", "reason"]) - - # ── Exclude pre-TOC lines from heading prediction ── - # When TOC is detected, lines/blocks before the first TOC area are typically - # cover/metadata (company name, version, classification marks), not real - # headings. Remove them before LLM judging to avoid misjudgment, then - # splice back with level=-1 after all processing is done. - pre_toc_rows = None - first_toc_start = None - if doc_type == "md": - first_toc_start = _resolve_first_toc_boundary(toc_hierarchies=toc_hierarchies) - elif doc_type == "docx": - first_toc_start = _resolve_first_toc_boundary( - toc_hierarchies=toc_hierarchies, - first_toc_ele_num=first_toc_ele_num, - ) - - if first_toc_start is not None and first_toc_start > 0: - pre_toc_mask = raw_preds["id"] < first_toc_start - if pre_toc_mask.any(): - pre_toc_rows = raw_preds[pre_toc_mask].copy() - pre_toc_rows["level"] = -1 - raw_preds = raw_preds[~pre_toc_mask].reset_index(drop=True) - if doc_type == "docx": - logger.info( - f"📌 Excluded {len(pre_toc_rows)} pre-TOC blocks " - f"(id < {first_toc_start}) from heading prediction" - ) - else: - logger.info( - f"📌 Excluded {len(pre_toc_rows)} pre-TOC lines " - f"(id < {first_toc_start}) from heading prediction" - ) - - # 2. Zone-based prediction when multiple TOCs exist - if ( - toc_hierarchies - and len(toc_hierarchies) > 1 - and doc_type in {"md", "docx"} - and smart_parse - ): - # Multiple TOCs divide the document into independent zones. - # Each zone gets its own naive + LLM pipeline with zone-specific TOC context. - coordinate_mode = "post_removal" if doc_type == "md" else "original" - zones = _compute_zone_boundaries( - toc_hierarchies, coordinate_mode=coordinate_mode - ) - logger.info( - f"🗂️ Zone-based prediction: {len(zones)} zones from {len(toc_hierarchies)} TOCs" - ) - - def _process_single_zone(zone_idx, zone_start, zone_end, zone_toc): - """Process a single zone independently. Returns (zone_idx, zone_heading_df).""" - # Extract rows belonging to this zone - if zone_end is not None: - zone_mask = (raw_preds["id"] >= zone_start) & ( - raw_preds["id"] <= zone_end - ) - else: - zone_mask = raw_preds["id"] >= zone_start - zone_preds = raw_preds[zone_mask].copy().reset_index(drop=True) - - if zone_preds.empty: - logger.warning(f" Zone {zone_idx}: empty, skipping") - return zone_idx, None - - zone_range_str = f"[{zone_start}, {zone_end or 'end'}]" - logger.info( - f" Zone {zone_idx}: {len(zone_preds)} rows, post-removal range {zone_range_str}" - ) - - # Independent naive + LLM prediction for this zone - zone_heading = est_hierarchies_naive( - zone_preds, smart_parse, output_dir=output_dir - ) - zone_heading = est_hierarchies_llm( - zone_heading, - prompt_limt, - toc_hierarchies=[zone_toc], # Single TOC for this zone - model_name=model_name, - output_dir=output_dir, - csv_suffix=f"_zone_{zone_idx}", - ) - valid_count = ( - len(zone_heading[zone_heading["level"] > 0]) - if not zone_heading.empty - else 0 - ) - logger.info(f" Zone {zone_idx}: ✅ {valid_count} valid headings") - return zone_idx, zone_heading - - if len(zones) == 1: - # Single zone: no parallel overhead - zone_start, zone_end, zone_toc = zones[0] - _, zone_heading = _process_single_zone(0, zone_start, zone_end, zone_toc) - zone_results = [zone_heading] if zone_heading is not None else [] - else: - # Multiple zones: parallel hierarchy prediction via gevent - logger.info( - f"Parallelizing zone hierarchy prediction for {len(zones)} zones" - ) - pool = GeventPool(size=len(zones)) - greenlets = [ - pool.spawn( - _process_single_zone, zone_idx, zone_start, zone_end, zone_toc - ) - for zone_idx, (zone_start, zone_end, zone_toc) in enumerate(zones) - ] - gevent.joinall(greenlets) - - # Collect results sorted by zone index to maintain document order - results = sorted( - [g.value for g in greenlets if g.value is not None], key=lambda r: r[0] - ) - zone_results = [ - heading_df for _, heading_df in results if heading_df is not None - ] - - if zone_results: - heading_preds = ( - pd.concat(zone_results, ignore_index=True) - .sort_values("id") - .reset_index(drop=True) - ) - else: - heading_preds = pd.DataFrame(columns=["id", "heading", "level", "reason"]) - logger.info("✅ Zone-based LLM hierarchy parsing completed") - else: - # Single-zone: current behavior - heading_preds = est_hierarchies_naive( - raw_preds, smart_parse, output_dir=output_dir - ) - if smart_parse: - heading_preds = est_hierarchies_llm( - heading_preds, - prompt_limt, - toc_hierarchies, - model_name=model_name, - output_dir=output_dir, - ) - logger.info("✅ LLM hierarchy parsing completed") - - # 3. final polishing for certain types - if doc_type in ["docx"]: - heading_preds = postprocess_headings(heading_preds, task="merge_continuous") - heading_preds = postprocess_headings(heading_preds, task="merge_short") - heading_preds = postprocess_headings(heading_preds, task="judge_negs") - logger.debug("Docx hiearchy detection postprocessing completed") - - if heading_preds["level"].eq(-1).all(): # if non are estimated as headings - logger.warning("⚠️ No valid headings estimated") - heading_preds = pd.DataFrame() - else: - heading_preds["level"] = ( - pd.to_numeric(heading_preds["level"], errors="coerce") - .fillna(-1) - .astype(int) - ) - - # process isolated nodes - try: - tree, node_to_id, _ = build_tree_from_dataframe(heading_preds) - processed_tree = remove_isolated_nodes(tree) - heading_preds = tree_to_dataframe(processed_tree, node_to_id, heading_preds) - except Exception as e: - logger.warning(f"Tree structure optimization failed, skipping: {e}") - - logger.info( - f"✅ Heading parsing completed, final {len(heading_preds[heading_preds['level'] > 0])} valid headings" - ) - - # ── Splice pre-TOC rows back ── - if pre_toc_rows is not None and not heading_preds.empty: - heading_preds = ( - pd.concat( - [pre_toc_rows[["id", "heading", "level", "reason"]], heading_preds], - ignore_index=True, - ) - .sort_values("id") - .reset_index(drop=True) - ) - logger.debug( - f"📌 Spliced {len(pre_toc_rows)} pre-TOC lines back into predictions" - ) - - # Save heading_preds as preds_5 - save_intermediate_csv(heading_preds, output_dir, "preds_5_final_output") - return heading_preds - - -def est_hierarchies_naive(raw_preds, proceed_smart=True, output_dir=None): - """Detect hierarchies by non-LLM - - Args: - raw_preds: raw data - proceed_smart: whether to proceed with smart parsing - output_dir: output directory, used to save intermediate results CSV - """ - logger.debug("🚀 non-llm parsing => recursive processing") - save_preds = raw_preds.copy() - - heading_preds = postprocess_headings(raw_preds, task="collapse") - save_preds.insert( - save_preds.columns.get_loc("level") + 1, - "lvl_cola", - heading_preds["level"].tolist(), - ) - - heading_preds = postprocess_headings(heading_preds, task="judge_negs") - save_preds.insert( - save_preds.columns.get_loc("lvl_cola") + 1, - "lvl_neg", - heading_preds["level"].tolist(), - ) - save_preds["reason"] = heading_preds["reason"] - - # mapping based on freq - if not proceed_smart: - heading_preds["level"] = heading_preds["level"].map( - lambda x: -1 if x == -2 else x - ) - heading_preds, lvl_mapping = build_level_mapping( - heading_preds, heading_preds["level"].tolist(), mode="freq" - ) - heading_preds = execute_level_mapping(heading_preds, lvl_mapping) - heading_preds.drop("origin_level", axis=1, inplace=True) - save_preds.insert( - save_preds.columns.get_loc("lvl_neg") + 1, - "lvl_map", - heading_preds["level"].tolist(), - ) - - return heading_preds - - -def est_hierarchies_llm( - raw_preds, - prompt_limt, - toc_hierarchies=None, - max_len=30, - max_depth=6, - model_name=None, - output_dir=None, - csv_suffix="", -): - """LLM-based hierarchy detection — first chunk via LLM, remaining chunks via reason-code mapping. - - When ``KB_LAYOUT_LLM_COMPACT_INPUT`` is enabled (default), consecutive - ``level == -1`` rows in ``raw_preds`` are folded into a single placeholder - row (``[N BODY LINES]``) before chunking. This shrinks the prompt, makes - most documents fit into a single chunk (skipping the lossy reason-code - mapping), and preserves the positional signal for the LLM. - - Strategy: - 1. (Optional) Compact raw_preds so consecutive body rows become placeholders. - 2. Send only the first chunk to LLM for hierarchy prediction. - 3. Collect ``{id -> level}`` from the LLM response (int ids only). - 4. For multi-chunk docs, extend that mapping via reason-code mapping on - chunks 1..N (placeholders excluded). - 5. Expand the id->level mapping back onto the ORIGINAL ``raw_preds``; - any row not present in the mapping defaults to ``level = -1``. - - Args: - raw_preds: raw data - prompt_limt: prompt character limit - toc_hierarchies: TOC hierarchies - max_len: maximum heading length (passed through to heading_tb_transfer) - max_depth: maximum hierarchy depth - model_name: LLM model name - output_dir: output directory, used to save intermediate results CSV - csv_suffix: suffix for intermediate CSV filenames - """ - model_name = _resolve_hierarchy_model_name(model_name) - if len(raw_preds) == 0: - return pd.DataFrame(columns=["id", "heading", "level", "reason"]) - - compact_enabled = os.environ.get( - "KB_LAYOUT_LLM_COMPACT_INPUT", "true" - ).strip().lower() in ("true", "1", "yes", "on") - preds_for_llm = _compact_for_llm(raw_preds) if compact_enabled else raw_preds.copy() - if compact_enabled: - placeholder_count = int(preds_for_llm["reason"].eq(PLACEHOLDER_REASON).sum()) - logger.info( - f"smart parse => compact input: {len(raw_preds)} → {len(preds_for_llm)} rows " - f"({placeholder_count} placeholder groups)" - ) - - # Short-circuit: if there are no heading candidates to judge (all rows were - # collapsed into placeholders, or raw_preds contains only level==-1 rows - # with compaction disabled), skip the LLM entirely and return raw_preds - # with all levels set to -1. - non_placeholder = ( - preds_for_llm[preds_for_llm["reason"].astype(str) != PLACEHOLDER_REASON] - if compact_enabled - else preds_for_llm - ) - if len(non_placeholder) == 0: - logger.info( - "smart parse => no heading candidates, skipping LLM hierarchy detection" - ) - fallback = raw_preds.copy()[["id", "heading", "level", "reason"]] - fallback["level"] = -1 - return fallback.sort_values("id").reset_index(drop=True) - - level_dfs, _raw_headings = heading_tb_transfer( - preds_for_llm, threshold=prompt_limt, max_start=max_len, max_end=5 - ) - chunk_sizes = [len(d) for d in level_dfs] - logger.info( - f"smart parse => {len(level_dfs)} chunk(s) | rows per chunk: {chunk_sizes} | " - f"threshold={prompt_limt} | max_start={max_len}" - ) - - # Pick the first chunk that actually contains heading candidates. When - # compaction is enabled a small prompt_limt may push a placeholder-only - # chunk to index 0 — using it would waste an LLM call and produce an empty - # mapping. Placeholder chunks that precede the chosen one contribute no - # reason-code signal (their ids map to -1 anyway). - basic_idx = 0 - for idx, chunk in enumerate(level_dfs): - if (chunk["reason"].astype(str) != PLACEHOLDER_REASON).any(): - basic_idx = idx - break - basic_df = level_dfs[basic_idx] - if basic_idx != 0: - logger.info( - f"smart parse => promoted chunk {basic_idx} as basic_df " - f"(chunks 0..{basic_idx - 1} contain only placeholders)" - ) - full_preds = None - try: - with stage_timer( - "heading.hierarchy_llm", - chunk_count=len(level_dfs), - base_chunk_rows=len(basic_df), - compact_enabled=compact_enabled, - source_row_count=len(raw_preds), - model_name=model_name, - ): - logger.debug("🚀 smart parse => interpreting hierarchy patterns...") - df4llm = basic_df.drop(columns=["reason"]).copy() - from .metadata_extractor import clean_md_text_for_llm - - # Keep formatting signals in `reason` / preliminary `level`, but let the LLM - # judge hierarchy from the semantic heading text instead of raw markdown markers. - df4llm["heading"] = df4llm["heading"].apply(clean_md_text_for_llm) - logger.debug(f"DataFrame transformation completed, rows: {len(df4llm)}") - - layout_res = hiearchy_llm( - df4llm, model_name, max_depth, toc_hierarchies, task="eval-headings" - ) - - # Build base_preds by aligning on basic_df["id"]: we always have one - # row per chunk-0 row in the rendered output, regardless of how many - # entries the LLM actually returned. Missing ids -> level=-1. - layout_level_by_id = {} - if isinstance(layout_res, list): - for item in layout_res: - if isinstance(item, dict) and "id" in item and "level" in item: - layout_level_by_id[item["id"]] = item["level"] - - def _level_for(rid): - if rid in layout_level_by_id: - return layout_level_by_id[rid] - try: - return layout_level_by_id.get(int(rid), -1) - except (TypeError, ValueError): - return -1 - - base_preds = ( - basic_df[["id", "heading", "reason"]].copy().reset_index(drop=True) - ) - base_preds.insert(2, "level", base_preds["id"].map(_level_for)) - - # Save base_preds as preds_3 (reflects what the LLM saw, compact or not) - save_intermediate_csv( - base_preds, output_dir, f"preds_3_llm_base{csv_suffix}" - ) - - # Collect {int_id -> level} from chunk-0 LLM output. Placeholder rows - # have non-integer ids and are skipped. - llm_levels = {} - for _, row in base_preds.iterrows(): - rid = row["id"] - if isinstance(rid, bool): - continue - if isinstance(rid, int): - llm_levels[rid] = row["level"] - - if len(level_dfs) > 1: - # Multi-chunk: build reason-code mapping from chunk-0 candidates and - # apply it to chunks 1..N to infer levels for headings beyond chunk 0. - # Placeholder rows are excluded from both the mapping source and the - # per-chunk application — they always map to level=-1 in the final df. - placeholder_mask_base = base_preds["reason"].eq(PLACEHOLDER_REASON) - figure_mask_base = base_preds["heading"].eq("Figure/Image") - exclude_mask_base = placeholder_mask_base | figure_mask_base - base_preds_for_mapping = base_preds[~exclude_mask_base].copy() - base_origin_for_mapping = basic_df.loc[ - ~exclude_mask_base.values, "level" - ].tolist() - - base_preds_for_mapping, lvl_mapping = build_level_mapping( - base_preds_for_mapping, base_origin_for_mapping, mode="freq" - ) - logger.debug( - f"mapping development finished: {len(lvl_mapping)} rules " - f"(placeholders and Figure/Image excluded)" - ) - - logger.debug( - f"mapping dataframe to levels across {len(level_dfs)} chunks..." - ) - lvl_mapping = handle_unseen_codes( - preds_for_llm, level_dfs, lvl_mapping, output_dir - ) - - for level_df in level_dfs: - placeholder_mask_chunk = level_df["reason"].eq(PLACEHOLDER_REASON) - figure_mask_chunk = level_df["heading"].eq("Figure/Image") - exclude_mask_chunk = placeholder_mask_chunk | figure_mask_chunk - non_excluded = level_df[~exclude_mask_chunk].copy() - if not non_excluded.empty: - non_excluded = execute_level_mapping(non_excluded, lvl_mapping) - for _, row in non_excluded.iterrows(): - rid = row["id"] - if isinstance(rid, bool): - continue - if isinstance(rid, int): - # Mapping may override chunk-0 LLM decisions when - # two rows share the same reason; accept that (the - # mapping is by construction the "representative" - # level for each reason-code). - llm_levels[rid] = row["level"] - logger.info( - f"multi-chunk mapping produced {len(llm_levels)} id→level entries" - ) - else: - logger.info( - "single chunk — skipping reason-code mapping, using LLM output directly" - ) - - # Expand back onto the original raw_preds: heading candidates take - # the LLM/mapping-assigned level; everything else is body text (-1). - full_preds = raw_preds.copy() - full_preds = full_preds[["id", "heading", "level", "reason"]] - - def _resolve_level(rid): - try: - int_id = int(rid) - except (TypeError, ValueError): - return -1 - lvl = llm_levels.get(int_id, -1) - try: - return int(lvl) - except (TypeError, ValueError): - return -1 - - full_preds["level"] = full_preds["id"].map(_resolve_level).astype(int) - - save_intermediate_csv( - full_preds, output_dir, f"preds_4_llm_final{csv_suffix}" - ) - - except Exception as e: - logger.warning(f"LLM-based parsing fails due to {e}, using non-llm pipeline...") - full_preds = est_hierarchies_naive(raw_preds.copy()) - return full_preds - - -def collapse_recursive(df, task, indices, merge_th=3, checked_pairs=None, depth=0): - """recursive collapse""" - if checked_pairs is None: - checked_pairs = set() - - if len(indices) < 2: - return - - for k in range(len(indices) - 1): - i, j = indices[k], indices[k + 1] - if (i, j) in checked_pairs: - continue - checked_pairs.add((i, j)) - - between = df.loc[i + 1 : j - 1] - i_txt = df.at[i, "heading"].strip() - j_txt = df.at[j, "heading"].strip() - - if task == "merge_short" and len(between) > 0: - between_lens = [count_cn_en(c) for c in between["heading"].tolist()] - between_lvls = [bl for bl in between["level"].tolist()] - i_half_len = int(count_cn_en(i_txt) / 2) - too_short = sum(between_lens) <= merge_th or sum(between_lens) < i_half_len - - if too_short and all( - bl == -1 for bl in between_lvls - ): # only non-headings can be merged - logger.debug( - f"⚠️ too short between {i}=>{i_txt[:15]} and {j}=>{j_txt[:15]} => merge to {i}" - ) - between_txts = [ - str(r["heading"]).strip() - for _, r in between.iterrows() - if isinstance(r.get("heading"), str) and r["heading"].strip() - ] - - if between_txts: - joined_txt = "\n".join(between_txts) - df.at[i, "heading"] = f"{i_txt} {joined_txt}" - - for idx in between.index: - df.at[idx, "level"] = -1 - df.at[idx, "reason"] = f"Merged into {i}" - logger.debug(f"\tmerged texts: {joined_txt[:50]}...") - - elif task == "collapse" and len(between) == 0: - logger.debug( - f"⚠️ Empty between i={i_txt[:15]}, j={j_txt[:15]} => set i.level=-1, j.level=Not Sure" - ) - df.at[i, "level"] = -2 # "Not Sure" sentinel (int-safe) - df.at[j, "level"] = -2 # "Not Sure" sentinel (int-safe) - - # ========== get subgroups for recursive tasks ========== - sub_between = between[between["level"] != -1] - code2sub = defaultdict(list) - for idx, row in sub_between.iterrows(): - level = row["level"] - reason = row["reason"] - if level != -1: - code2sub[(level, reason)].append(idx) - - for _, sub_indices in code2sub.items(): - collapse_recursive( - df, task, sub_indices, merge_th, checked_pairs, depth + 1 - ) - - -def postprocess_headings(df, task, max_depth=-1): - """postprocess headings""" - if task == "judge_negs": - for i, row in df.iterrows(): - neg_code = remove_by_conditions(row["heading"], include_punc=True) - if any(x > 0 for x in neg_code): - current_code = str(df.loc[i, "reason"]) - - neg_match = re.search(r"(.*NEG\s*)\[[^\]]*\](.*)", current_code) - if neg_match: - update_code = f"{neg_match.group(1)}{neg_code}{neg_match.group(2)}" - else: - update_code = f"{current_code} NEG {neg_code}" - - df.loc[i, "level"] = -1 - df.loc[i, "reason"] = update_code - return df - - elif task == "merge_continuous": - denoised_rows = [] - punc_pattern = re.compile(r'[.,!?;:，。！？；：）】〕｝〉》’”"]$') - - i = 0 - while i < len(df): - row = df.iloc[i] - current_content = str(row["heading"]).strip() - current_level = row["level"] - - j = i + 1 - while j < len(df): - next_row = df.iloc[j] - next_content = str(next_row["heading"]).strip() - next_level = next_row["level"] - - # Skip merge if ID is not continuous (indicates table/image was skipped in between) - expected_id = row["id"] + (j - i) - if next_row["id"] != expected_id: - break - - # both current and next rows are not heading & current row has no punctuation -> merge - current_not_punc = not punc_pattern.search(current_content[-2:]) - if (current_level == -1 and next_level == -1) and current_not_punc: - current_content += " " + next_content - j += 1 - else: - break - - merge_row = row.copy() - merge_row["heading"] = current_content - denoised_rows.append(tuple(merge_row)) - i = j - return pd.DataFrame(denoised_rows, columns=["id", "heading", "level", "reason"]) - - elif task == "merge_short" or task == "collapse": - group2indices = defaultdict(list) - for idx, row in df.iterrows(): - level = row["level"] - reason = row["reason"] - if level != -1: - group2indices[(level, reason)].append(idx) - - checked_pairs = set() - for _, indices in group2indices.items(): - collapse_recursive( - df, task, indices, merge_th=3, checked_pairs=checked_pairs, depth=0 - ) - - if task == "merge_short": - drop_between = df.index[ - df["reason"].astype(str).str.startswith("Merged into", na=False) - ].tolist() - if drop_between: - logger.debug( - f"🛠️ Delete rows labeled as merged into, total {len(drop_between)} rows" - ) - df.drop(drop_between, inplace=True) - df.reset_index(drop=True, inplace=True) - return df - - else: - return None - - -# def parse_outline_hier(markdown_text): -# lines = markdown_text.strip().splitlines() -# stack = [] -# root = [] -# for line in lines: -# line = line.replace('markdown', '') # handle possible unexpected outputs -# if not line.strip(): -# continue - -# stripped = line.lstrip() -# indent = len(line) - len(stripped) -# match = re.match(r"[-*+] (.+)", stripped) -# if not match: -# continue - -# title = match.group(1).strip() -# node = {"chapter": title, "children": [], 'serial': 1} -# level = indent // 2 # Two spaces per level, adjustable if needed. -# if level == 0: -# node['serial'] = len(root)+1 -# root.append(node) -# stack = [(level, node)] -# else: -# while stack and stack[-1][0] >= level: -# stack.pop() -# if stack: -# parent = stack[-1][1] -# node['serial'] = len(parent['children']) + 1 -# parent["children"].append(node) -# stack.append((level, node)) -# return root - - -# def outline_to_markdown(nodes, level=0, path=""): -# rows = [] -# def traverse(node_list, level, path_prefix): -# for node in node_list: -# split_char = settings.SPLIT_CHAR or "/" -# current_path = f"{path_prefix} {split_char} {node['chapter']}" if path_prefix else node['chapter'] -# rows.append({ -# "path": current_path, -# "title": node["chapter"], -# "thoughts": node.get("thoughts", "").strip(), -# "level": level -# }) -# if node.get("children"): -# traverse(node["children"], level + 1, current_path) -# traverse(nodes, level, path) -# return pd.DataFrame(rows) diff --git a/apps/worker/app/services/document_parser/md_parser.py b/apps/worker/app/services/document_parser/md_parser.py deleted file mode 100755 index ab85751bc..000000000 --- a/apps/worker/app/services/document_parser/md_parser.py +++ /dev/null @@ -1,870 +0,0 @@ -# pyright: reportArgumentType=false, reportAssignmentType=false, reportOptionalIterable=false, reportOptionalMemberAccess=false, reportOptionalOperand=false, reportOptionalSubscript=false -import json -import os -import re -import shutil -from pathlib import Path - -import gevent -import pandas as pd -from app.services.common.kb_utils import ( - find_matches_parsing, - gen_str_codes, - get_str_time, - process_dup_paths_df, -) -from app.services.document_parser.html_parser import ( - first_cols_rows_html, - merge_html_tables, -) -from app.services.document_parser.image_parser import ( - MD_IMAGE_PATTERN, - _get_vision_client, - ask_image, - detect_summary_img_md, - perceptual_hash, -) -from app.services.document_parser.layout_parser import md_heading_match, pred_titles -from app.services.document_parser.stage_profiler import stage_timer -from app.services.document_parser.table_parser import ( - extract_tables_by_forms, - identify_tables, - sanitize_table_name_from_header, -) -from app.services.document_parser.toc_parser import detect_tocs_in_texts -from app.services.document_parser.txt_parser import ( - extract_title_keywords_summary, - split_title_summary, -) -from gevent.pool import Pool as GeventPool -from loguru import logger - -from shared.core.config import settings -from shared.utils.chunk_refs import build_chunk_ref, has_chunk_ref -from shared.utils.file_utils import path_handle -from shared.utils.text_utils import tokenize2stw_remove - - -def resolve_workspace_image_path( - candidate_path: Path, workspace_path: Path -) -> Path | None: - """Return the candidate only when it exists inside the current job workspace.""" - resolved_path = candidate_path.resolve(strict=False) - try: - resolved_path.relative_to(workspace_path) - except ValueError: - return None - return resolved_path if resolved_path.exists() else None - - -def resolve_markdown_image_source_path(output_dir: str, img_path: str) -> Path | None: - """Handle local absolute refs and container cwd-relative refs safely.""" - if not img_path: - return None - - workspace_path = Path(output_dir).resolve() - raw_path = Path(img_path).expanduser() - candidate_paths = ( - [raw_path] - if raw_path.is_absolute() - else [ - workspace_path / raw_path, - Path.cwd() / raw_path, - ] - ) - - for candidate_path in candidate_paths: - resolved_path = resolve_workspace_image_path(candidate_path, workspace_path) - if resolved_path is not None: - return resolved_path - - return None - - -def find_surround_context(md_lines, lid): - def is_skip(line): - s = line.strip() - is_image = re.findall(MD_IMAGE_PATTERN, line, flags=re.IGNORECASE) - is_table, _, _ = identify_tables(line) - return not s or is_image or is_table - - n = len(md_lines) - prev_text = "" - for i in range(max(lid - 5, 0), lid): - if not is_skip(md_lines[i]): - prev_text = md_lines[i].strip() - break - - next_text = "" - for i in range(lid + 1, min(lid + 6, n)): - if not is_skip(md_lines[i]): - next_text = md_lines[i].strip() - break - return f"{prev_text} {next_text}".strip() - - -def heading_md_relocate(md_lines, heading_preds): - """Relocate markdown headings based on predicted levels (sxjg simplified logic)""" - - def remove_hash(txt): - return re.sub(r"^\s*(#+)\s*", "", txt) - - for lid, line_txt in enumerate(md_lines): - pred_level_df = heading_preds[heading_preds["id"] == lid] - - if pred_level_df.empty: # if the line does not enter predicting - line_txt = remove_hash(line_txt) - else: - pred_level = pred_level_df["level"].iloc[0] - if pred_level < 0: - line_txt = remove_hash(line_txt) - else: - # sxjg simplified: remove all #, then add correct number of # - clean_text = line_txt.lstrip("#").lstrip() - line_txt = f"{'#' * int(pred_level)} {clean_text}" - # update lines - md_lines[lid] = line_txt.strip() - - md_lines = [line for line in md_lines if line.strip() != ""] - return md_lines # note the length=original md_lines but contents/level are updated - - -def eval_md_headings( - md_lines, - source_type, - toc_hierarchies=None, - smart_parse=False, - model_name=None, - output_dir=None, - layout_json_path=None, -): - """Evaluate markdown headings with optional TOC hierarchies context""" - heading_preds = pred_titles( - md_lines, - source_type, - toc_hierarchies=toc_hierarchies, - enable_regx=True, - smart_parse=smart_parse, - model_name=model_name, - output_dir=output_dir, - layout_json_path=layout_json_path, - ) - - if len(heading_preds) == 0: - lines_with_heading = md_lines - else: - lines_with_heading = heading_md_relocate(md_lines, heading_preds) - return lines_with_heading - - -def clean_md_table_lines(table_lines, start_line_num): - expected_columns = table_lines[0].count("|") - 1 - cleaned_lines = [] - error_lines = [] # To record line numbers that need cleaning - - for i, line in enumerate(table_lines): - line_columns = line.count("|") - 1 - current_line_num = ( - start_line_num + i - ) # Calculate the current line number in the original file - if line_columns == expected_columns: - cleaned_lines.append(line) - else: - error_lines.append(current_line_num) - if line_columns > expected_columns: - parts = line.split("|") - cleaned_line = "|".join( - parts[: expected_columns + 1] - ) # If there are more columns, combine them (or drop extra columns) - cleaned_lines.append(cleaned_line) - elif line_columns < expected_columns: - # If there are fewer columns, pad the line (or you could skip it) - cleaned_line = line + "|" * (expected_columns - line_columns) - cleaned_lines.append(cleaned_line) - return cleaned_lines, error_lines - - -def replace_chunk_ref_in_rows(df_list, old_path: str, new_path: str) -> None: - """Rewrite readable chunk refs after deferred image/table renames.""" - old_ref = build_chunk_ref(old_path) - new_ref = build_chunk_ref(new_path) - if not old_ref or old_ref == new_ref: - return - - for row in df_list: - if len(row) > 0 and isinstance(row[0], str): - row[0] = row[0].replace(old_ref, new_ref) - if len(row) > 1 and row[1] == old_path: - row[1] = new_path - if len(row) > 2 and isinstance(row[2], str): - row[2] = row[2].replace(old_ref, new_ref) - if len(row) > 8 and isinstance(row[8], str): - row[8] = row[8].replace(old_ref, new_ref) - - -def update_df_list( - df_list, - content_items, - path, - llm_paras, - time_stamp, - page_nums="", - summary_len=1500, - skip_llm=False, -): - """Flush accumulated content_items into a chunk row in df_list. - - Args: - content_items: list of content strings. Each item is either pure text - or an IMAGE/TABLE ref block. know_id is generated from pure text - items only (deterministic), while full content includes all items. - skip_llm: if True, skip inline LLM calls (deferred to parallel batch). - """ - # Separate pure text from IMAGE/TABLE ref blocks for deterministic know_id - text_items = [item for item in content_items if not has_chunk_ref(str(item))] - pure_text = "\n".join(text_items).strip() - bottom_content = "\n".join(content_items).strip() - - match_type = find_matches_parsing(bottom_content, path) - know_id_source = pure_text if pure_text else f"{path or ''}::{page_nums or ''}" - know_id = gen_str_codes(know_id_source) - bottom_tokens = tokenize2stw_remove([bottom_content], llm_paras["stopwords"]) - - keywords = "" - summary = "" - needs_llm = ( - not skip_llm and len(bottom_content) > summary_len and llm_paras["summary_txt"] - ) - if needs_llm: - _title, keywords, summary = extract_title_keywords_summary( - bottom_content, max_keywords=3, summary_len=summary_len - ) - - df_list.append( - [ - bottom_content, - path, - match_type, - len(bottom_content), - keywords, - summary, - know_id, - bottom_tokens, - "", - time_stamp, - page_nums, - ] - ) - return df_list - - -def parse_md( - output_dir, - source_type, - file_path=None, - md_lines=None, - base_llm_paras=None, - relative_root=None, -): - if md_lines is None and file_path is not None: - from shared.utils.CommonHelperSync import is_remote, load_file_bytes - - if is_remote(file_path): - file_bytes = load_file_bytes(file_path) - md_content = file_bytes.decode("utf-8") - md_lines = md_content.splitlines() - else: - with open(file_path, "r", encoding="utf-8") as file: - md_lines = file.readlines() - - md_lines = [line.strip() for line in md_lines if line.strip() != ""] - - # Preprocess: merge multi-line HTML tables into single lines - md_lines = merge_html_tables(md_lines) - - # Detect TOC using async LLM-based detection - toc_model_name = ( - base_llm_paras.get("model_name", settings.NORMOL_MODEL) - if base_llm_paras - else settings.NORMOL_MODEL - ) - hierarchy_model_name = ( - (base_llm_paras.get("hierarchy_model_name") or toc_model_name) - if base_llm_paras - else (settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL) - ) - - with stage_timer( - "md.detect_toc", line_count=len(md_lines), model_name=toc_model_name - ): - toc_hierarchies, md_lines = detect_tocs_in_texts( - md_lines, - model_name=toc_model_name, - hierarchy_model_name=hierarchy_model_name, - ) - - # Save toc_hierarchies.json to output_dir (will be included in final zip package) - if toc_hierarchies: - toc_json_path = os.path.join(output_dir, "toc_hierarchies.json") - with open(toc_json_path, "w", encoding="utf-8") as f: - json.dump(toc_hierarchies, f, ensure_ascii=False, indent=2) - logger.info(f"Saved TOC hierarchies to {toc_json_path}") - - # Clean old artifacts to prevent accumulation across debug runs. - # In production each job uses a fresh workspace so rmtree never triggers. - tb_dir = os.path.join(output_dir, "tables") - if os.path.isdir(tb_dir): - shutil.rmtree(tb_dir) - os.makedirs(tb_dir, exist_ok=True) - img_dir = os.path.join(output_dir, "images") - if os.path.isdir(img_dir): - # Only remove parse_md's own output (image-N-*) from previous runs - for fname in os.listdir(img_dir): - if re.match(r"^image-\d+", fname): - os.remove(os.path.join(img_dir, fname)) - os.makedirs(img_dir, exist_ok=True) - - # initialize vars - split_char = settings.SPLIT_CHAR or "/" - df_list = [] - path_stack = [] # sxjg: uses (heading, level) tuples - inner_paths = [] - error_line_numbers = [] - table_lines = [] - current_pg_num = 0 - chunk_pages = set() # collect all page numbers seen during current chunk - base_level = None - content_items = [] - # Use relative_root as initial path (not absolute output_dir) - path = relative_root if relative_root else "" - table_count = 1 - img_count = 1 - path_counter = {} # Track path occurrences for deduplication - deferred_llm_tasks = [] # Collected during loop, executed in parallel after - _seen_images: dict[str, dict] = {} # sha256_hex -> cached image info for dedup - - # Find layout.json path - layout_json_path = os.path.join(output_dir, "layout.json") - if not os.path.exists(layout_json_path): - layout_json_path = None - logger.debug("layout.json not found, META features will not be added") - - # estimate hierarchies with toc_hierarchies context - with stage_timer( - "md.predict_headings", - line_count=len(md_lines), - smart_parse=base_llm_paras["smart_title_parse"], - model_name=hierarchy_model_name, - ): - lines_with_heading = eval_md_headings( - md_lines, - source_type, - toc_hierarchies=toc_hierarchies, - smart_parse=base_llm_paras["smart_title_parse"], - model_name=hierarchy_model_name, - output_dir=output_dir, - layout_json_path=layout_json_path, - ) - - time_stamp = get_str_time() - logger.debug("Parsing md data... total_lines={}", len(lines_with_heading)) - for i, line in enumerate(lines_with_heading): - if "" in line: - if "page" in line or "Slide number" in line: - # Parse actual page number from marker: - pg_match = re.search(r"page\s+(\d+)", line) - if pg_match: - current_pg_num = int(pg_match.group(1)) - else: - current_pg_num += 1 # fallback for Slide number or old format - chunk_pages.add(current_pg_num) - continue - - last_context = find_surround_context( - lines_with_heading, i - ) # record the previous and next line which is not table/image - current_heading, current_heading_level = md_heading_match(line, as_is=False) - - if ( - not current_heading_level == -1 - ): # indicate a new path should be evaluated or added - if content_items: # record contents of the last path and reset content - # Build page_nums from collected pages during this chunk - chunk_page_str = ( - ",".join(str(p) for p in sorted(chunk_pages)) if chunk_pages else "" - ) - df_list = update_df_list( - df_list, - content_items, - path, - base_llm_paras, - time_stamp, - page_nums=chunk_page_str, - skip_llm=True, - ) - content_items = [] - chunk_pages = set() # reset for next chunk - if current_pg_num > 0: - chunk_pages.add( - current_pg_num - ) # carry current page into next chunk - elif path and path != (relative_root or ""): - # Consecutive headings with no body text between them: - # Create a placeholder chunk so the previous heading's path - chunk_page_str = ( - ",".join(str(p) for p in sorted(chunk_pages)) if chunk_pages else "" - ) - df_list = update_df_list( - df_list, - [], - path, - base_llm_paras, - time_stamp, - page_nums=chunk_page_str, - skip_llm=True, - ) - - # update path based on path name and level - if base_level is None: - base_level = current_heading_level - elif current_heading_level < base_level: - base_level = current_heading_level - - adjusted_level = current_heading_level - base_level + 1 - path_stack = [(h, lvl) for h, lvl in path_stack if lvl < adjusted_level] - - # Build tentative path to check for duplicates - # Sanitize heading: replace split_char in heading text to prevent path corruption - current_heading = ( - current_heading.replace(split_char, "∕") - if split_char in current_heading - else current_heading - ) - tentative_heading = current_heading - tentative_names = [h for h, lvl in path_stack] + [tentative_heading] - tentative_path_parts = [relative_root] if relative_root else [] - tentative_path_parts.extend(tentative_names) - tentative_path = split_char.join(tentative_path_parts) - - # Deduplicate: if path already exists, add suffix - if tentative_path in path_counter: - path_counter[tentative_path] += 1 - suffix = path_counter[tentative_path] - current_heading = ( - f"{current_heading}_{suffix}" # Modify heading with suffix - ) - else: - path_counter[tentative_path] = 1 - - path_stack.append((current_heading, adjusted_level)) - - # Extract pure heading names for path construction - heading_names = [h for h, lvl in path_stack] - # Use relative_root as prefix - path_parts = [relative_root] if relative_root else [] - path_parts.extend(heading_names) - inner_paths.append(split_char.join(heading_names)) - path = split_char.join(path_parts) # path with relative root - - else: # no path change, remain in the same hierarchy - # a. handle lines containing images (LLM deferred to post-loop parallel batch) - img_name_context = path_handle(last_context[:10], mode="clean_single") - img_name = f"image-{str(img_count)}-{img_name_context}" - # Always skip inline LLM — vision calls are deferred to parallel batch - imgs = detect_summary_img_md(line, last_context, output_dir, mode=False) - - for img_path, img_title, img_summary in imgs: - img_suffix = os.path.splitext(img_path)[-1] - update_img_path = os.path.join(img_dir, f"{img_name}{img_suffix}") - - # Check if source image file exists before renaming - source_path = resolve_markdown_image_source_path(output_dir, img_path) - if source_path is None or not source_path.exists(): - logger.warning(f"Image file not found, skipping rename: {img_path}") - img_count += 1 - continue - - # Document-level dedup: perceptual hash for visual duplicates - with open(source_path, "rb") as f: - img_binary_hash = perceptual_hash(f.read()) - - if img_binary_hash in _seen_images: - cached = _seen_images[img_binary_hash] - content_items.append(cached["img_content"]) - df_list.append( - [ - cached["img_content"], - cached["relative_img_path"], - "image", - len(cached["img_content"]), - "", - cached["img_summary_field"], - cached["temp_uid"], - "", - "", - time_stamp, - str(current_pg_num) if current_pg_num > 0 else "", - ] - ) - logger.debug( - f"Skipped duplicate image (hash={img_binary_hash[:12]}...)" - ) - # Remove unused source file since we reuse the cached image - try: - source_path.unlink() - except OSError: - pass - continue - - os.rename(source_path, update_img_path) - - # Image index (always present) - image_index = f"image-{img_count}" - - # Fallback: LLM summary -> last_context -> None - effective_summary = img_summary or last_context or None - - # Deterministic know_id: use image binary hash - temp_uid = gen_str_codes(img_binary_hash) - relative_img_path = f"images/{img_name}{img_suffix}" - img_ref = build_chunk_ref(relative_img_path) - - # Build img_summary_field for df_list: image-n + optional summary - if effective_summary: - img_summary_field = f"{image_index}\n{effective_summary}" - else: - img_summary_field = image_index - - # Build image_ref for content: optional summary + image path ref - if effective_summary: - img_content = f"\n{effective_summary}\n{img_ref}\n" - else: - img_content = f"\n{img_ref}\n" - - content_items.append(img_content) - - df_list.append( - [ - img_content, - relative_img_path, - "image", - len(img_content), - "", - img_summary_field, - temp_uid, - "", - "", - time_stamp, - str(current_pg_num) if current_pg_num > 0 else "", - ] - ) - - # Cache result for document-level dedup - _seen_images[img_binary_hash] = { - "relative_img_path": relative_img_path, - "img_content": img_content, - "img_summary_field": img_summary_field, - "temp_uid": temp_uid, - } - - if base_llm_paras["summary_image"]: - # Store img_dir, img_name, img_suffix for post-loop rename (mirrors table deferred task) - deferred_llm_tasks.append( - ( - "image", - len(df_list) - 1, - relative_img_path, - img_dir, - img_name, - img_suffix, - ) - ) - img_count += 1 - - # TODO for large and dense tables, such as "Epstein flight logs", - # integrate tabula-py as an independent extraction path to solve VLM hallucinations and misplacement - # b. handle lines containing tables - tb_bool, form, _ = identify_tables(line) - if tb_bool: - if form == "html": - # each line is a complete table - process immediately - tb_str = line - elif form == "md": - # For MD tables, accumulate lines until table ends - table_lines.append(line) - if i + 1 >= len(lines_with_heading): - tb_bool_next = False - else: - tb_bool_next, _, _ = identify_tables( - lines_with_heading[i + 1].strip() - ) - - if not tb_bool_next or i == len(lines_with_heading) - 1: - cleaned_table_lines, error_lines = clean_md_table_lines( - table_lines, start_line_num=i - ) - tb_str = "\n".join(cleaned_table_lines) - error_line_numbers.extend(error_lines) - tb_str = extract_tables_by_forms(tb_str, form="md") - else: - continue # Keep accumulating MD table lines - else: - continue # Unknown form, skip - - # Extract first row and first column for fallback file naming only - first_row_text, first_col_text = first_cols_rows_html(tb_str) - - # Table index (always present) - table_index = f"table-{table_count}" - - # LLM title + keywords + summary deferred to post-loop parallel batch - llm_title = None - llm_summary = None - tb_keywords = "" - - # Build tb_summary for df_list: table-n + optional LLM summary - if llm_summary: - tb_summary = f"{table_index}\n{llm_summary}" - else: - tb_summary = table_index - - raw_tb_name = ( - sanitize_table_name_from_header(first_row_text) - if first_row_text - else "" - ) - # Use LLM title for filename when available, fallback to sanitized header - effective_name = llm_title if llm_title else raw_tb_name - tb_name = path_handle( - f"table-{str(table_count)} {effective_name}", mode="clean_single" - ) - temp_uid = gen_str_codes((tb_str + str(table_count))) - - relative_tb_path = f"tables/{tb_name}.html" - tb_ref = build_chunk_ref(relative_tb_path) - - # Build table_ref for content: optional LLM summary + table path ref - if llm_summary: - content_items.append(f"\n{llm_summary}\n{tb_ref}\n") - else: - content_items.append(f"\n{tb_ref}\n") - tb_path = os.path.join(tb_dir, f"{tb_name}.html") - # Add border to HTML tables for consistent display - tb_str_with_border = tb_str.replace( - "
", "
" - ).replace("
0 else "", - ] - ) - if base_llm_paras["summary_table"]: - deferred_llm_tasks.append( - ( - "table", - len(df_list) - 1, - tb_str, - tb_dir, - tb_name, - table_count - 1, - ) - ) - table_lines = [] # Reset table_lines after storing the DataFrame - table_count += 1 - - # c. handle plain texts - if len(imgs) == 0 and not tb_bool: - content_items.append(line.strip()) - if current_pg_num > 0: - chunk_pages.add(current_pg_num) # track page for this content line - - if content_items: # handle the remaining contents, append them to the last section - chunk_page_str = ( - ",".join(str(p) for p in sorted(chunk_pages)) if chunk_pages else "" - ) - df_list = update_df_list( - df_list, - content_items, - path, - base_llm_paras, - time_stamp, - page_nums=chunk_page_str, - skip_llm=True, - ) - - # Collect text chunk deferred tasks (entries needing summary/keywords) - summary_len = 1500 - if base_llm_paras.get("summary_txt"): - for idx, entry in enumerate(df_list): - marker = entry[2] # col 2: match_type / img_id / table_id - if isinstance(marker, str) and marker.strip().split("\n", 1)[0].lower() in { - "image", - "table", - }: - continue - if len(entry[0]) > summary_len and not entry[4] and not entry[5]: - deferred_llm_tasks.append(("text", idx, entry[0])) - - # ── Post-loop: execute all deferred LLM calls in parallel via gevent ── - if deferred_llm_tasks: - image_task_count = sum(1 for task in deferred_llm_tasks if task[0] == "image") - table_task_count = sum(1 for task in deferred_llm_tasks if task[0] == "table") - text_task_count = sum(1 for task in deferred_llm_tasks if task[0] == "text") - logger.info( - f"Running {len(deferred_llm_tasks)} deferred summary LLM calls in parallel" - ) - max_concurrent = getattr(settings, "SUMMARY_LLM_MAX_CONCURRENT", 8) - - with stage_timer( - "md.deferred_summaries", - total_tasks=len(deferred_llm_tasks), - image_tasks=image_task_count, - table_tasks=table_task_count, - text_tasks=text_task_count, - max_concurrent=min(max_concurrent, len(deferred_llm_tasks)), - ): - - def _run_deferred(task): - task_type, idx = task[0], task[1] - try: - if task_type == "image": - relative_path = task[2] - client = _get_vision_client() - # TODO: Risk of missing text content if MinerU outputted a pure text image. - # Consider adding judge-image-type and OCR fallback as done in image_parser.parse_image. - llm_resp = ask_image(client, output_dir, paths_=[relative_path]) - if llm_resp: - img_title, img_summary = split_title_summary(llm_resp) - else: - img_title, img_summary = None, None - return idx, task_type, (img_title, img_summary) - elif task_type == "table": - tb_html = task[2] - title, kw, summary = extract_title_keywords_summary( - tb_html, max_keywords=3 - ) - return idx, task_type, (title, kw, summary) - elif task_type == "text": - text_content = task[2] - _, kw, summary = extract_title_keywords_summary( - text_content, max_keywords=3, summary_len=summary_len - ) - return idx, task_type, (kw, summary) - except Exception as e: - logger.warning( - f"Deferred {task_type} LLM call failed for idx={idx}: {e}" - ) - return idx, task_type, None - - pool = GeventPool(size=min(max_concurrent, len(deferred_llm_tasks))) - greenlets = [pool.spawn(_run_deferred, task) for task in deferred_llm_tasks] - gevent.joinall(greenlets) - - # Build a lookup from deferred task list: idx -> original task tuple - deferred_by_idx = {task[1]: task for task in deferred_llm_tasks} - - for g in greenlets: - if g.value is None: - continue - idx, task_type, result = g.value - if result is None: - continue - if task_type == "image": - img_title, img_summary = result - entry = df_list[idx] - if img_summary: - image_index = entry[5].split("\n")[0] if entry[5] else "image" - entry[5] = f"{image_index}\n{img_summary}" - # Rename image file if LLM provided a better title (mirrors table rename logic) - if img_title: - orig_task = deferred_by_idx[idx] - i_dir, old_img_name, i_suffix = ( - orig_task[3], - orig_task[4], - orig_task[5], - ) - safe_title = path_handle(img_title, mode="clean_single") - # Derive image index number from old_img_name (e.g. "image-3-xxx" -> "3") - img_num_match = re.match(r"image-(\d+)", old_img_name) - img_num = ( - img_num_match.group(1) - if img_num_match - else ( - old_img_name.split("-")[1] - if "-" in old_img_name - else "0" - ) - ) - new_img_name = path_handle( - f"image-{img_num}-{safe_title}", mode="clean_single" - ) - old_path = os.path.join(i_dir, f"{old_img_name}{i_suffix}") - new_path = os.path.join(i_dir, f"{new_img_name}{i_suffix}") - if old_path != new_path and os.path.exists(old_path): - os.rename(old_path, new_path) - new_relative_path = f"images/{new_img_name}{i_suffix}" - replace_chunk_ref_in_rows( - df_list, entry[1], new_relative_path - ) - entry[1] = new_relative_path - elif task_type == "table": - title, kw, summary = result - entry = df_list[idx] - entry[4] = kw if isinstance(kw, str) else "" - if summary: - table_index = ( - entry[5] - if "\n" not in entry[5] - else entry[5].split("\n")[0] - ) - entry[5] = f"{table_index}\n{summary}" - # Rename table file if LLM provided a better title - if title: - orig_task = deferred_by_idx[idx] - t_dir, old_tb_name, t_count = ( - orig_task[3], - orig_task[4], - orig_task[5], - ) - safe_title = ( - sanitize_table_name_from_header(title) if title else "" - ) - new_tb_name = path_handle( - f"table-{t_count} {safe_title}", mode="clean_single" - ) - old_path = os.path.join(t_dir, f"{old_tb_name}.html") - new_path = os.path.join(t_dir, f"{new_tb_name}.html") - if old_path != new_path and os.path.exists(old_path): - os.rename(old_path, new_path) - new_relative_path = f"tables/{new_tb_name}.html" - replace_chunk_ref_in_rows( - df_list, entry[1], new_relative_path - ) - entry[1] = new_relative_path - elif task_type == "text": - kw, summary = result - df_list[idx][4] = kw if isinstance(kw, str) else "" - df_list[idx][5] = summary if isinstance(summary, str) else "" - - logger.info( - f"Completed {len(deferred_llm_tasks)} deferred summary LLM calls" - ) - - with stage_timer("md.build_dataframe", row_count=len(df_list)): - doc_df = pd.DataFrame(df_list, columns=settings.ALL_DF_COLS.split(",")) - doc_df = process_dup_paths_df(doc_df) - - return doc_df diff --git a/apps/worker/app/services/document_parser/mineru_pdf_service.py b/apps/worker/app/services/document_parser/mineru_pdf_service.py deleted file mode 100644 index 24a2145cc..000000000 --- a/apps/worker/app/services/document_parser/mineru_pdf_service.py +++ /dev/null @@ -1,764 +0,0 @@ -# pyright: reportUnusedExpression=false -import os -import time -from typing import Any, Callable, Optional - -import requests -from app.services.document_parser.mineru_quota_manager import get_mineru_quota_manager -from app.services.document_parser.parser_log_utils import truncate_log_value -from loguru import logger -from requests.adapters import HTTPAdapter -from urllib3.util.retry import Retry - -from shared.core.config import settings -from shared.core.constants import APIConstants -from shared.core.exceptions.domain_exceptions import ( - MinerUServiceException, - PDFParsingException, - StorageServiceException, - TimeoutException, - UnavailableException, -) -from shared.core.exceptions.knowhere_exception import KnowhereException -from shared.utils.CommonHelperSync import is_remote -from shared.utils.FileDownUpUtils import s3_download_extract_zip - -MINERU_UPLOAD_TIMEOUT = ( - settings.MINERU_UPLOAD_CONNECT_TIMEOUT, - settings.MINERU_UPLOAD_READ_TIMEOUT, -) - - -def _build_mineru_session() -> requests.Session: - session = requests.Session() - # Hybrid rate-limit control: urllib3 handles transient 429s with backoff - # (respects Retry-After header); application-level handling in callers - # covers persistent rate limits with Redis token marking and pool rotation. - retry_strategy = Retry( - total=settings.MINERU_UPLOAD_RETRY_TOTAL, - backoff_factor=settings.MINERU_UPLOAD_RETRY_BACKOFF_FACTOR, - status_forcelist=[429, 502, 503, 504], - allowed_methods=["GET", "POST", "PUT"], - raise_on_status=False, - ) - adapter = HTTPAdapter( - max_retries=retry_strategy, - pool_connections=1, - pool_maxsize=settings.MINERU_POOL_MAXSIZE, - ) - session.mount("https://", adapter) - session.mount("http://", adapter) - return session - - -_mineru_session: Optional[requests.Session] = None - - -def get_mineru_session() -> requests.Session: - global _mineru_session - if _mineru_session is None: - _mineru_session = _build_mineru_session() - return _mineru_session - - -def get_mineru_headers(api_key: str) -> dict[str, str]: - return { - "Content-Type": "application/json", - "Authorization": f"Bearer {api_key}", - } - - -def _mineru_logger(step: str, **fields: Any): - return logger.bind(service="mineru", step=step, **fields) - - -def _should_use_mineru_s3_url_mode(s3_key: Optional[str]) -> bool: - if settings.FORCE_MINERU_UPLOAD_ENABLED: - return False - - return settings.ENVIRONMENT != "development" and s3_key is not None - - -def _log_mineru_url_mode_storage_fallback( - operation: str, - s3_key: str, - local_file_path: Optional[str], - exc: Exception, -) -> None: - _mineru_logger( - "url_mode_storage_fallback", - operation=operation, - source_s3_key=s3_key, - local_file_path=local_file_path, - error_type=type(exc).__name__, - error_message=truncate_log_value(exc), - ).warning( - "MinerU URL-mode storage preparation failed. Falling back to direct upload." - ) - - -def _log_mineru_url_mode_ingestion_fallback( - operation: str, - s3_key: str, - pdf_url: str, - exc: Exception, -) -> None: - _mineru_logger( - "url_mode_ingestion_fallback", - operation=operation, - source_s3_key=s3_key, - source_kind="remote_url" if is_remote(pdf_url) else "local_file", - source_path=None if is_remote(pdf_url) else pdf_url, - error_type=type(exc).__name__, - error_message=truncate_log_value(exc), - ).warning("MinerU URL-mode ingestion setup failed. Falling back to direct upload.") - - -def _inspect_mineru_source_s3_key(s3_key: Optional[str]) -> tuple[Optional[str], bool]: - """Inspect whether URL mode can reuse or prepare the requested S3 source key.""" - if not _should_use_mineru_s3_url_mode(s3_key): - return None, False - - assert s3_key is not None - from app.services.storage.sync_storage_service import verify_s3_file_exists - - try: - existing_file = verify_s3_file_exists(s3_key, settings.S3_BUCKET_NAME) - except Exception as exc: - _log_mineru_url_mode_storage_fallback( - operation="verify_source_object", - s3_key=s3_key, - local_file_path=None, - exc=exc, - ) - return None, False - - if existing_file.get("exists"): - _mineru_logger( - "url_mode_source_reused", - source_s3_key=s3_key, - ).info("Reusing existing S3 source for MinerU URL mode") - return s3_key, True - - return None, True - - -def get_existing_mineru_source_s3_key(s3_key: Optional[str]) -> Optional[str]: - """Return an existing S3 source key for URL mode, or None if it is unavailable.""" - existing_s3_key, _ = _inspect_mineru_source_s3_key(s3_key) - return existing_s3_key - - -def resolve_mineru_source_s3_key( - s3_key: Optional[str], - local_file_path: Optional[str] = None, -) -> Optional[str]: - """Resolve an S3 source key for URL mode, uploading a local file if needed.""" - existing_s3_key, can_prepare_url_mode = _inspect_mineru_source_s3_key(s3_key) - if existing_s3_key is not None: - return existing_s3_key - - if not can_prepare_url_mode: - return None - - if local_file_path is None or is_remote(local_file_path): - return None - - assert s3_key is not None - from app.services.storage.sync_storage_service import upload_to_s3 - - try: - upload_to_s3(local_file_path, s3_key, settings.S3_BUCKET_NAME) - except Exception as exc: - _log_mineru_url_mode_storage_fallback( - operation="upload_source_object", - s3_key=s3_key, - local_file_path=local_file_path, - exc=exc, - ) - return None - - _mineru_logger( - "url_mode_source_uploaded", - source_s3_key=s3_key, - local_file_path=local_file_path, - ).info("Uploaded local PDF to S3 for MinerU URL mode") - return s3_key - - -def _get_retry_after_seconds( - response: requests.Response, default_retry_after: int -) -> int: - """Parse Retry-After header with sane bounds for worker backoff.""" - retry_after_header = response.headers.get("Retry-After") - if retry_after_header: - try: - return max( - 1, - min( - int(retry_after_header), settings.MINERU_RATE_LIMIT_MAX_RETRY_AFTER - ), - ) - except ValueError: - logger.debug(f"Invalid MinerU Retry-After header: {retry_after_header}") - - return max(1, min(default_retry_after, settings.MINERU_RATE_LIMIT_MAX_RETRY_AFTER)) - - -def _raise_mineru_unavailable( - token_id: str, response: requests.Response, operation: str -) -> None: - retry_after = _get_retry_after_seconds( - response, settings.MINERU_TOKEN_COOLDOWN_SECONDS - ) - quota_manager = get_mineru_quota_manager() - quota_manager.mark_rate_limited(token_id, retry_after) - _mineru_logger( - "rate_limited", - operation=operation, - token_id=token_id, - status_code=response.status_code, - retry_after=retry_after, - ).warning("MinerU request rate-limited") - raise UnavailableException( - internal_message=f"MinerU rate limited during {operation}", - retry_after=retry_after, - limit=settings.MINERU_TOKEN_RPM_LIMIT, - period="minute", - user_message="Document processing is busy right now. Please retry shortly.", - ) - - -def _polling_interval_for_state(state: str, attempt: int) -> float: - """Return seconds to sleep before the next poll. - - Tuned so that ``WORKER_CONCURRENCY`` scales safely with MinerU limits. - With a 4-token pool (300 RPM each), total budget is 1200 req/min: - - 150 tasks × (60 s / 15 s) = 600 req/min total - 600 / 4 tokens = 150 req/min per token ← leaves headroom - - Observed data (Logfire, 2026-03-08 dev batch): - - 99 % of tasks never enter ``running``; lifecycle is - ``waiting-file`` → ``done`` in 2-4 s on MinerU's side. - - Longest observed task: 21 s (5 poll attempts). - - Peak burst: 351 concurrent tasks, 308 req/min → rate-limited. - """ - if state == "pending": - return min(20.0, 5.0 + attempt * 1.5) - if state == "running": - return 10.0 - # waiting-file, converting, unknown, etc. - return 15.0 - - -def _get_batch_status(data: dict[str, Any]) -> Optional[dict[str, Any]]: - extract_result = data.get("data", {}).get("extract_result") - if isinstance(extract_result, list): - return extract_result[0] if extract_result else None - return extract_result - - -def poll_mineru_task( - status_url: str, - task_id: str, - output_dir: str, - get_status: Callable[[dict[str, Any]], Optional[dict[str, Any]]], - preferred_token_id: Optional[str] = None, -) -> None: - quota_manager = get_mineru_quota_manager() - polling_logger = _mineru_logger( - "poll_status", - operation="poll_status", - task_id=task_id, - preferred_token_id=preferred_token_id, - ) - - max_polling_attempts = 120 - polling_interval = 5.0 - max_wait_time = 6000 - - start_time = time.time() - attempt = 0 - last_token_id: Optional[str] = None - last_state: Optional[str] = None - - polling_logger.info("Starting MinerU polling") - - while attempt < max_polling_attempts: - if time.time() - start_time > max_wait_time: - polling_logger.bind( - attempt=attempt + 1, - max_polling_attempts=max_polling_attempts, - max_wait_time=max_wait_time, - ).warning("MinerU polling timed out") - raise TimeoutException( - internal_message=f"PDF parsing timed out, exceeded {max_wait_time} seconds", - retry_after=60, - user_message="PDF parsing timed out. Please try again.", - ) - - try: - logger.debug( - f"parse_pdfs status_url: {status_url} " - f"(attempt {attempt + 1}/{max_polling_attempts})" - ) - lease = quota_manager.acquire_request( - operation="poll_status", - preferred_token_id=preferred_token_id, - ) - if lease.token_id != last_token_id: - polling_logger.bind( - token_id=lease.token_id, - attempt=attempt + 1, - ).info("Acquired MinerU token for polling") - last_token_id = lease.token_id - - response = get_mineru_session().get( - status_url, - headers=get_mineru_headers(lease.api_key), - timeout=settings.MINERU_API_TIMEOUT, - ) - - if response.status_code == 429: - # urllib3 already retried with backoff — if we still see 429, - # mark the token and let Celery handle task-level retry. - _raise_mineru_unavailable( - lease.token_id, response, operation="poll_status" - ) - - if response.status_code == 200: - response_json = response.json() - if response_json.get("code") != 0: - response_message = str(response_json.get("msg") or "Unknown error") - if "rate limit" in response_message.lower(): - quota_manager.mark_rate_limited( - lease.token_id, - settings.MINERU_TOKEN_COOLDOWN_SECONDS, - ) - raise UnavailableException( - internal_message=( - f"MinerU rate limited during poll_status: {response_message}" - ), - retry_after=settings.MINERU_TOKEN_COOLDOWN_SECONDS, - limit=lease.rpm_limit, - period="minute", - user_message="Document processing is busy right now. Please retry shortly.", - ) - raise MinerUServiceException( - internal_message=f"MinerU API Error: {response_message}" - ) - - status = get_status(response_json) - if not status: - polling_logger.bind( - token_id=lease.token_id, - attempt=attempt + 1, - ).warning("Received empty MinerU status payload") - time.sleep(polling_interval) - attempt += 1 - continue - - state = status.get("state", "unknown") - if state != last_state: - # polling_logger.bind( - # token_id=lease.token_id, - # attempt=attempt + 1, - # state=state, - # ).info("MinerU status changed") - last_state = state - - if state == "done": - s3_download_extract_zip( - status["full_zip_url"], - dest_dir=output_dir, - keep_exts=(".md", ".jpg", ".jpeg", ".png", ".gif", ".json"), - exclude_patterns=("content_list", "middle.json", "model.json"), - ) - polling_logger.bind(token_id=lease.token_id).info( - "MinerU parsing completed" - ) - break - - if state == "running": - if "extract_progress" in status: - try: - ( - status["extract_progress"]["extracted_pages"] - / status["extract_progress"]["total_pages"] - ) - # polling_logger.bind( - # token_id=lease.token_id, - # progress=progress, - # ).info("MinerU parsing progress updated") - except (KeyError, ZeroDivisionError): - polling_logger.bind(token_id=lease.token_id).info( - "MinerU parsing in progress" - ) - else: - polling_logger.bind(token_id=lease.token_id).info( - "MinerU parsing in progress" - ) - elif state == "failed": - error_message = status.get("err_msg", "Unknown error") - polling_logger.bind( - token_id=lease.token_id, - error_message=error_message, - ).error("MinerU parsing reported failed state") - raise PDFParsingException( - user_message="Failed to parse the PDF file", - internal_message=f"MinerU failed with state 'failed': {error_message}", - ) - elif state == "pending": - polling_logger.bind(token_id=lease.token_id).debug( - "MinerU parsing pending" - ) - elif state == "waiting-file": - polling_logger.bind(token_id=lease.token_id).debug( - "MinerU waiting for file queueing" - ) - elif state == "converting": - polling_logger.bind(token_id=lease.token_id).debug( - "MinerU converting file" - ) - else: - polling_logger.bind( - token_id=lease.token_id, - state=state, - ).warning("MinerU returned unknown state") - - time.sleep(_polling_interval_for_state(state, attempt)) - attempt += 1 - else: - polling_logger.bind( - token_id=lease.token_id, - attempt=attempt + 1, - status_code=response.status_code, - ).warning("MinerU status query failed") - time.sleep(polling_interval * 2) - attempt += 1 - - except requests.RequestException as exc: - polling_logger.bind( - attempt=attempt + 1, - error_message=str(exc), - ).warning("MinerU polling network request failed") - time.sleep(polling_interval * 2) - attempt += 1 - except KnowhereException: - raise - except Exception as exc: - polling_logger.bind( - attempt=attempt + 1, - error_message=str(exc), - ).error("Unexpected error during MinerU polling") - raise PDFParsingException( - user_message="An unexpected error occurred while parsing the PDF", - internal_message=str(exc), - original_exception=exc, - ) - - if attempt >= max_polling_attempts: - raise TimeoutException( - internal_message=( - f"minerU PDF parsing timed out after {max_polling_attempts} attempts, " - f"Task ID: {task_id}" - ), - retry_after=60, - user_message="PDF parsing timed out. Please try again.", - ) - - -def _request_upload_target(pdf_url: str, filename: str) -> tuple[str, str, str]: - base_url = settings.MINERU_URL - quota_manager = get_mineru_quota_manager() - upload_logger = _mineru_logger( - "upload_url", - operation="upload_url", - filename=filename, - source_kind="remote_url" if is_remote(pdf_url) else "local_file", - ) - url = f"{base_url}/file-urls/batch" - payload = { - "files": [ - { - "name": filename, - "is_ocr": True, - } - ], - "enable_formula": True, - "enable_table": True, - "language": "auto", - "model_version": "vlm", - } - - upload_logger.info("Requesting MinerU upload URL") - lease = quota_manager.acquire_request(operation="upload_url") - upload_logger.bind(token_id=lease.token_id).info( - "Acquired MinerU token for upload URL" - ) - response = get_mineru_session().post( - url, - headers=get_mineru_headers(lease.api_key), - json=payload, - timeout=settings.MINERU_API_TIMEOUT, - ) - if response.status_code == 429: - _raise_mineru_unavailable(lease.token_id, response, operation="upload_url") - if response.status_code != 200: - upload_logger.bind( - token_id=lease.token_id, - status_code=response.status_code, - ).error("Failed to get MinerU upload URL") - raise MinerUServiceException( - internal_message=f"Failed to get upload URL: {response.text}", - status_code=response.status_code, - ) - - result = response.json() - if result.get("code") != 0: - response_message = str(result.get("msg", "Unknown error")) - if "rate limit" in response_message.lower(): - quota_manager.mark_rate_limited( - lease.token_id, - settings.MINERU_TOKEN_COOLDOWN_SECONDS, - ) - upload_logger.bind( - token_id=lease.token_id, - retry_after=settings.MINERU_TOKEN_COOLDOWN_SECONDS, - error_message=response_message, - ).warning("MinerU upload URL request hit rate limit") - raise UnavailableException( - internal_message=f"MinerU rate limited during upload_url: {response_message}", - retry_after=settings.MINERU_TOKEN_COOLDOWN_SECONDS, - limit=lease.rpm_limit, - period="minute", - user_message="Document processing is busy right now. Please retry shortly.", - ) - upload_logger.bind( - token_id=lease.token_id, - error_message=response_message, - ).error("MinerU upload URL request returned API error") - raise MinerUServiceException( - internal_message=f"MinerU API error: {response_message}" - ) - - batch_id = result["data"]["batch_id"] - upload_url = result["data"]["file_urls"][0] - upload_logger.bind(token_id=lease.token_id, batch_id=batch_id).info( - "Received MinerU upload URL" - ) - return batch_id, upload_url, lease.token_id - - -def _upload_file_to_mineru( - pdf_url: str, filename: str, upload_url: str, token_id: str -) -> None: - upload_logger = _mineru_logger( - "file_upload", - operation="file_upload", - filename=filename, - token_id=token_id, - source_kind="remote_url" if is_remote(pdf_url) else "local_file", - ) - - if is_remote(pdf_url): - import tempfile - - upload_logger.info("Downloading remote source file before MinerU upload") - try: - download_response = get_mineru_session().get( - pdf_url, - stream=True, - timeout=APIConstants.S3_FILE_DOWNLOAD_TIMEOUT, - ) - download_response.raise_for_status() - - with tempfile.NamedTemporaryFile( - delete=False, suffix=os.path.splitext(filename)[1] - ) as temp_file: - for chunk in download_response.iter_content(chunk_size=8192): - temp_file.write(chunk) - temp_path = temp_file.name - - upload_logger.bind(temp_file_path=temp_path).info( - "Uploading staged file to MinerU" - ) - with open(temp_path, "rb") as file_obj: - upload_response = get_mineru_session().put( - upload_url, - data=file_obj, - timeout=MINERU_UPLOAD_TIMEOUT, - ) - - os.unlink(temp_path) - except requests.RequestException as exc: - upload_logger.bind(error_message=str(exc)).error( - "Failed to stage remote source file for MinerU" - ) - raise StorageServiceException( - internal_message=f"Failed to download remote file: {exc}" - ) - else: - upload_logger.bind(local_path=pdf_url).info("Uploading local file to MinerU") - try: - with open(pdf_url, "rb") as file_obj: - try: - upload_response = get_mineru_session().put( - upload_url, - data=file_obj, - timeout=MINERU_UPLOAD_TIMEOUT, - ) - except requests.RequestException as exc: - upload_logger.bind(error_message=str(exc)).error( - "Failed to upload local file to MinerU" - ) - raise MinerUServiceException( - internal_message=f"Failed to upload file to MinerU: {exc}", - original_exception=exc, - ) from exc - except OSError as exc: - upload_logger.bind(error_message=str(exc)).error( - "Failed to read local file for MinerU upload" - ) - raise StorageServiceException( - internal_message=f"Failed to read local file: {exc}", - original_exception=exc, - ) from exc - - if upload_response.status_code != 200: - upload_logger.bind(status_code=upload_response.status_code).error( - "MinerU file upload failed" - ) - raise MinerUServiceException( - internal_message=f"Failed to upload file to MinerU: {upload_response.text}", - status_code=upload_response.status_code, - ) - - upload_logger.info("MinerU file upload completed, switching to polling") - - -def _submit_url_task(presigned_url: str, filename: str) -> tuple[str, str]: - """Submit a URL-based extraction task to MinerU. - - Uses the /extract/task/batch endpoint so MinerU fetches the file - directly from our S3 via presigned URL, skipping the OSS upload hop. - - Returns (batch_id, token_id). - """ - base_url = settings.MINERU_URL - quota_manager = get_mineru_quota_manager() - submit_logger = _mineru_logger( - "submit_url_task", - operation="submit_url_task", - filename=filename, - ) - - url = f"{base_url}/extract/task/batch" - payload = { - "files": [{"url": presigned_url}], - "is_ocr": True, - "enable_formula": True, - "enable_table": True, - "language": "auto", - "model_version": "vlm", - } - - submit_logger.info("Submitting URL-based MinerU extraction task") - lease = quota_manager.acquire_request(operation="submit_url_task") - submit_logger.bind(token_id=lease.token_id).info( - "Acquired MinerU token for URL task submission" - ) - - response = get_mineru_session().post( - url, - headers=get_mineru_headers(lease.api_key), - json=payload, - timeout=settings.MINERU_API_TIMEOUT, - ) - - if response.status_code == 429: - _raise_mineru_unavailable(lease.token_id, response, operation="submit_url_task") - - if response.status_code != 200: - submit_logger.bind( - token_id=lease.token_id, - status_code=response.status_code, - ).error("MinerU URL task submission failed") - raise MinerUServiceException( - internal_message=f"URL task submission failed: {response.text}", - status_code=response.status_code, - ) - - result = response.json() - if result.get("code") != 0: - response_message = str(result.get("msg", "Unknown error")) - if "rate limit" in response_message.lower(): - quota_manager.mark_rate_limited( - lease.token_id, - settings.MINERU_TOKEN_COOLDOWN_SECONDS, - ) - raise UnavailableException( - internal_message=f"MinerU rate limited during submit_url_task: {response_message}", - retry_after=settings.MINERU_TOKEN_COOLDOWN_SECONDS, - limit=lease.rpm_limit, - period="minute", - user_message="Document processing is busy right now. Please retry shortly.", - ) - raise MinerUServiceException( - internal_message=f"MinerU API error: {response_message}" - ) - - batch_id = result["data"]["batch_id"] - submit_logger.bind(token_id=lease.token_id, batch_id=batch_id).info( - "MinerU URL task submitted" - ) - return batch_id, lease.token_id - - -def parse_via_full( - pdf_url: str, - filename: str, - output_dir: str, - s3_key: Optional[str] = None, -) -> None: - resolved_s3_key = resolve_mineru_source_s3_key( - s3_key=s3_key, - local_file_path=None if is_remote(pdf_url) else pdf_url, - ) - - if resolved_s3_key is not None: - try: - from app.services.storage.sync_storage_service import generate_download_url - - presigned = generate_download_url( - resolved_s3_key, expires_in=settings.MINERU_URL_MODE_PRESIGN_EXPIRY - ) - presigned_url = presigned["download_url"] - _mineru_logger("ingestion_mode", mode="s3_url").info( - "Using S3 URL mode for MinerU ingestion" - ) - batch_id, token_id = _submit_url_task(presigned_url, filename) - except Exception as exc: - _log_mineru_url_mode_ingestion_fallback( - operation="start_url_mode_ingestion", - s3_key=resolved_s3_key, - pdf_url=pdf_url, - exc=exc, - ) - resolved_s3_key = None - - if resolved_s3_key is None: - _mineru_logger("ingestion_mode", mode="direct_upload").info( - "Using direct upload mode for MinerU ingestion" - ) - batch_id, upload_url, token_id = _request_upload_target(pdf_url, filename) - _upload_file_to_mineru(pdf_url, filename, upload_url, token_id) - - poll_mineru_task( - status_url=f"{settings.MINERU_URL}/extract-results/batch/{batch_id}", - task_id=batch_id, - output_dir=output_dir, - get_status=_get_batch_status, - preferred_token_id=token_id, - ) diff --git a/apps/worker/app/services/document_parser/orchestration/__init__.py b/apps/worker/app/services/document_parser/orchestration/__init__.py new file mode 100644 index 000000000..6f12d8497 --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/__init__.py @@ -0,0 +1 @@ +"""Document parser orchestration modules.""" diff --git a/apps/worker/app/services/document_parser/orchestration/format_adapters.py b/apps/worker/app/services/document_parser/orchestration/format_adapters.py new file mode 100644 index 000000000..6dc3409c2 --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/format_adapters.py @@ -0,0 +1,219 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Protocol + +from app.services.document_parser.orchestration.parse_output import ParseOutput +from app.services.document_parser.orchestration.parse_session import ParseSession + + +class DocumentParseAdapter(Protocol): + @property + def document_format(self) -> object: + """Document format handled by this adapter.""" + + def parse(self, session: ParseSession) -> ParseOutput: + """Parse a document session into a stable parser output object.""" + raise NotImplementedError + + +@dataclass(frozen=True) +class FragmentParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> ParseOutput: + from app.services.document_parser.formats.fragment.parser import parse_fragment + + full_output_dir, _relative_root, parsed_df = parse_fragment( + session.fragment_content, + filename=session.filename, + output_dir=session.output_dir, + base_llm_paras=session.base_llm_paras, + ) + return ParseOutput(output_dir=full_output_dir, parsed_df=parsed_df) + + +@dataclass(frozen=True) +class TextParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> ParseOutput: + from app.services.document_parser.formats.markdown.parser import parse_md + from app.services.document_parser.formats.text.parser import parse_texts + + text_lines = parse_texts(file_path=session.file_full_path, baseurl=session.base_url) + parsed_df = parse_md( + session.full_output_dir, + source_type="md", + md_lines=text_lines, + base_llm_paras=session.base_llm_paras, + relative_root=session.relative_root, + ) + return ParseOutput(output_dir=session.full_output_dir, parsed_df=parsed_df) + + +@dataclass(frozen=True) +class ImageParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> ParseOutput: + from app.services.document_parser.formats.image.parser import parse_image + + parsed_df = parse_image( + session.file_full_path, + filename=session.filename, + output_dir=session.full_output_dir, + baseurl=session.base_url, + base_llm_paras=session.base_llm_paras, + relative_root=session.relative_root, + ) + return ParseOutput(output_dir=session.full_output_dir, parsed_df=parsed_df) + + +@dataclass(frozen=True) +class PdfParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> ParseOutput: + from app.services.document_parser.formats.pdf.parser import parse_pdfs + + parsed_df = parse_pdfs( + session.file_full_path, + filename=session.filename, + output_dir=session.full_output_dir, + base_llm_paras=session.base_llm_paras, + profile=session.profile, + relative_root=session.relative_root, + s3_key=session.s3_key, + ) + return ParseOutput(output_dir=session.full_output_dir, parsed_df=parsed_df) + + +@dataclass(frozen=True) +class DocParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> ParseOutput: + from app.services.document_parser.conversion.legacy_converter import doc_to_docx + + converted_docx_path, _ = doc_to_docx( + session.file_full_path, + outdir=session.full_output_dir, + ) + return _parse_docx_path(converted_docx_path, session) + + +@dataclass(frozen=True) +class DocxParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> ParseOutput: + return _parse_docx_path(session.file_full_path, session) + + +@dataclass(frozen=True) +class XlsParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> ParseOutput: + from app.services.document_parser.conversion.legacy_converter import xls_to_xlsx + + converted_xlsx_path, _ = xls_to_xlsx( + session.file_full_path, + outdir=session.full_output_dir, + ) + return _parse_xlsx_path(converted_xlsx_path, session) + + +@dataclass(frozen=True) +class XlsxParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> ParseOutput: + return _parse_xlsx_path(session.file_full_path, session) + + +@dataclass(frozen=True) +class PptxParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> ParseOutput: + from app.services.document_parser.formats.pptx.parser import parse_pptx + + parsed_df = parse_pptx( + session.file_full_path, + filename=session.filename, + output_dir=session.full_output_dir, + base_llm_paras=session.base_llm_paras, + strategy="to_pdf_api", + job_id=session.job_id, + relative_root=session.relative_root, + baseurl=session.base_url, + ) + return ParseOutput(output_dir=session.full_output_dir, parsed_df=parsed_df) + + +@dataclass(frozen=True) +class MarkdownParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> ParseOutput: + from app.services.document_parser.formats.markdown.parser import parse_md + + parsed_df = parse_md( + session.full_output_dir, + source_type="md", + file_path=session.file_full_path, + base_llm_paras=session.base_llm_paras, + relative_root=session.relative_root, + ) + return ParseOutput(output_dir=session.full_output_dir, parsed_df=parsed_df) + + +@dataclass(frozen=True) +class JsonParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> ParseOutput: + return ParseOutput(output_dir=session.full_output_dir, parsed_df=None) + + +def _parse_docx_path( + docx_path: str, + session: ParseSession, +) -> ParseOutput: + from app.services.document_parser.formats.docx.parser import convert_doc2dics, parse_docx + + parsed_structure, dataframe_list = parse_docx( + docx_path, + session.base_llm_paras, + session.full_output_dir, + session.filename, + session.base_url, + relative_root=session.relative_root, + ) + parsed_df = convert_doc2dics( + parsed_structure, + dataframe_list, + session.full_output_dir, + base_llm_paras=session.base_llm_paras, + relative_root=session.relative_root, + ) + return ParseOutput(output_dir=session.full_output_dir, parsed_df=parsed_df) + + +def _parse_xlsx_path( + xlsx_path: str, + session: ParseSession, +) -> ParseOutput: + from app.services.document_parser.formats.excel.table_parser import parse_xlsx + + parsed_df = parse_xlsx( + xlsx_path, + session.filename, + session.full_output_dir, + session.base_url, + base_llm_paras=session.base_llm_paras, + relative_root=session.relative_root, + ) + return ParseOutput(output_dir=session.full_output_dir, parsed_df=parsed_df) diff --git a/apps/worker/app/services/document_parser/orchestration/format_router.py b/apps/worker/app/services/document_parser/orchestration/format_router.py new file mode 100644 index 000000000..b94cf2787 --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/format_router.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +import os +from enum import Enum + +from app.services.document_parser.orchestration import format_adapters +from app.services.document_parser.orchestration.format_adapters import ( + DocumentParseAdapter, +) +from shared.core.exceptions.domain_exceptions import ValidationException + + +class DocumentFormat(str, Enum): + TEXT = "text" + FRAGMENT = "fragment" + IMAGE = "image" + PDF = "pdf" + DOC = "doc" + DOCX = "docx" + XLS = "xls" + XLSX = "xlsx" + PPTX = "pptx" + MARKDOWN = "markdown" + JSON = "json" + + +SUPPORTED_FILE_TYPES: tuple[str, ...] = ( + ".txt", + ".fragment", + ".png", + ".jpg", + ".jpeg", + ".pdf", + ".doc", + ".docx", + ".xls", + ".xlsx", + ".pptx", + ".md", + ".json", +) + + +def resolve_document_format(file_path: str) -> DocumentFormat: + extension = os.path.splitext(file_path)[1].lower() + if extension == ".fragment": + return DocumentFormat.FRAGMENT + if extension == ".txt": + return DocumentFormat.TEXT + if extension in (".png", ".jpg", ".jpeg"): + return DocumentFormat.IMAGE + if extension == ".pdf": + return DocumentFormat.PDF + if extension == ".doc": + return DocumentFormat.DOC + if extension == ".docx": + return DocumentFormat.DOCX + if extension == ".xls": + return DocumentFormat.XLS + if extension == ".xlsx": + return DocumentFormat.XLSX + if extension == ".pptx": + return DocumentFormat.PPTX + if extension == ".md": + return DocumentFormat.MARKDOWN + if extension == ".json": + return DocumentFormat.JSON + + raise ValidationException( + user_message=f"Unsupported file type: {extension}", + violations=[ + { + "field": "file_type", + "description": f"Must be one of: {', '.join(SUPPORTED_FILE_TYPES)}", + } + ], + ) + + +def get_document_parse_adapter(document_format: DocumentFormat) -> DocumentParseAdapter: + adapter_by_format: dict[DocumentFormat, DocumentParseAdapter] = { + DocumentFormat.FRAGMENT: format_adapters.FragmentParseAdapter(document_format), + DocumentFormat.TEXT: format_adapters.TextParseAdapter(document_format), + DocumentFormat.IMAGE: format_adapters.ImageParseAdapter(document_format), + DocumentFormat.PDF: format_adapters.PdfParseAdapter(document_format), + DocumentFormat.DOC: format_adapters.DocParseAdapter(document_format), + DocumentFormat.DOCX: format_adapters.DocxParseAdapter(document_format), + DocumentFormat.XLS: format_adapters.XlsParseAdapter(document_format), + DocumentFormat.XLSX: format_adapters.XlsxParseAdapter(document_format), + DocumentFormat.PPTX: format_adapters.PptxParseAdapter(document_format), + DocumentFormat.MARKDOWN: format_adapters.MarkdownParseAdapter(document_format), + DocumentFormat.JSON: format_adapters.JsonParseAdapter(document_format), + } + return adapter_by_format[document_format] diff --git a/apps/worker/app/services/document_parser/orchestration/parse_input.py b/apps/worker/app/services/document_parser/orchestration/parse_input.py new file mode 100644 index 000000000..98dd9fa2c --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/parse_input.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass(frozen=True) +class ParseOptions: + llm_histories: int = 5 + smart_title_parse: bool = True + summary_image: bool = True + summary_table: bool = True + summary_txt: bool = True + stopwords: list[str] | None = None + doc_type: str = "auto" + add_frag_desc: str = "" + + +@dataclass(frozen=True) +class ParseInput: + file_full_path: str + filename: str + output_dir: str + internal_output_filename: str + job_id: str | None = None + options: ParseOptions = field(default_factory=ParseOptions) + base_url: str = "" + fragment_content: str = "" + s3_key: str | None = None diff --git a/apps/worker/app/services/document_parser/orchestration/parse_output.py b/apps/worker/app/services/document_parser/orchestration/parse_output.py new file mode 100644 index 000000000..8e70fd56d --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/parse_output.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import pandas as pd + + +@dataclass(frozen=True) +class ParseOutput: + """Parser adapter output.""" + + output_dir: str + parsed_df: pd.DataFrame | None + + @property + def rows_count(self) -> int: + if self.parsed_df is None: + return 0 + return len(self.parsed_df) + + def with_dataframe(self, parsed_df: pd.DataFrame | None) -> ParseOutput: + return ParseOutput(output_dir=self.output_dir, parsed_df=parsed_df) diff --git a/apps/worker/app/services/document_parser/orchestration/parse_pipeline.py b/apps/worker/app/services/document_parser/orchestration/parse_pipeline.py new file mode 100644 index 000000000..909a31197 --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/parse_pipeline.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +from app.services.document_parser.orchestration.parse_input import ParseInput +from app.services.document_parser.orchestration.parse_output import ParseOutput +from app.services.document_parser.orchestration.parse_session import build_parse_session +from app.services.document_parser.orchestration.postprocess import apply_parse_postprocess +from app.services.document_parser.orchestration.route_parse import route_document_parse + + +ParsePipelineResult = ParseOutput + + +def run_parse_pipeline(parse_input: ParseInput) -> ParsePipelineResult: + """Run parser session building, format routing, and output postprocessing.""" + session = build_parse_session(parse_input) + parsed_output = route_document_parse(session) + processed_df = apply_parse_postprocess( + parsed_output.output_dir, + parsed_output.parsed_df, + ) + return parsed_output.with_dataframe(processed_df) diff --git a/apps/worker/app/services/document_parser/orchestration/parse_session.py b/apps/worker/app/services/document_parser/orchestration/parse_session.py new file mode 100644 index 000000000..3fe956f1a --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/parse_session.py @@ -0,0 +1,205 @@ +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Any + +from app.services.document_parser.formats.atlas.classifier import classify_atlas_with_vlm +from app.services.document_parser.orchestration.path_segment import ( + build_parser_path_segment, +) +from app.services.document_parser.orchestration.parse_input import ParseInput +from app.services.document_parser.profiling.doc_profiler import profile_document +from app.services.document_parser.support.stage_profiler import stage_timer +from loguru import logger + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ValidationException + +PDF_PAGE_LIMIT = 600 + + +@dataclass(frozen=True) +class ParseSession: + base_llm_paras: dict[str, object] + base_url: str + file_full_path: str + filename: str + fragment_content: str + full_output_dir: str + internal_output_filename: str + job_id: str | None + output_dir: str + profile: Any + relative_root: str + s3_key: str | None + + @classmethod + def from_input( + cls, + *, + parse_input: ParseInput, + base_llm_paras: dict[str, object], + full_output_dir: str, + profile: Any, + relative_root: str, + ) -> "ParseSession": + return cls( + base_llm_paras=base_llm_paras, + base_url=parse_input.base_url, + file_full_path=parse_input.file_full_path, + filename=parse_input.filename, + fragment_content=parse_input.fragment_content, + full_output_dir=full_output_dir, + internal_output_filename=parse_input.internal_output_filename, + job_id=parse_input.job_id, + output_dir=parse_input.output_dir, + profile=profile, + relative_root=relative_root, + s3_key=parse_input.s3_key, + ) + + +def build_parse_session(parse_input: ParseInput) -> ParseSession: + """Build the parser routing session from explicit parse inputs.""" + parse_options = parse_input.options + base_llm_paras = { + "llm_histories": parse_options.llm_histories, + "smart_title_parse": parse_options.smart_title_parse, + "summary_image": parse_options.summary_image, + "summary_table": parse_options.summary_table, + "summary_txt": parse_options.summary_txt, + "stopwords": parse_options.stopwords, + "doc_type": parse_options.doc_type, + "frag_desc": parse_options.add_frag_desc, + "model_name": settings.NORMOL_MODEL, + "hierarchy_model_name": settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL, + } + + logger.debug(f"baseurl: {parse_input.base_url}") + logger.debug(f"file_full_path: {parse_input.file_full_path}") + + relative_root, full_output_dir = _resolve_output_paths( + filename=parse_input.filename, + internal_output_filename=parse_input.internal_output_filename, + output_dir=parse_input.output_dir, + ) + logger.debug(f"relative_root: {relative_root}") + logger.debug(f"full_output_dir: {full_output_dir}") + + with stage_timer("document.profile", filename=parse_input.filename): + profile = profile_document( + parse_input.file_full_path, + parse_input.internal_output_filename, + ) + logger.info(f"📋 DocProfile: {profile.summary()}") + logger.debug(f"📋 Reasoning: {profile.reasoning}") + + if profile.atlas_candidate and profile.doc_category not in ("atlas", "ppt_converted"): + logger.info( + f"🔍 Atlas candidate detected, running VLM visual check for {parse_input.filename}" + ) + with stage_timer("document.atlas_vlm_check", filename=parse_input.filename): + vlm_is_atlas = classify_atlas_with_vlm(parse_input.file_full_path) + if vlm_is_atlas: + profile.doc_category = "atlas" + profile.reasoning += " | vlm_confirmed_atlas=True" + logger.info(f"✅ VLM confirmed atlas for {parse_input.filename}") + else: + profile.reasoning += " | vlm_confirmed_atlas=False" + logger.info( + f"ℹ️ VLM rejected atlas for {parse_input.filename}, routing as generic" + ) + + if profile.file_type == "pdf" and profile.page_count > PDF_PAGE_LIMIT: + raise ValidationException( + user_message=( + f"Document too large: {profile.page_count} pages exceeds the {PDF_PAGE_LIMIT}-page limit. " + "Please split the document and upload in smaller batches." + ), + violations=[ + { + "field": "page_count", + "description": f"PDF has {profile.page_count} pages, limit is {PDF_PAGE_LIMIT}", + } + ], + ) + + if profile.doc_category == "atlas": + filename, internal_output_filename, relative_root, full_output_dir = ( + _rename_atlas_output( + filename=parse_input.filename, + internal_output_filename=parse_input.internal_output_filename, + output_dir=parse_input.output_dir, + ) + ) + logger.info(f"📐 Atlas output renamed: {filename}") + parse_input = ParseInput( + file_full_path=parse_input.file_full_path, + filename=filename, + output_dir=parse_input.output_dir, + internal_output_filename=internal_output_filename, + job_id=parse_input.job_id, + options=parse_input.options, + base_url=parse_input.base_url, + fragment_content=parse_input.fragment_content, + s3_key=parse_input.s3_key, + ) + + return ParseSession.from_input( + parse_input=parse_input, + base_llm_paras=base_llm_paras, + full_output_dir=full_output_dir, + profile=profile, + relative_root=relative_root, + ) + + +def _rename_atlas_output( + *, + filename: str, + internal_output_filename: str, + output_dir: str, +) -> tuple[str, str, str, str]: + name_base, _ = os.path.splitext(filename) + internal_name_base, _ = os.path.splitext(internal_output_filename) + atlas_filename = name_base + ".atlas" + atlas_internal_filename = internal_name_base + ".atlas" + relative_root, full_output_dir = _resolve_output_paths( + filename=atlas_filename, + internal_output_filename=atlas_internal_filename, + output_dir=output_dir, + ) + return atlas_filename, atlas_internal_filename, relative_root, full_output_dir + + +def _resolve_output_paths( + *, + filename: str, + internal_output_filename: str, + output_dir: str, +) -> tuple[str, str]: + filename_segment = build_parser_path_segment(filename) + internal_filename_segment = build_parser_path_segment( + internal_output_filename, + default=filename_segment, + ) + relative_root = filename_segment + + full_output_dir = os.path.join( + output_dir, + internal_filename_segment, + ) + resolved_output_dir = os.path.realpath(output_dir) + resolved_full_output_dir = os.path.realpath(full_output_dir) + if ( + os.path.commonpath([resolved_output_dir, resolved_full_output_dir]) + != resolved_output_dir + ): + raise ValueError( + f"Parser output directory escaped task workspace: {full_output_dir}" + ) + os.makedirs(resolved_full_output_dir, exist_ok=True) + + logger.debug(f"internal_output_root: {internal_filename_segment}") + return relative_root, resolved_full_output_dir diff --git a/apps/worker/app/services/document_parser/orchestration/path_segment.py b/apps/worker/app/services/document_parser/orchestration/path_segment.py new file mode 100644 index 000000000..ed03bf9eb --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/path_segment.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +import os + +from app.services.common.file_utils import path_handle + + +def build_parser_path_segment(value: str | None, default: str = "document") -> str: + """Map parser-owned names to one safe task-local path segment.""" + raw_value = str(value or "").strip() + raw_segment = os.path.basename(raw_value) if raw_value else default + sanitized_segment = path_handle(raw_segment, mode="clean_single") + if not isinstance(sanitized_segment, str): + return default + + segment = sanitized_segment.strip() + if segment in {"", ".", ".."}: + return default + return segment diff --git a/apps/worker/app/services/document_parser/orchestration/postprocess.py b/apps/worker/app/services/document_parser/orchestration/postprocess.py new file mode 100644 index 000000000..0fbb9f5c8 --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/postprocess.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import os +import re + +import pandas as pd +from app.services.document_parser.assets.image_compressor import ( + apply_rename_map_to_dataframe, + compress_output_images, +) +from app.services.document_parser.support.stage_profiler import stage_timer +from loguru import logger + + +def apply_parse_postprocess( + output_dir: str, + parsed_df: pd.DataFrame | None, +) -> pd.DataFrame | None: + """Apply output cleanup and image compression after parsing.""" + logger.debug(f"full_output_dir: {output_dir}") + + with stage_timer("document.cleanup_unreferenced_images", output_dir=output_dir): + cleanup_unreferenced_images(output_dir) + + with stage_timer("document.compress_images", output_dir=output_dir): + compress_stats = compress_output_images(output_dir) + if compress_stats.processed > 0: + logger.info( + f"📦 Image compression: {compress_stats.processed} processed " + f"({compress_stats.converted_png_to_jpg} PNG→JPG, " + f"{compress_stats.resized} resized), " + f"{compress_stats.bytes_before / 1024 / 1024:.1f}MB → " + f"{compress_stats.bytes_after / 1024 / 1024:.1f}MB" + ) + if compress_stats.rename_map and parsed_df is not None: + return apply_rename_map_to_dataframe(parsed_df, compress_stats.rename_map) + + return parsed_df + + +def cleanup_unreferenced_images(output_dir: str) -> int: + """Remove UUID-named images that are not referenced by final parsed output.""" + image_dir = os.path.join(output_dir, "images") + if not os.path.isdir(image_dir): + return 0 + + uuid_pattern = re.compile( + r"^[a-f0-9]{64}\.(?:jpg|jpeg|png|gif|webp)$", + re.IGNORECASE, + ) + removed_count = 0 + + for filename in os.listdir(image_dir): + if not uuid_pattern.match(filename): + continue + + file_path = os.path.join(image_dir, filename) + try: + os.remove(file_path) + removed_count += 1 + logger.debug(f"Removed unreferenced image: {filename}") + except OSError as exc: + logger.warning(f"Failed to remove {filename}: {exc}") + + if removed_count > 0: + logger.info( + f"Cleaned up {removed_count} unreferenced UUID-named images from {image_dir}" + ) + + return removed_count diff --git a/apps/worker/app/services/document_parser/orchestration/route_parse.py b/apps/worker/app/services/document_parser/orchestration/route_parse.py new file mode 100644 index 000000000..68bb57e73 --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/route_parse.py @@ -0,0 +1,13 @@ +from app.services.document_parser.orchestration.format_router import ( + get_document_parse_adapter, + resolve_document_format, +) +from app.services.document_parser.orchestration.parse_output import ParseOutput +from app.services.document_parser.orchestration.parse_session import ParseSession + + +def route_document_parse(session: ParseSession) -> ParseOutput: + """Route a parser session to the correct adapter and return its output.""" + document_format = resolve_document_format(session.file_full_path) + adapter = get_document_parse_adapter(document_format) + return adapter.parse(session) diff --git a/apps/worker/app/services/document_parser/parse_service.py b/apps/worker/app/services/document_parser/parse_service.py index ad5058e47..ba3d9f6ec 100644 --- a/apps/worker/app/services/document_parser/parse_service.py +++ b/apps/worker/app/services/document_parser/parse_service.py @@ -1,77 +1,18 @@ -# pyright: reportArgumentType=false, reportReturnType=false -""" -main parsing service -""" +"""Stable parser seam backed by dedicated orchestration modules.""" -import os -import re - -import pandas as pd -from app.services.document_parser.atlas_classifier import classify_atlas_with_vlm - -# document_parser imports -from app.services.document_parser.doc_profiler import profile_document -from app.services.document_parser.stage_profiler import stage_timer -from loguru import logger - -from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import ( - ValidationException, +from app.services.document_parser.orchestration.parse_input import ParseInput, ParseOptions +from app.services.document_parser.orchestration.parse_pipeline import ( + ParsePipelineResult, + run_parse_pipeline, ) -from shared.utils.file_utils import path_handle - - -def cleanup_unreferenced_images(output_dir: str) -> int: - """ - Clean up unreferenced UUID-named images from the images directory. - - After document parsing (PDF, DOCX, PPTX, etc.), the images/ directory may contain: - 1. Processed images: renamed with semantic names like 'image-0-xxx.jpg' - 2. Unreferenced images: UUID-named (64-char hex) that were parsed as tables/formulas - - This function removes the unreferenced UUID-named images to reduce final package size. - - Args: - output_dir: The full output directory path - Returns: - Number of files removed - """ - img_dir = os.path.join(output_dir, "images") - if not os.path.isdir(img_dir): - return 0 - # UUID pattern: 64 hex characters followed by image extension - uuid_pattern = re.compile( - r"^[a-f0-9]{64}\.(?:jpg|jpeg|png|gif|webp)$", re.IGNORECASE - ) - - removed_count = 0 - for filename in os.listdir(img_dir): - if uuid_pattern.match(filename): - file_path = os.path.join(img_dir, filename) - try: - os.remove(file_path) - removed_count += 1 - logger.debug(f"Removed unreferenced image: {filename}") - except OSError as e: - logger.warning(f"Failed to remove {filename}: {e}") - - if removed_count > 0: - logger.info( - f"Cleaned up {removed_count} unreferenced UUID-named images from {img_dir}" - ) - - return removed_count - - -def checkerboard_inject_parse( +def checkerboard_parse_output( file_full_path: str, filename: str, output_dir: str, internal_output_filename: str, job_id: str | None = None, - kb_dir: str = "Default_Root", llm_histories: int = 5, smart_title_parse: bool = True, summary_image: bool = True, @@ -83,375 +24,26 @@ def checkerboard_inject_parse( base_url: str = "", fragment_content: str = "", s3_key: str | None = None, -) -> tuple[str, pd.DataFrame | None]: - """ - main parsing function - - Args: - file_full_path: source file path (local or URL) - filename: file name - output_dir: output directory (absolute path, caller provides) - kb_dir: sub-directory name - llm_histories: retained for downstream LLM settings - smart_title_parse: enable smart heading parsing - summary_image: enable image summaries - summary_table: enable table summaries - summary_txt: enable text summaries - stopwords: optional stopword list - doc_type: parser document type hint - add_frag_desc: extra fragment description - base_url: optional source base URL - fragment_content: raw fragment content - job_id: optional job identifier used for parser artifacts - internal_output_filename: normalized internal folder name for on-disk output - s3_key: optional S3 key for downstream parsers - - Returns: - tuple: (output_dir, parsed_df) - - output_dir: directory path after parsing - - parsed_df: parsed content DataFrame - """ - # Build base_llm_paras from explicit parameters - base_llm_paras = { - "llm_histories": llm_histories, - "smart_title_parse": smart_title_parse, - "summary_image": summary_image, - "summary_table": summary_table, - "summary_txt": summary_txt, - "stopwords": stopwords, - "doc_type": doc_type, - "frag_desc": add_frag_desc, - "model_name": settings.NORMOL_MODEL, - "hierarchy_model_name": settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL, - } - - baseurl = base_url - - logger.debug(f"baseurl: {baseurl}") - logger.debug(f"file_full_path: {file_full_path}") - - # ========== Path handling ========== - split_char = settings.SPLIT_CHAR or "/" - - # Develop relative root path for chunk path field - kb_dir_parts = kb_dir.split(split_char) - if filename and "images" not in kb_dir_parts: - relative_root = "/".join(kb_dir_parts + [filename]) - else: - relative_root = "/".join(kb_dir_parts) - - if internal_output_filename and "images" not in kb_dir_parts: - internal_relative_root = "/".join(kb_dir_parts + [internal_output_filename]) - else: - internal_relative_root = "/".join(kb_dir_parts) - - # Develop full output directory (output_dir + relative_root) - full_output_dir = os.path.join( - output_dir, internal_relative_root.replace("/", os.sep) +) -> ParsePipelineResult: + """Run the stable parser seam and return the parser output contract.""" + parse_input = ParseInput( + file_full_path=file_full_path, + filename=filename, + internal_output_filename=internal_output_filename, + job_id=job_id, + output_dir=output_dir, + options=ParseOptions( + add_frag_desc=add_frag_desc, + doc_type=doc_type, + llm_histories=llm_histories, + smart_title_parse=smart_title_parse, + stopwords=stopwords, + summary_image=summary_image, + summary_table=summary_table, + summary_txt=summary_txt, + ), + base_url=base_url, + fragment_content=fragment_content, + s3_key=s3_key, ) - full_output_dir = path_handle(full_output_dir, mode="sanitize") - os.makedirs(full_output_dir, exist_ok=True) - - logger.debug(f"relative_root: {relative_root}") - logger.debug(f"internal_relative_root: {internal_relative_root}") - logger.debug(f"full_output_dir: {full_output_dir}") - - file_path_lower = file_full_path.lower() - parsed_df = None - - # ── Agentic Profiler: classify document before routing ── - with stage_timer("document.profile", filename=filename): - profile = profile_document(file_full_path, internal_output_filename) - logger.info(f"📋 DocProfile: {profile.summary()}") - logger.debug(f"📋 Reasoning: {profile.reasoning}") - - # ── VLM second-pass: confirm atlas_candidate with visual check ── - # Heuristics can miss atlases that have a rich OCR text layer on top of - # scanned drawing pages (avg_text_density too high). VLM sees the actual - # page layout and makes the final call. - if profile.atlas_candidate and profile.doc_category not in ( - "atlas", - "ppt_converted", - ): - logger.info( - f"🔍 Atlas candidate detected, running VLM visual check for {filename}" - ) - with stage_timer("document.atlas_vlm_check", filename=filename): - vlm_is_atlas = classify_atlas_with_vlm(file_full_path) - if vlm_is_atlas: - profile.doc_category = "atlas" - profile.reasoning += " | vlm_confirmed_atlas=True" - logger.info(f"✅ VLM confirmed atlas for {filename}") - else: - profile.reasoning += " | vlm_confirmed_atlas=False" - logger.info(f"ℹ️ VLM rejected atlas for {filename}, routing as generic") - - # ── Page count guard: reject oversized PDFs before routing ── - PDF_PAGE_LIMIT = 600 - if profile and profile.file_type == "pdf" and profile.page_count > PDF_PAGE_LIMIT: - raise ValidationException( - user_message=( - f"Document too large: {profile.page_count} pages exceeds the {PDF_PAGE_LIMIT}-page limit. " - f"Please split the document and upload in smaller batches." - ), - violations=[ - { - "field": "page_count", - "description": f"PDF has {profile.page_count} pages, limit is {PDF_PAGE_LIMIT}", - } - ], - ) - - # Atlas routing: rename output folder from .pdf → .atlas for easy filtering - if profile and profile.doc_category == "atlas": - name_base, _ = os.path.splitext(filename) - internal_name_base, _ = os.path.splitext(internal_output_filename) - filename = name_base + ".atlas" - internal_output_filename = internal_name_base + ".atlas" - relative_root = "/".join(kb_dir_parts + [filename]) - internal_relative_root = "/".join(kb_dir_parts + [internal_output_filename]) - full_output_dir = os.path.join( - output_dir, internal_relative_root.replace("/", os.sep) - ) - full_output_dir = path_handle(full_output_dir, mode="sanitize") - os.makedirs(full_output_dir, exist_ok=True) - logger.info(f"📐 Atlas output renamed: {filename}") - - if ".fragment" in file_path_lower: - logger.debug("file type is fragment") - from app.services.document_parser.fragment_parser import parse_fragment - - full_output_dir, relative_root, parsed_df = parse_fragment( - fragment_content, - filename=filename, - output_dir=output_dir, - kb_dir=kb_dir, - base_llm_paras=base_llm_paras, - ) - - elif ".txt" in file_path_lower: - logger.debug("file type is txt") - from app.services.document_parser.md_parser import parse_md - from app.services.document_parser.txt_parser import parse_texts - - txt_lines = parse_texts(file_path=file_full_path, baseurl=baseurl) - parsed_df = parse_md( - full_output_dir, - source_type="md", - md_lines=txt_lines, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - ) - - elif ( - ".png" in file_path_lower - or ".jpg" in file_path_lower - or ".jpeg" in file_path_lower - ): - logger.debug("file type is image") - from app.services.document_parser.image_parser import parse_image - - parsed_df = parse_image( - file_full_path, - filename=filename, - output_dir=full_output_dir, - baseurl=baseurl, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - ) - - elif ".pdf" in file_path_lower: - logger.debug("file type is pdf") - from app.services.document_parser.pdf_parser import parse_pdfs - - if filename and file_full_path: - parsed_df = parse_pdfs( - file_full_path, - filename=filename, - output_dir=full_output_dir, - base_llm_paras=base_llm_paras, - profile=profile, - relative_root=relative_root, - s3_key=s3_key, - ) - - elif ".doc" in file_path_lower and ".docx" not in file_path_lower: - logger.debug("file type is doc") - from app.services.document_parser.doc_parser import convert_doc2dics, parse_docx - from app.services.document_parser.legacy_converter import doc_to_docx - - if filename and file_full_path: - converted_docx_path, _ = doc_to_docx(file_full_path, outdir=full_output_dir) - parsed_structure, df_list = parse_docx( - converted_docx_path, - base_llm_paras, - full_output_dir, - filename, - baseurl, - relative_root=relative_root, - ) - parsed_df = convert_doc2dics( - parsed_structure, - df_list, - full_output_dir, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - ) - - elif ".docx" in file_path_lower: - logger.debug("file type is docx") - from app.services.document_parser.doc_parser import convert_doc2dics, parse_docx - - if filename and file_full_path: - parsed_structure, df_list = parse_docx( - file_full_path, - base_llm_paras, - full_output_dir, - filename, - baseurl, - relative_root=relative_root, - ) - parsed_df = convert_doc2dics( - parsed_structure, - df_list, - full_output_dir, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - ) - - elif ".xls" in file_path_lower and ".xlsx" not in file_path_lower: - logger.debug("file type is xls") - from app.services.document_parser.legacy_converter import xls_to_xlsx - from app.services.document_parser.table_parser import parse_xlsx - - if filename and file_full_path: - converted_xlsx_path, _ = xls_to_xlsx(file_full_path, outdir=full_output_dir) - parsed_df = parse_xlsx( - converted_xlsx_path, - filename, - full_output_dir, - baseurl, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - ) - - elif ".xlsx" in file_path_lower: - logger.debug("file type is xlsx") - from app.services.document_parser.table_parser import parse_xlsx - - if filename and file_full_path: - parsed_df = parse_xlsx( - file_full_path, - filename, - full_output_dir, - baseurl, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - ) - - elif ".pptx" in file_path_lower: - logger.debug("file type is pptx") - from app.services.document_parser.pptx_parser import parse_pptx - - if filename and file_full_path: - # ====== iLoveAPI PPTX → PDF → MinerU (default production route) ====== - parsed_df = parse_pptx( - file_full_path, - filename=filename, - output_dir=full_output_dir, - base_llm_paras=base_llm_paras, - strategy="to_pdf_api", - job_id=job_id, - relative_root=relative_root, - baseurl=baseurl, - ) - - # ====== [EXPERIMENTAL] Directly send PPTX to MinerU via parse_pdfs ====== - # Uncomment the block below (and comment out parse_pptx above) to bypass iLoveAPI - # from app.services.document_parser.pdf_parser import parse_pdfs - # parsed_df = parse_pdfs( - # file_full_path, - # filename=filename, - # output_dir=full_output_dir, - # base_llm_paras=base_llm_paras, - # profile=profile, - # relative_root=relative_root, - # s3_key=s3_key - # ) - - elif ".md" in file_path_lower: - logger.debug("file type is md") - from app.services.document_parser.md_parser import parse_md - - if filename and file_full_path: - parsed_df = parse_md( - full_output_dir, - source_type="md", - file_path=file_full_path, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - ) - - elif ".json" in file_path_lower: - logger.debug("file type is json") - # JSON parsing not yet implemented - - else: - # Unsupported file type - file_ext = os.path.splitext(file_full_path)[1].lower() - supported_types = [ - ".txt", - ".fragment", - ".png", - ".jpg", - ".jpeg", - ".pdf", - ".doc", - ".docx", - ".xls", - ".xlsx", - ".pptx", - ".md", - ".json", - ] - raise ValidationException( - user_message=f"Unsupported file type: {file_ext}", - violations=[ - { - "field": "file_type", - "description": f"Must be one of: {', '.join(supported_types)}", - } - ], - ) - - logger.debug(f"full_output_dir: {full_output_dir}") - - # Post-processing: clean up unreferenced UUID-named images - with stage_timer( - "document.cleanup_unreferenced_images", output_dir=full_output_dir - ): - cleanup_unreferenced_images(full_output_dir) - - # Post-processing: compress output images (PNG→JPEG, resize oversized) - from app.services.document_parser.image_compressor import ( - apply_rename_map_to_dataframe, - compress_output_images, - ) - - with stage_timer("document.compress_images", output_dir=full_output_dir): - compress_stats = compress_output_images(full_output_dir) - if compress_stats.processed > 0: - logger.info( - f"📦 Image compression: {compress_stats.processed} processed " - f"({compress_stats.converted_png_to_jpg} PNG→JPG, " - f"{compress_stats.resized} resized), " - f"{compress_stats.bytes_before / 1024 / 1024:.1f}MB → " - f"{compress_stats.bytes_after / 1024 / 1024:.1f}MB" - ) - # Update DataFrame references when PNG→JPG conversions occurred - if compress_stats.rename_map and parsed_df is not None: - parsed_df = apply_rename_map_to_dataframe( - parsed_df, compress_stats.rename_map - ) - - return full_output_dir, parsed_df + return run_parse_pipeline(parse_input) diff --git a/apps/worker/app/services/document_parser/profiling/doc_profile_model.py b/apps/worker/app/services/document_parser/profiling/doc_profile_model.py new file mode 100644 index 000000000..ad1a12e65 --- /dev/null +++ b/apps/worker/app/services/document_parser/profiling/doc_profile_model.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import gc +import json +import os +from dataclasses import asdict, dataclass, field +from typing import List, Literal, Optional + +from loguru import logger + + +@dataclass +class DocProfile: + """Document profile data contract used by parser routing.""" + + file_type: str = "" + route: Literal["fast", "standard"] = "standard" + decision_band: Literal["safe_fast", "gray_zone", "safe_standard"] = "safe_standard" + scan_type: Optional[Literal["electronic", "scanned", "mixed"]] = None + doc_category: Literal["generic", "atlas", "ppt_converted"] = "generic" + page_count: int = 0 + avg_text_density: float = 0.0 + avg_image_coverage: float = 0.0 + has_tables: bool = False + has_embedded_fonts: bool = False + is_multi_column: bool = False + is_degraded_electronic: bool = False + sample_text: str = "" + has_significant_images: bool = False + significant_image_count: int = 0 + max_image_coverage_on_page: float = 0.0 + pages_with_significant_images: int = 0 + large_image_page_ratio: float = 0.0 + table_signal_pages: int = 0 + table_signal_strength: float = 0.0 + complex_pages: int = 0 + complex_page_ratio: float = 0.0 + max_drawing_count: int = 0 + min_text_density_page: float = 0.0 + text_density_std: float = 0.0 + estimated_fast_benefit: float = 0.0 + estimated_risk_score: float = 0.0 + atlas_candidate: bool = False + page_details: List[dict] = field(default_factory=list) + reasoning: str = "" + + def to_dict(self) -> dict: + data = asdict(self) + data.pop("page_details", None) + data.pop("sample_text", None) + return data + + def summary(self) -> str: + parts = ( + f"[{self.file_type.upper()}] route={self.route}, band={self.decision_band}, " + f"scan={self.scan_type}, category={self.doc_category}, " + f"pages={self.page_count}, text_density={self.avg_text_density:.0f}, " + f"img_coverage={self.avg_image_coverage:.1%}, " + f"risk={self.estimated_risk_score:.2f}, gain={self.estimated_fast_benefit:.2f}" + ) + if self.is_degraded_electronic: + parts += ", degraded=True" + return parts + + +def publish_profile_result(queue, profile: DocProfile) -> None: + gc.collect() + queue.put({"ok": True, "profile": asdict(profile)}) + + +def save_profile_metadata(profile: DocProfile, output_dir: str) -> None: + profile_path = os.path.join(output_dir, "profile.json") + with open(profile_path, "w", encoding="utf-8") as file_obj: + json.dump(profile.to_dict(), file_obj, ensure_ascii=False, indent=2) + logger.debug(f"Profile metadata saved to {profile_path}") diff --git a/apps/worker/app/services/document_parser/doc_profiler.py b/apps/worker/app/services/document_parser/profiling/doc_profile_pdf.py similarity index 83% rename from apps/worker/app/services/document_parser/doc_profiler.py rename to apps/worker/app/services/document_parser/profiling/doc_profile_pdf.py index 8fbefea76..32a13840a 100644 --- a/apps/worker/app/services/document_parser/doc_profiler.py +++ b/apps/worker/app/services/document_parser/profiling/doc_profile_pdf.py @@ -1,105 +1,16 @@ # pyright: reportAttributeAccessIssue=false, reportOperatorIssue=false -""" -Agentic Document Profiler +from __future__ import annotations -Before data enters the pipeline, use lightweight analysis (~50ms) to generate -DocProfile, driving routing decisions and type annotations. - -Usage: - from app.services.document_parser.doc_profiler import profile_document - profile = profile_document("/path/to/file.pdf") -""" - -import gc -import json import math -import os -from dataclasses import asdict, dataclass, field -from typing import Any, List, Literal, Optional +from typing import Any -from app.services.document_parser.pymupdf_subprocess import run_in_child_process, worker +from app.services.document_parser.profiling.doc_profile_model import ( + DocProfile, + publish_profile_result, +) +from app.services.document_parser.formats.pdf.pymupdf_subprocess import run_in_child_process, worker from loguru import logger - -@dataclass -class DocProfile: - """Profile data structure.""" - - # Basic information - file_type: str = "" - - # Routing decision - route: Literal["fast", "standard"] = "standard" - decision_band: Literal["safe_fast", "gray_zone", "safe_standard"] = "safe_standard" - - # Document type - scan_type: Optional[Literal["electronic", "scanned", "mixed"]] = None - doc_category: Literal["generic", "atlas", "ppt_converted"] = "generic" - - # Raw features - page_count: int = 0 - avg_text_density: float = 0.0 - avg_image_coverage: float = 0.0 - has_tables: bool = False - has_embedded_fonts: bool = False - is_multi_column: bool = False - is_degraded_electronic: bool = False - sample_text: str = "" - - # Image complexity - has_significant_images: bool = False - significant_image_count: int = 0 - max_image_coverage_on_page: float = 0.0 - pages_with_significant_images: int = 0 - large_image_page_ratio: float = 0.0 - - # Table complexity - table_signal_pages: int = 0 - table_signal_strength: float = 0.0 - - # Page complexity - complex_pages: int = 0 - complex_page_ratio: float = 0.0 - max_drawing_count: int = 0 - min_text_density_page: float = 0.0 - text_density_std: float = 0.0 - - # Aggregated decision scores - estimated_fast_benefit: float = 0.0 - estimated_risk_score: float = 0.0 - - # Atlas VLM second-pass flag - # True when heuristics suggest atlas-like layout but confidence is not high enough - # to commit without visual confirmation from a VLM. - atlas_candidate: bool = False - - # Page details for debug - page_details: List[dict] = field(default_factory=list) - - # Reasoning - reasoning: str = "" - - def to_dict(self) -> dict: - """Convert to dict (excluding page_details/sample_text to reduce size).""" - data = asdict(self) - data.pop("page_details", None) - data.pop("sample_text", None) - return data - - def summary(self) -> str: - """One-line summary.""" - parts = ( - f"[{self.file_type.upper()}] route={self.route}, band={self.decision_band}, " - f"scan={self.scan_type}, category={self.doc_category}, " - f"pages={self.page_count}, text_density={self.avg_text_density:.0f}, " - f"img_coverage={self.avg_image_coverage:.1%}, " - f"risk={self.estimated_risk_score:.2f}, gain={self.estimated_fast_benefit:.2f}" - ) - if self.is_degraded_electronic: - parts += ", degraded=True" - return parts - - # Thresholds SCAN_TEXT_THRESHOLD = 50 SCAN_IMAGE_COVERAGE_MIN = 0.6 @@ -335,11 +246,6 @@ def _classify_route(profile: DocProfile) -> tuple[str, str, float, float, list[s ) -def _publish_profile_result(queue, profile: DocProfile) -> None: - """Release Python-side wrappers before publishing the profile result.""" - gc.collect() - queue.put({"ok": True, "profile": asdict(profile)}) - @worker def _profile_pdf_worker(queue, file_path: str) -> None: @@ -353,7 +259,7 @@ def _profile_pdf_worker(queue, file_path: str) -> None: doc = pymupdf.open(file_path) except Exception as exc: profile.reasoning = f"Cannot open file: {exc}" - _publish_profile_result(queue, profile) + publish_profile_result(queue, profile) return profile.page_count = doc.page_count @@ -362,7 +268,7 @@ def _profile_pdf_worker(queue, file_path: str) -> None: profile.reasoning = "Empty file (0 pages)" doc.close() del doc - _publish_profile_result(queue, profile) + publish_profile_result(queue, profile) return if doc.page_count <= 50: @@ -530,11 +436,6 @@ def _profile_pdf_worker(queue, file_path: str) -> None: drawing_table_signal = line_like_items >= TABLE_DRAWING_LINE_THRESHOLD and ( (horizontal_line_items >= 2 and vertical_line_items >= 2) or rect_items >= TABLE_DRAWING_RECT_THRESHOLD - or ( - line_like_items >= TABLE_DRAWING_STRONG_THRESHOLD - and horizontal_line_items >= 3 - and vertical_line_items >= 3 - ) ) # NOTE: # `page.find_tables()` produces too many false positives on Word / Writer @@ -753,10 +654,10 @@ def _profile_pdf_worker(queue, file_path: str) -> None: reasons.extend(route_reasons) profile.reasoning = " | ".join(reasons) - _publish_profile_result(queue, profile) + publish_profile_result(queue, profile) -def _profile_pdf(file_path: str) -> DocProfile: +def profile_pdf(file_path: str) -> DocProfile: """Profile a PDF by running PyMuPDF analysis in a spawned child process.""" result = run_in_child_process(_profile_pdf_worker, file_path, timeout=300) profile = DocProfile(**result["profile"]) @@ -768,37 +669,3 @@ def _profile_pdf(file_path: str) -> DocProfile: f"gain={profile.estimated_fast_benefit:.2f}" ) return profile - - -def profile_document(file_path: str, filename: str = "") -> DocProfile: - """ - General document profiling entry point. - - Args: - file_path: Local file path - filename: File name (used to infer type) - - Returns: - DocProfile - """ - if not filename: - filename = os.path.basename(file_path) - - ext = os.path.splitext(filename)[1].lower() - if ext == ".pdf": - return _profile_pdf(file_path) - - return DocProfile( - file_type=ext.lstrip("."), - route="standard", - decision_band="safe_standard", - reasoning=f"Non-PDF format ({ext}), using default route", - ) - - -def save_profile_metadata(profile: DocProfile, output_dir: str): - """Save profile to output_dir/profile.json.""" - profile_path = os.path.join(output_dir, "profile.json") - with open(profile_path, "w", encoding="utf-8") as file_obj: - json.dump(profile.to_dict(), file_obj, ensure_ascii=False, indent=2) - logger.debug(f"Profile metadata saved to {profile_path}") diff --git a/apps/worker/app/services/document_parser/profiling/doc_profiler.py b/apps/worker/app/services/document_parser/profiling/doc_profiler.py new file mode 100644 index 000000000..184f13642 --- /dev/null +++ b/apps/worker/app/services/document_parser/profiling/doc_profiler.py @@ -0,0 +1,41 @@ +""" +Agentic Document Profiler + +Before data enters the pipeline, use lightweight analysis (~50ms) to generate +DocProfile, driving routing decisions and type annotations. + +Usage: + from app.services.document_parser.profiling.doc_profiler import profile_document + profile = profile_document("/path/to/file.pdf") +""" + +import os + +from app.services.document_parser.profiling.doc_profile_model import DocProfile +from app.services.document_parser.profiling.doc_profile_pdf import profile_pdf + + +def profile_document(file_path: str, filename: str = "") -> DocProfile: + """ + General document profiling entry point. + + Args: + file_path: Local file path + filename: File name (used to infer type) + + Returns: + DocProfile + """ + if not filename: + filename = os.path.basename(file_path) + + ext = os.path.splitext(filename)[1].lower() + if ext == ".pdf": + return profile_pdf(file_path) + + return DocProfile( + file_type=ext.lstrip("."), + route="standard", + decision_band="safe_standard", + reasoning=f"Non-PDF format ({ext}), using default route", + ) diff --git a/apps/worker/app/services/document_parser/providers/mineru/client.py b/apps/worker/app/services/document_parser/providers/mineru/client.py new file mode 100644 index 000000000..b86f77531 --- /dev/null +++ b/apps/worker/app/services/document_parser/providers/mineru/client.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +from typing import Any, Optional + +import requests +from app.services.document_parser.providers.mineru.quota_manager import get_mineru_quota_manager +from loguru import logger +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import UnavailableException + + +def build_mineru_session() -> requests.Session: + session = requests.Session() + retry_strategy = Retry( + total=settings.MINERU_UPLOAD_RETRY_TOTAL, + backoff_factor=settings.MINERU_UPLOAD_RETRY_BACKOFF_FACTOR, + status_forcelist=[429, 502, 503, 504], + allowed_methods=["GET", "POST", "PUT"], + raise_on_status=False, + ) + adapter = HTTPAdapter( + max_retries=retry_strategy, + pool_connections=1, + pool_maxsize=settings.MINERU_POOL_MAXSIZE, + ) + session.mount("https://", adapter) + session.mount("http://", adapter) + return session + + +_mineru_session: Optional[requests.Session] = None + + +def get_mineru_session() -> requests.Session: + global _mineru_session + if _mineru_session is None: + _mineru_session = build_mineru_session() + return _mineru_session + + +def get_mineru_headers(api_key: str) -> dict[str, str]: + return { + "Content-Type": "application/json", + "Authorization": f"Bearer {api_key}", + } + + +def mineru_logger(step: str, **fields: Any): + return logger.bind(service="mineru", step=step, **fields) + + +def get_retry_after_seconds( + response: requests.Response, default_retry_after: int +) -> int: + retry_after_header = response.headers.get("Retry-After") + if retry_after_header: + try: + return max( + 1, + min( + int(retry_after_header), settings.MINERU_RATE_LIMIT_MAX_RETRY_AFTER + ), + ) + except ValueError: + logger.debug(f"Invalid MinerU Retry-After header: {retry_after_header}") + + return max(1, min(default_retry_after, settings.MINERU_RATE_LIMIT_MAX_RETRY_AFTER)) + + +def raise_mineru_unavailable( + token_id: str, response: requests.Response, operation: str +) -> None: + retry_after = get_retry_after_seconds( + response, settings.MINERU_TOKEN_COOLDOWN_SECONDS + ) + quota_manager = get_mineru_quota_manager() + quota_manager.mark_rate_limited(token_id, retry_after) + mineru_logger( + "rate_limited", + operation=operation, + token_id=token_id, + status_code=response.status_code, + retry_after=retry_after, + ).warning("MinerU request rate-limited") + raise UnavailableException( + internal_message=f"MinerU rate limited during {operation}", + retry_after=retry_after, + limit=settings.MINERU_TOKEN_RPM_LIMIT, + period="minute", + user_message="Document processing is busy right now. Please retry shortly.", + ) diff --git a/apps/worker/app/services/document_parser/providers/mineru/pdf_service.py b/apps/worker/app/services/document_parser/providers/mineru/pdf_service.py new file mode 100644 index 000000000..4a3b6ec50 --- /dev/null +++ b/apps/worker/app/services/document_parser/providers/mineru/pdf_service.py @@ -0,0 +1,442 @@ +import os +from typing import Optional + +import requests +from app.services.document_parser.providers.mineru.client import ( + get_mineru_headers, + get_mineru_session, + mineru_logger, + raise_mineru_unavailable, +) +from app.services.document_parser.providers.mineru.quota_manager import get_mineru_quota_manager +from app.services.document_parser.providers.mineru.task_polling import ( + get_batch_status, + poll_mineru_task, +) +from app.services.document_parser.support.parser_log_utils import truncate_log_value + +from shared.core.config import settings +from shared.core.constants import APIConstants +from shared.core.exceptions.domain_exceptions import ( + MinerUServiceException, + StorageServiceException, + UnavailableException, +) +from shared.services.storage.job_file_storage import JobFileStorage +from app.services.common.file_loading import is_remote + +MINERU_UPLOAD_TIMEOUT = ( + settings.MINERU_UPLOAD_CONNECT_TIMEOUT, + settings.MINERU_UPLOAD_READ_TIMEOUT, +) + + +def _should_use_mineru_s3_url_mode(s3_key: Optional[str]) -> bool: + if settings.FORCE_MINERU_UPLOAD_ENABLED: + return False + + return settings.ENVIRONMENT != "development" and s3_key is not None + + +def _log_mineru_url_mode_storage_fallback( + operation: str, + s3_key: str, + local_file_path: Optional[str], + exc: Exception, +) -> None: + mineru_logger( + "url_mode_storage_fallback", + operation=operation, + source_s3_key=s3_key, + local_file_path=local_file_path, + error_type=type(exc).__name__, + error_message=truncate_log_value(exc), + ).warning( + "MinerU URL-mode storage preparation failed. Falling back to direct upload." + ) + + +def _log_mineru_url_mode_ingestion_fallback( + operation: str, + s3_key: str, + pdf_url: str, + exc: Exception, +) -> None: + mineru_logger( + "url_mode_ingestion_fallback", + operation=operation, + source_s3_key=s3_key, + source_kind="remote_url" if is_remote(pdf_url) else "local_file", + source_path=None if is_remote(pdf_url) else pdf_url, + error_type=type(exc).__name__, + error_message=truncate_log_value(exc), + ).warning("MinerU URL-mode ingestion setup failed. Falling back to direct upload.") + + +def _inspect_mineru_source_s3_key(s3_key: Optional[str]) -> tuple[Optional[str], bool]: + """Inspect whether URL mode can reuse or prepare the requested S3 source key.""" + if not _should_use_mineru_s3_url_mode(s3_key): + return None, False + + assert s3_key is not None + try: + existing_file = JobFileStorage().verify_upload_exists(s3_key) + except Exception as exc: + _log_mineru_url_mode_storage_fallback( + operation="verify_source_object", + s3_key=s3_key, + local_file_path=None, + exc=exc, + ) + return None, False + + if existing_file.get("exists"): + mineru_logger( + "url_mode_source_reused", + source_s3_key=s3_key, + ).info("Reusing existing S3 source for MinerU URL mode") + return s3_key, True + + return None, True + + +def get_existing_mineru_source_s3_key(s3_key: Optional[str]) -> Optional[str]: + """Return an existing S3 source key for URL mode, or None if it is unavailable.""" + existing_s3_key, _ = _inspect_mineru_source_s3_key(s3_key) + return existing_s3_key + + +def resolve_mineru_source_s3_key( + s3_key: Optional[str], + local_file_path: Optional[str] = None, +) -> Optional[str]: + """Resolve an S3 source key for URL mode, uploading a local file if needed.""" + existing_s3_key, can_prepare_url_mode = _inspect_mineru_source_s3_key(s3_key) + if existing_s3_key is not None: + return existing_s3_key + + if not can_prepare_url_mode: + return None + + if local_file_path is None or is_remote(local_file_path): + return None + + assert s3_key is not None + try: + JobFileStorage().upload_source_file(local_file_path, s3_key) + except Exception as exc: + _log_mineru_url_mode_storage_fallback( + operation="upload_source_object", + s3_key=s3_key, + local_file_path=local_file_path, + exc=exc, + ) + return None + + mineru_logger( + "url_mode_source_uploaded", + source_s3_key=s3_key, + local_file_path=local_file_path, + ).info("Uploaded local PDF to S3 for MinerU URL mode") + return s3_key + + + +def _request_upload_target(pdf_url: str, filename: str) -> tuple[str, str, str]: + base_url = settings.MINERU_URL + quota_manager = get_mineru_quota_manager() + upload_logger = mineru_logger( + "upload_url", + operation="upload_url", + filename=filename, + source_kind="remote_url" if is_remote(pdf_url) else "local_file", + ) + url = f"{base_url}/file-urls/batch" + payload = { + "files": [ + { + "name": filename, + "is_ocr": True, + } + ], + "enable_formula": True, + "enable_table": True, + "language": "auto", + "model_version": "vlm", + } + + upload_logger.info("Requesting MinerU upload URL") + lease = quota_manager.acquire_request(operation="upload_url") + upload_logger.bind(token_id=lease.token_id).info( + "Acquired MinerU token for upload URL" + ) + response = get_mineru_session().post( + url, + headers=get_mineru_headers(lease.api_key), + json=payload, + timeout=settings.MINERU_API_TIMEOUT, + ) + if response.status_code == 429: + raise_mineru_unavailable(lease.token_id, response, operation="upload_url") + if response.status_code != 200: + upload_logger.bind( + token_id=lease.token_id, + status_code=response.status_code, + ).error("Failed to get MinerU upload URL") + raise MinerUServiceException( + internal_message=f"Failed to get upload URL: {response.text}", + status_code=response.status_code, + ) + + result = response.json() + if result.get("code") != 0: + response_message = str(result.get("msg", "Unknown error")) + if "rate limit" in response_message.lower(): + quota_manager.mark_rate_limited( + lease.token_id, + settings.MINERU_TOKEN_COOLDOWN_SECONDS, + ) + upload_logger.bind( + token_id=lease.token_id, + retry_after=settings.MINERU_TOKEN_COOLDOWN_SECONDS, + error_message=response_message, + ).warning("MinerU upload URL request hit rate limit") + raise UnavailableException( + internal_message=f"MinerU rate limited during upload_url: {response_message}", + retry_after=settings.MINERU_TOKEN_COOLDOWN_SECONDS, + limit=lease.rpm_limit, + period="minute", + user_message="Document processing is busy right now. Please retry shortly.", + ) + upload_logger.bind( + token_id=lease.token_id, + error_message=response_message, + ).error("MinerU upload URL request returned API error") + raise MinerUServiceException( + internal_message=f"MinerU API error: {response_message}" + ) + + batch_id = result["data"]["batch_id"] + upload_url = result["data"]["file_urls"][0] + upload_logger.bind(token_id=lease.token_id, batch_id=batch_id).info( + "Received MinerU upload URL" + ) + return batch_id, upload_url, lease.token_id + + +def _upload_file_to_mineru( + pdf_url: str, filename: str, upload_url: str, token_id: str +) -> None: + upload_logger = mineru_logger( + "file_upload", + operation="file_upload", + filename=filename, + token_id=token_id, + source_kind="remote_url" if is_remote(pdf_url) else "local_file", + ) + + if is_remote(pdf_url): + import tempfile + + upload_logger.info("Downloading remote source file before MinerU upload") + try: + download_response = get_mineru_session().get( + pdf_url, + stream=True, + timeout=APIConstants.S3_FILE_DOWNLOAD_TIMEOUT, + ) + download_response.raise_for_status() + + with tempfile.NamedTemporaryFile( + delete=False, suffix=os.path.splitext(filename)[1] + ) as temp_file: + for chunk in download_response.iter_content(chunk_size=8192): + temp_file.write(chunk) + temp_path = temp_file.name + + upload_logger.bind(temp_file_path=temp_path).info( + "Uploading staged file to MinerU" + ) + with open(temp_path, "rb") as file_obj: + upload_response = get_mineru_session().put( + upload_url, + data=file_obj, + timeout=MINERU_UPLOAD_TIMEOUT, + ) + + os.unlink(temp_path) + except requests.RequestException as exc: + upload_logger.bind(error_message=str(exc)).error( + "Failed to stage remote source file for MinerU" + ) + raise StorageServiceException( + internal_message=f"Failed to download remote file: {exc}" + ) + else: + upload_logger.bind(local_path=pdf_url).info("Uploading local file to MinerU") + try: + with open(pdf_url, "rb") as file_obj: + try: + upload_response = get_mineru_session().put( + upload_url, + data=file_obj, + timeout=MINERU_UPLOAD_TIMEOUT, + ) + except requests.RequestException as exc: + upload_logger.bind(error_message=str(exc)).error( + "Failed to upload local file to MinerU" + ) + raise MinerUServiceException( + internal_message=f"Failed to upload file to MinerU: {exc}", + original_exception=exc, + ) from exc + except OSError as exc: + upload_logger.bind(error_message=str(exc)).error( + "Failed to read local file for MinerU upload" + ) + raise StorageServiceException( + internal_message=f"Failed to read local file: {exc}", + original_exception=exc, + ) from exc + + if upload_response.status_code != 200: + upload_logger.bind(status_code=upload_response.status_code).error( + "MinerU file upload failed" + ) + raise MinerUServiceException( + internal_message=f"Failed to upload file to MinerU: {upload_response.text}", + status_code=upload_response.status_code, + ) + + upload_logger.info("MinerU file upload completed, switching to polling") + + +def _submit_url_task(presigned_url: str, filename: str) -> tuple[str, str]: + """Submit a URL-based extraction task to MinerU. + + Uses the /extract/task/batch endpoint so MinerU fetches the file + directly from our S3 via presigned URL, skipping the OSS upload hop. + + Returns (batch_id, token_id). + """ + base_url = settings.MINERU_URL + quota_manager = get_mineru_quota_manager() + submit_logger = mineru_logger( + "submit_url_task", + operation="submit_url_task", + filename=filename, + ) + + url = f"{base_url}/extract/task/batch" + payload = { + "files": [{"url": presigned_url}], + "is_ocr": True, + "enable_formula": True, + "enable_table": True, + "language": "auto", + "model_version": "vlm", + } + + submit_logger.info("Submitting URL-based MinerU extraction task") + lease = quota_manager.acquire_request(operation="submit_url_task") + submit_logger.bind(token_id=lease.token_id).info( + "Acquired MinerU token for URL task submission" + ) + + response = get_mineru_session().post( + url, + headers=get_mineru_headers(lease.api_key), + json=payload, + timeout=settings.MINERU_API_TIMEOUT, + ) + + if response.status_code == 429: + raise_mineru_unavailable(lease.token_id, response, operation="submit_url_task") + + if response.status_code != 200: + submit_logger.bind( + token_id=lease.token_id, + status_code=response.status_code, + ).error("MinerU URL task submission failed") + raise MinerUServiceException( + internal_message=f"URL task submission failed: {response.text}", + status_code=response.status_code, + ) + + result = response.json() + if result.get("code") != 0: + response_message = str(result.get("msg", "Unknown error")) + if "rate limit" in response_message.lower(): + quota_manager.mark_rate_limited( + lease.token_id, + settings.MINERU_TOKEN_COOLDOWN_SECONDS, + ) + raise UnavailableException( + internal_message=f"MinerU rate limited during submit_url_task: {response_message}", + retry_after=settings.MINERU_TOKEN_COOLDOWN_SECONDS, + limit=lease.rpm_limit, + period="minute", + user_message="Document processing is busy right now. Please retry shortly.", + ) + raise MinerUServiceException( + internal_message=f"MinerU API error: {response_message}" + ) + + batch_id = result["data"]["batch_id"] + submit_logger.bind(token_id=lease.token_id, batch_id=batch_id).info( + "MinerU URL task submitted" + ) + return batch_id, lease.token_id + + +def parse_via_full( + pdf_url: str, + filename: str, + output_dir: str, + s3_key: Optional[str] = None, +) -> None: + batch_id: str | None = None + token_id: str | None = None + resolved_s3_key = resolve_mineru_source_s3_key( + s3_key=s3_key, + local_file_path=None if is_remote(pdf_url) else pdf_url, + ) + + if resolved_s3_key is not None: + try: + presigned = JobFileStorage().generate_upload_download_url( + resolved_s3_key, expires_in=settings.MINERU_URL_MODE_PRESIGN_EXPIRY + ) + presigned_url = presigned["download_url"] + mineru_logger("ingestion_mode", mode="s3_url").info( + "Using S3 URL mode for MinerU ingestion" + ) + batch_id, token_id = _submit_url_task(presigned_url, filename) + except Exception as exc: + _log_mineru_url_mode_ingestion_fallback( + operation="start_url_mode_ingestion", + s3_key=resolved_s3_key, + pdf_url=pdf_url, + exc=exc, + ) + resolved_s3_key = None + + if resolved_s3_key is None: + mineru_logger("ingestion_mode", mode="direct_upload").info( + "Using direct upload mode for MinerU ingestion" + ) + batch_id, upload_url, token_id = _request_upload_target(pdf_url, filename) + _upload_file_to_mineru(pdf_url, filename, upload_url, token_id) + + if batch_id is None or token_id is None: + raise MinerUServiceException( + internal_message="MinerU task setup completed without a batch id or token" + ) + + poll_mineru_task( + status_url=f"{settings.MINERU_URL}/extract-results/batch/{batch_id}", + task_id=batch_id, + output_dir=output_dir, + get_status=get_batch_status, + preferred_token_id=token_id, + ) diff --git a/apps/worker/app/services/document_parser/mineru_quota_manager.py b/apps/worker/app/services/document_parser/providers/mineru/quota_manager.py similarity index 95% rename from apps/worker/app/services/document_parser/mineru_quota_manager.py rename to apps/worker/app/services/document_parser/providers/mineru/quota_manager.py index 830330b09..f6fe88f3c 100644 --- a/apps/worker/app/services/document_parser/mineru_quota_manager.py +++ b/apps/worker/app/services/document_parser/providers/mineru/quota_manager.py @@ -14,7 +14,7 @@ SyncRedisService, SyncRedisServiceFactory, ) -from shared.utils.quota_manager import BaseQuotaManager, TokenConfig, TokenLease +from shared.services.quota.token_pool import BaseQuotaManager, TokenConfig, TokenLease # Backward-compatible aliases so existing imports keep working MinerUTokenConfig = TokenConfig diff --git a/apps/worker/app/services/document_parser/providers/mineru/task_polling.py b/apps/worker/app/services/document_parser/providers/mineru/task_polling.py new file mode 100644 index 000000000..fa6ee7775 --- /dev/null +++ b/apps/worker/app/services/document_parser/providers/mineru/task_polling.py @@ -0,0 +1,247 @@ +from __future__ import annotations + +import time +from collections.abc import Callable +from typing import Any, Optional + +import requests +from app.services.document_parser.providers.mineru.client import ( + get_mineru_headers, + get_mineru_session, + mineru_logger, + raise_mineru_unavailable, +) +from app.services.document_parser.providers.mineru.quota_manager import get_mineru_quota_manager +from loguru import logger + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ( + MinerUServiceException, + PDFParsingException, + TimeoutException, + UnavailableException, +) +from shared.core.exceptions.knowhere_exception import KnowhereException +from shared.utils.zip_download import download_and_extract_zip + + +def get_batch_status(data: dict[str, Any]) -> Optional[dict[str, Any]]: + extract_result = data.get("data", {}).get("extract_result") + if isinstance(extract_result, list): + return extract_result[0] if extract_result else None + return extract_result + + +def get_polling_interval_for_state(state: str, attempt: int) -> float: + """Return seconds to sleep before the next MinerU status poll.""" + if state == "pending": + return min(20.0, 5.0 + attempt * 1.5) + if state == "running": + return 10.0 + return 15.0 + + +def poll_mineru_task( + status_url: str, + task_id: str, + output_dir: str, + get_status: Callable[[dict[str, Any]], Optional[dict[str, Any]]], + preferred_token_id: Optional[str] = None, +) -> None: + quota_manager = get_mineru_quota_manager() + polling_logger = mineru_logger( + "poll_status", + operation="poll_status", + task_id=task_id, + preferred_token_id=preferred_token_id, + ) + + max_polling_attempts = 120 + polling_interval = 5.0 + max_wait_time = 6000 + + start_time = time.time() + attempt = 0 + last_token_id: Optional[str] = None + last_state: Optional[str] = None + + polling_logger.info("Starting MinerU polling") + + while attempt < max_polling_attempts: + if time.time() - start_time > max_wait_time: + polling_logger.bind( + attempt=attempt + 1, + max_polling_attempts=max_polling_attempts, + max_wait_time=max_wait_time, + ).warning("MinerU polling timed out") + raise TimeoutException( + internal_message=f"PDF parsing timed out, exceeded {max_wait_time} seconds", + retry_after=60, + user_message="PDF parsing timed out. Please try again.", + ) + + try: + logger.debug( + f"parse_pdfs status_url: {status_url} " + f"(attempt {attempt + 1}/{max_polling_attempts})" + ) + lease = quota_manager.acquire_request( + operation="poll_status", + preferred_token_id=preferred_token_id, + ) + if lease.token_id != last_token_id: + polling_logger.bind( + token_id=lease.token_id, + attempt=attempt + 1, + ).info("Acquired MinerU token for polling") + last_token_id = lease.token_id + + response = get_mineru_session().get( + status_url, + headers=get_mineru_headers(lease.api_key), + timeout=settings.MINERU_API_TIMEOUT, + ) + + if response.status_code == 429: + raise_mineru_unavailable( + lease.token_id, response, operation="poll_status" + ) + + if response.status_code == 200: + response_json = response.json() + if response_json.get("code") != 0: + response_message = str(response_json.get("msg") or "Unknown error") + if "rate limit" in response_message.lower(): + quota_manager.mark_rate_limited( + lease.token_id, + settings.MINERU_TOKEN_COOLDOWN_SECONDS, + ) + raise UnavailableException( + internal_message=( + f"MinerU rate limited during poll_status: {response_message}" + ), + retry_after=settings.MINERU_TOKEN_COOLDOWN_SECONDS, + limit=lease.rpm_limit, + period="minute", + user_message="Document processing is busy right now. Please retry shortly.", + ) + raise MinerUServiceException( + internal_message=f"MinerU API Error: {response_message}" + ) + + status = get_status(response_json) + if not status: + polling_logger.bind( + token_id=lease.token_id, + attempt=attempt + 1, + ).warning("Received empty MinerU status payload") + time.sleep(polling_interval) + attempt += 1 + continue + + state = status.get("state", "unknown") + if state != last_state: + last_state = state + + if state == "done": + download_and_extract_zip( + status["full_zip_url"], + dest_dir=output_dir, + keep_exts=(".md", ".jpg", ".jpeg", ".png", ".gif", ".json"), + exclude_patterns=("content_list", "middle.json", "model.json"), + ) + polling_logger.bind(token_id=lease.token_id).info( + "MinerU parsing completed" + ) + break + + if state == "running": + if "extract_progress" in status: + try: + extracted_pages = status["extract_progress"][ + "extracted_pages" + ] + total_pages = status["extract_progress"]["total_pages"] + progress = extracted_pages / total_pages + polling_logger.bind( + token_id=lease.token_id, + extracted_pages=extracted_pages, + total_pages=total_pages, + progress_pct=round(progress * 100, 1), + ).debug("MinerU parsing in progress") + except (KeyError, ZeroDivisionError): + polling_logger.bind(token_id=lease.token_id).info( + "MinerU parsing in progress" + ) + else: + polling_logger.bind(token_id=lease.token_id).info( + "MinerU parsing in progress" + ) + elif state == "failed": + error_message = status.get("err_msg", "Unknown error") + polling_logger.bind( + token_id=lease.token_id, + error_message=error_message, + ).error("MinerU parsing reported failed state") + raise PDFParsingException( + user_message="Failed to parse the PDF file", + internal_message=f"MinerU failed with state 'failed': {error_message}", + ) + elif state == "pending": + polling_logger.bind(token_id=lease.token_id).debug( + "MinerU parsing pending" + ) + elif state == "waiting-file": + polling_logger.bind(token_id=lease.token_id).debug( + "MinerU waiting for file queueing" + ) + elif state == "converting": + polling_logger.bind(token_id=lease.token_id).debug( + "MinerU converting file" + ) + else: + polling_logger.bind( + token_id=lease.token_id, + state=state, + ).warning("MinerU returned unknown state") + + time.sleep(get_polling_interval_for_state(state, attempt)) + attempt += 1 + else: + polling_logger.bind( + token_id=lease.token_id, + attempt=attempt + 1, + status_code=response.status_code, + ).warning("MinerU status query failed") + time.sleep(polling_interval * 2) + attempt += 1 + + except requests.RequestException as exc: + polling_logger.bind( + attempt=attempt + 1, + error_message=str(exc), + ).warning("MinerU polling network request failed") + time.sleep(polling_interval * 2) + attempt += 1 + except KnowhereException: + raise + except Exception as exc: + polling_logger.bind( + attempt=attempt + 1, + error_message=str(exc), + ).error("Unexpected error during MinerU polling") + raise PDFParsingException( + user_message="An unexpected error occurred while parsing the PDF", + internal_message=str(exc), + original_exception=exc, + ) + + if attempt >= max_polling_attempts: + raise TimeoutException( + internal_message=( + f"minerU PDF parsing timed out after {max_polling_attempts} attempts, " + f"Task ID: {task_id}" + ), + retry_after=60, + user_message="PDF parsing timed out. Please try again.", + ) diff --git a/apps/worker/app/services/document_parser/structure/heading_candidates.py b/apps/worker/app/services/document_parser/structure/heading_candidates.py new file mode 100644 index 000000000..f9d060349 --- /dev/null +++ b/apps/worker/app/services/document_parser/structure/heading_candidates.py @@ -0,0 +1,488 @@ +from __future__ import annotations + +import re +import unicodedata +from collections import defaultdict +from typing import Any + +import pandas as pd +from docx.oxml.ns import qn +from loguru import logger +from pandas import Index + +from app.services.document_parser.support.text_helpers import count_cn_en + +HEADING_COLUMNS = Index(["id", "heading", "level", "reason"]) + + +def get_max_lvl(code_str: str): + match = re.search(r"\[([^]]+)]", code_str) + if not match: + return "Sure" + + nums = [int(item.strip()) for item in match.group(1).split(",")] + max_value = int(max(nums)) + return max_value if max_value > 1 else -2 + + +def judge_by_conditions( + text, + scope: int = 20, + return_detail: bool = False, + cn_special_index: int = 12, + **legacy_options: Any, +): + legacy_cn_special_index = legacy_options.pop("CN_SPECIAL_IDX", None) + if legacy_cn_special_index is not None: + cn_special_index = int(legacy_cn_special_index) + if legacy_options: + unknown_options = ", ".join(sorted(legacy_options)) + raise TypeError(f"Unknown heading condition option(s): {unknown_options}") + + text = text.replace("\u3000", " ") + text = unicodedata.normalize("NFKC", text)[:scope] + + pos_regex_conditions = [ + r"^\d+(?:\s*\.\s*\d+)+(?![、，。！？；：])(?=\s|$|\w|[一-龥])", + r"^\d、\s{0,4}(?=\S|$)", + r"^\d+\.(?!\d)\s{0,4}(?=\S)", + r"^[0-9]{1,2}\s{1,8}(?=\S)", + r"^\d+(?:\.\d+)*、\s*(?=[A-Za-z一-龥])", + r"^[一二三四五六七八九十百千万]+、\s{0,4}(?=\S|$)", + r"^[一二三四五六七八九十百千万]+(?:\s*\.[一二三四五六七八九十百千万\d]+)+", + r"^[一二三四五六七八九十百千万]+(?=\s|$)", + r"^[\(\（]\s*\d+(?:\.\d+)*(?!\.0)\s*[\)\）]", + r"^\d+(?:\.\d+)*(?!\.0)\s*[\)\）]", + r"^[\(\（]\s*[一二三四五六七八九十百千万]+(?:\.[一二三四五六七八九十百千万\d]+)*\s*[\)\）]", + r"^[一二三四五六七八九十百千万]+(?:\.[一二三四五六七八九十百千万\d]+)*\s*[\)\）]", + r"^第[一二三四五六七八九十百千万\d]+(?:\.[一二三四五六七八九十百千万\d]+)*(章|节|条|部分|款|目|项|编|篇|卷|辑)?(?=$|\s|[A-Za-z0-9\u4e00-\u9fa5])", + r"^[A-Za-z](?:\.\d+)*[\.、](?=\s*\S)", + r"^[\(\（]\s*[A-Za-z](?:\.\d+)*(?!\.0)\s*[\)\）]", + r"^[A-Za-z](?:\.\d+)*(?!\.0)\s*[\)\）]", + r"^((附件|附录|附表|附图)|(?i:appendix))[\s_\-—]{0,4}(?:\[)?[一二三四五六七八九十A-Za-z\d]", + ] + + pos_triggered_code = [] + reason_suffix_parts = [] + + for index, regex in enumerate(pos_regex_conditions): + match = re.match(regex, text) + if match: + matched_text = match.group(0) + count = sum(matched_text.count(symbol) for symbol in ".-") + 1 + if index == cn_special_index and return_detail: + unit_match = re.search(r"(章|节|条|部分|款|目|项|编|篇|卷|辑)", matched_text) + if unit_match: + reason_suffix_parts.append(f"[CN:{unit_match.group(1)}]") + pos_triggered_code.append(count) + else: + pos_triggered_code.append(0) + + if return_detail: + reason_suffix = " ".join(reason_suffix_parts) if reason_suffix_parts else "" + return pos_triggered_code, { + "reason_suffix": f" {reason_suffix}" if reason_suffix else "" + } + return pos_triggered_code + + +def remove_by_conditions(text, include_punc: bool = False): + neg_conditions = [ + r"^\d{3,}", + r"(?i)(^https?://\S+|^www\.\S+|^P\.S|^\b\d{0,2}\s*(?:a\.m|p\.m)\b)", + ( + r"(?:" + r"\$[^$]*\\[A-Za-z]+(?:\s*\{[^{}]*\})?[^$]*\$" + r"|" + r"\\(?:times|div|cdot|pm|mp|leq|geq|neq|approx|equiv|sim|infty" + r"|sum|prod|int|sqrt|frac|mathrm|mathbf|mathit|mathcal" + r"|text(?:bf|it|rm)?|alpha|beta|gamma|delta|epsilon|theta" + r"|lambda|mu|sigma|pi|omega|partial|nabla" + r"|left|right|begin|end|overline|underline|hat|vec|tilde)\b" + r")" + ), + r"^0\.\d+[\u4e00-\u9fa5A-Za-z\S]*", + r"^\d*\.\d+$", + r"[。！；].+", + ( + r"^\d+\.?\d*\s{0,2}" + r"(?:mm|cm|km|nm|μm|inch(?:es)?|ft|yd|mi" + r"|kg|mg|μg|lb|oz" + r"|kPa|MPa|GPa|Pa|psi|bar" + r"|°[CFK]" + r"|Hz|kHz|MHz|GHz" + r"|mol|mL|dL|dB|Nm|kN|MN|kW|MW|GW|hp|rpm|cc|cal|kcal)\b" + ), + ] + + neg_triggered_code = [] + for regex in neg_conditions: + neg_triggered_code.append(1 if re.search(regex, text) else 0) + + if include_punc: + neg_triggered_code.append(1 if re.search(r"[.,;，。；]$", text) else 0) + else: + neg_triggered_code.append(0) + + return neg_triggered_code + + +def md_heading_match(line, as_is: bool = True): + match = re.match(r"^\s*(#+)\s*(.*)$", line) + if not match: + return line, -1 + + level = len(match.group(1)) + return (line, level) if as_is else (line.lstrip("#").strip(), level) + + +def filter_markdown_headings( + md_lines: list[str], + num_pos: int = 17, + num_neg: int = 7, + layout_json_path: str | None = None, +) -> pd.DataFrame: + meta_ctx = None + if layout_json_path: + try: + from app.services.document_parser.structure.metadata_extractor import MetadataContext + + meta_ctx = MetadataContext(md_lines, layout_json_path) + except Exception as exc: + logger.warning(f"Failed to create MetadataContext: {exc}") + + raw_candidates = [] + for line_index, line in enumerate(md_lines): + line = line.strip() + if not line: + continue + + if _is_non_heading_markdown_line(line): + est_level = -1 + zero_pos_code = [0] * num_pos + zero_neg_code = [0] * num_neg + reason = f"POS {zero_pos_code} NEG {zero_neg_code}" + if meta_ctx: + reason += " META [0, 0, 0]" + line = "Figure/Image" + else: + est_level, reason, line = _estimate_markdown_heading_level(line, meta_ctx) + + raw_candidates.append((line_index, line, est_level, reason)) + + return pd.DataFrame(raw_candidates, columns=HEADING_COLUMNS, index=None) + + +def filter_document_headings( + heading_infos: list[tuple[Any, Any, str]], + *, + enable_regex: bool = True, +) -> pd.DataFrame: + raw_candidates = [] + logger.debug("Filtering docx heading candidates... total_items={}", len(heading_infos)) + + for element_id, paragraph, text in heading_infos: + reason = "" + est_level = None + style_level = _find_docx_style_level(paragraph) + setting_level = _find_docx_outline_level(paragraph) + + if style_level is not None: + est_level = style_level + reason = f"style-{style_level}" + elif setting_level is not None: + est_level = setting_level + reason = f"outline-{setting_level}" + + is_bold = 1 if _is_bold_docx_paragraph(paragraph) else 0 + + if enable_regex: + pos_code, detail_info = judge_by_conditions(text, return_detail=True) + neg_code = remove_by_conditions(text) + + if any(value > 0 for value in neg_code): + code_level = -1 + code_reason = ( + f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}" + ) + elif any(value > 0 for value in pos_code) and all( + value == 0 for value in neg_code + ): + code_level = get_max_lvl(str(pos_code)) + code_reason = ( + f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}" + ) + else: + code_level = -1 + code_reason = f"POS {pos_code} NEG {neg_code}" + + if is_bold: + code_reason += f" META [0, 0, {is_bold}]" + + if est_level is None: + est_level = code_level + reason = code_reason + else: + reason = f"{reason} AND {code_reason}" + + raw_candidates.append((element_id, text, est_level, reason)) + + candidates = pd.DataFrame(raw_candidates, columns=HEADING_COLUMNS, index=None) + if candidates.empty: + return pd.DataFrame(columns=HEADING_COLUMNS) + + candidates = postprocess_headings(candidates, task="merge_continuous") + return postprocess_headings(candidates, task="merge_short") + + +def postprocess_headings(df: pd.DataFrame, task: str, max_depth: int = -1) -> pd.DataFrame: + if task == "judge_negs": + return _judge_negative_headings(df) + + if task == "merge_continuous": + return _merge_continuous_non_headings(df) + + if task == "merge_short" or task == "collapse": + return _collapse_heading_groups(df, task) + + return df + + +def _is_non_heading_markdown_line(line: str) -> bool: + return ( + ("" in line) + or line.startswith("|") + or line.startswith("
") + or ("![" in line and "](" in line) + ) + + +def _estimate_markdown_heading_level(line: str, meta_ctx: Any | None): + from app.services.document_parser.structure.metadata_extractor import detect_and_strip_md_bold + + line_clean, hash_level = md_heading_match(line, as_is=False) + stripped_line, is_full_bold = detect_and_strip_md_bold(line_clean) + pos_code, detail_info = judge_by_conditions(stripped_line, return_detail=True) + neg_code = remove_by_conditions(stripped_line) + + if any(value > 0 for value in neg_code): + code_level = -1 + code_reason = f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}" + elif any(value > 0 for value in pos_code) and all(value == 0 for value in neg_code): + code_level = get_max_lvl(str(pos_code)) + code_reason = f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}" + else: + code_level = -1 + code_reason = f"POS {pos_code} NEG {neg_code}" + + if meta_ctx: + size_rank, occurrence = meta_ctx.get_meta_for_line(line_clean) + bold_value = 1 if is_full_bold else 0 + code_reason += meta_ctx.format_meta_suffix(size_rank, occurrence, bold_value) + elif is_full_bold: + code_reason += " META [0, 0, 1]" + + if hash_level <= 0: + return code_level, code_reason, line_clean + + if isinstance(code_level, int): + est_level = max(hash_level, code_level) + else: + est_level = code_level + return est_level, f"{hash_level}# AND {code_reason}", line_clean + + +def _find_docx_style_level(paragraph: Any): + try: + style_name = paragraph.style.name + except Exception: + style_name = "normal" + + if not (style_name.startswith("Heading") or style_name.startswith("标题")): + return None + + try: + return int(style_name.split(" ")[1]) + except Exception: + return -2 + + +def _find_docx_outline_level(paragraph: Any): + paragraph_properties = paragraph._element.find(qn("w:pPr")) + if paragraph_properties is None: + return None + + outline_level = paragraph_properties.find(qn("w:outlineLvl")) + if outline_level is None: + return None + + return int(outline_level.get(qn("w:val"))) + 1 + + +def _is_bold_docx_paragraph(paragraph: Any): + if paragraph.runs and all(run.bold for run in paragraph.runs if run.text.strip()): + return True + return None + + +def _judge_negative_headings(df: pd.DataFrame) -> pd.DataFrame: + for index, row in df.iterrows(): + neg_code = remove_by_conditions(row["heading"], include_punc=True) + if any(value > 0 for value in neg_code): + current_code = str(df.loc[index, "reason"]) + + neg_match = re.search(r"(.*NEG\s*)\[[^\]]*\](.*)", current_code) + if neg_match: + updated_code = f"{neg_match.group(1)}{neg_code}{neg_match.group(2)}" + else: + updated_code = f"{current_code} NEG {neg_code}" + + df.loc[index, "level"] = -1 + df.loc[index, "reason"] = updated_code + return df + + +def _merge_continuous_non_headings(df: pd.DataFrame) -> pd.DataFrame: + denoised_rows = [] + punc_pattern = re.compile(r'[.,!?;:，。！？；：）】〕｝〉》’”"]$') + + index = 0 + while index < len(df): + row = df.iloc[index] + current_content = str(row["heading"]).strip() + current_level = row["level"] + + next_index = index + 1 + while next_index < len(df): + next_row = df.iloc[next_index] + next_content = str(next_row["heading"]).strip() + next_level = next_row["level"] + + expected_id = row["id"] + (next_index - index) + if next_row["id"] != expected_id: + break + + current_not_punc = not punc_pattern.search(current_content[-2:]) + if (current_level == -1 and next_level == -1) and current_not_punc: + current_content += " " + next_content + next_index += 1 + else: + break + + merged_row = row.copy() + merged_row["heading"] = current_content + denoised_rows.append(tuple(merged_row)) + index = next_index + + return pd.DataFrame(denoised_rows, columns=HEADING_COLUMNS) + + +def _collapse_heading_groups(df: pd.DataFrame, task: str) -> pd.DataFrame: + group_to_indices = defaultdict(list) + for index, row in df.iterrows(): + level = row["level"] + reason = row["reason"] + if level != -1: + group_to_indices[(level, reason)].append(index) + + checked_pairs = set() + for _, indices in group_to_indices.items(): + _collapse_recursive(df, task, indices, merge_threshold=3, checked_pairs=checked_pairs) + + if task == "merge_short": + drop_between = df.index[ + df["reason"].astype(str).str.startswith("Merged into", na=False) + ].tolist() + if drop_between: + logger.debug( + f"🛠️ Delete rows labeled as merged into, total {len(drop_between)} rows" + ) + df.drop(drop_between, inplace=True) + df.reset_index(drop=True, inplace=True) + return df + + +def _collapse_recursive( + df: pd.DataFrame, + task: str, + indices: list[int], + merge_threshold: int = 3, + checked_pairs: set[tuple[int, int]] | None = None, +) -> None: + if checked_pairs is None: + checked_pairs = set() + + if len(indices) < 2: + return + + for group_index in range(len(indices) - 1): + index, next_index = indices[group_index], indices[group_index + 1] + if (index, next_index) in checked_pairs: + continue + checked_pairs.add((index, next_index)) + + between = df.loc[index + 1 : next_index - 1] + current_text = df.at[index, "heading"].strip() + next_text = df.at[next_index, "heading"].strip() + + if task == "merge_short" and len(between) > 0: + _merge_short_between_headings( + df, between, index, next_index, current_text, merge_threshold + ) + elif task == "collapse" and len(between) == 0: + logger.debug( + f"⚠️ Empty between i={current_text[:15]}, j={next_text[:15]} => set i.level=-1, j.level=Not Sure" + ) + df.at[index, "level"] = -2 + df.at[next_index, "level"] = -2 + + sub_between = between[between["level"] != -1] + code_to_sub = defaultdict(list) + for row_index, row in sub_between.iterrows(): + level = row["level"] + reason = row["reason"] + if level != -1: + code_to_sub[(level, reason)].append(row_index) + + for _, sub_indices in code_to_sub.items(): + _collapse_recursive(df, task, sub_indices, merge_threshold, checked_pairs) + + +def _merge_short_between_headings( + df: pd.DataFrame, + between: pd.DataFrame, + index: int, + next_index: int, + current_text: str, + merge_threshold: int, +) -> None: + between_lengths = [count_cn_en(content) for content in between["heading"].tolist()] + between_levels = [level for level in between["level"].tolist()] + half_current_length = int(count_cn_en(current_text) / 2) + too_short = ( + sum(between_lengths) <= merge_threshold + or sum(between_lengths) < half_current_length + ) + + if not too_short or not all(level == -1 for level in between_levels): + return + + next_text = df.at[next_index, "heading"].strip() + logger.debug( + f"⚠️ too short between {index}=>{current_text[:15]} and {next_index}=>{next_text[:15]} => merge to {index}" + ) + between_texts = [ + heading_text.strip() + for _, row in between.iterrows() + if isinstance(heading_text := row["heading"], str) and heading_text.strip() + ] + + joined_text = "" + if between_texts: + joined_text = "\n".join(between_texts) + df.at[index, "heading"] = f"{current_text} {joined_text}" + + for row_index in between.index: + df.at[row_index, "level"] = -1 + df.at[row_index, "reason"] = f"Merged into {index}" + logger.debug(f"\tmerged texts: {joined_text[:50]}...") diff --git a/apps/worker/app/services/document_parser/structure/heading_hierarchy.py b/apps/worker/app/services/document_parser/structure/heading_hierarchy.py new file mode 100644 index 000000000..20a9b170e --- /dev/null +++ b/apps/worker/app/services/document_parser/structure/heading_hierarchy.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Literal + +import pandas as pd + +from app.services.document_parser.structure.layout_parser import pred_titles + + +@dataclass(frozen=True) +class HeadingHierarchyInput: + infos: Any + doc_type: Literal["pptx", "md", "docx"] + toc_hierarchies: Any | None = None + prompt_limit: int = 4000 + enable_regex: bool = True + smart_parse: bool = False + model_name: str | None = None + output_dir: str | None = None + layout_json_path: str | None = None + first_toc_ele_num: int | None = None + + +def predict_heading_hierarchy(heading_input: HeadingHierarchyInput) -> pd.DataFrame: + return pred_titles( + heading_input.infos, + doc_type=heading_input.doc_type, + toc_hierarchies=heading_input.toc_hierarchies, + prompt_limt=heading_input.prompt_limit, + enable_regx=heading_input.enable_regex, + smart_parse=heading_input.smart_parse, + model_name=heading_input.model_name, + output_dir=heading_input.output_dir, + layout_json_path=heading_input.layout_json_path, + first_toc_ele_num=heading_input.first_toc_ele_num, + ) diff --git a/apps/worker/app/services/document_parser/structure/heading_llm_executor.py b/apps/worker/app/services/document_parser/structure/heading_llm_executor.py new file mode 100644 index 000000000..6f314416b --- /dev/null +++ b/apps/worker/app/services/document_parser/structure/heading_llm_executor.py @@ -0,0 +1,610 @@ +# pyright: reportArgumentType=false, reportAssignmentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportOptionalSubscript=false, reportReturnType=false +from __future__ import annotations + +import os +import re +from collections import Counter +from collections.abc import Callable +from typing import Any + +import pandas as pd +from app.services.document_parser.structure.metadata_extractor import clean_md_text_for_llm +from app.services.document_parser.support.stage_profiler import stage_timer +from app.services.document_parser.support.text_helpers import count_cn_en, truncate_text_by_tokens +from loguru import logger + +from shared.core.exceptions.domain_exceptions import WorkerHandlingException + +PLACEHOLDER_REASON = "__PLACEHOLDER__" + +HierarchyJudge = Callable[..., list[dict[str, Any]]] +FallbackHierarchy = Callable[[pd.DataFrame], pd.DataFrame] +SaveIntermediateCsv = Callable[[pd.DataFrame, str | None, str], None] + + +def build_level_mapping( + df: pd.DataFrame, origin_lvls: list[int], mode: str = "max" +) -> tuple[pd.DataFrame, dict[str, dict[str, Any]]]: + mapped_df = df.copy() + mapped_df["origin_level"] = origin_lvls + + mapping = mapped_df.groupby("reason")["level"].apply(list).to_dict() + + processed_mapping: dict[str, dict[str, Any]] = {} + for reason, lvls in mapping.items(): + positive_lvls = [lvl for lvl in lvls if lvl > -1] + counts = Counter(lvls) + + if not positive_lvls: + mapped_lvl = -1 + elif mode == "max": + mapped_lvl = max(positive_lvls) + elif mode == "freq": + mapped_lvl = counts.most_common(1)[0][0] + else: + raise WorkerHandlingException( + internal_message=f"wrong input mode: {mode}. Must be 'max' or 'freq'" + ) + + processed_mapping[reason] = { + "lvls": lvls, + "positive_lvls": positive_lvls, + "freqs": dict(counts), + "mapped_lvl": mapped_lvl, + } + return mapped_df, processed_mapping + + +def execute_level_mapping( + df: pd.DataFrame, mapping: dict[str, dict[str, Any]] +) -> pd.DataFrame: + def map_row(row: pd.Series) -> int: + reason = row["reason"] + if reason in mapping: + return int(mapping[reason]["mapped_lvl"]) + return int(row["level"]) + + mapped_df = df.copy() + origin_est_lvls = mapped_df["level"].tolist() + mapped_df["level"] = mapped_df.apply(map_row, axis=1) + mapped_df["origin_level"] = origin_est_lvls + return mapped_df + + +def extract_non_neg_code(reason_str: str) -> str: + """Extract the non-NEG code from a heading reason string.""" + if not reason_str or not isinstance(reason_str, str): + return "" + neg_match = re.search(r"\s*NEG\s*\[[^\]]*\]", reason_str) + if neg_match: + before_neg = reason_str[: neg_match.start()] + after_neg = reason_str[neg_match.end() :] + return (before_neg + after_neg).strip() + return reason_str.strip() + + +def build_non_neg_mapping(lvl_mapping: dict[str, dict[str, Any]]) -> dict[str, int]: + non_neg_levels: dict[str, list[int]] = {} + for reason, info in lvl_mapping.items(): + non_neg_code = extract_non_neg_code(reason) + mapped_lvl = int(info.get("mapped_lvl", -1)) + if non_neg_code: + non_neg_levels.setdefault(non_neg_code, []).append(mapped_lvl) + + non_neg_mapping: dict[str, int] = {} + for non_neg_code, levels in non_neg_levels.items(): + positive_levels = [lvl for lvl in levels if lvl > -1] + if positive_levels: + level_counts = Counter(positive_levels) + non_neg_mapping[non_neg_code] = level_counts.most_common(1)[0][0] + else: + non_neg_mapping[non_neg_code] = -1 + + return non_neg_mapping + + +def handle_unseen_codes( + df: pd.DataFrame, + level_dfs: list[pd.DataFrame], + lvl_mapping: dict[str, dict[str, Any]], + output_dir: str | None = None, + window_half_size: int = 10, + strategy: str = "double_mapping", +) -> dict[str, dict[str, Any]]: + """Extend first-chunk reason mapping to reason codes only seen in later chunks.""" + + def extract_reason_signature(reason: str) -> str: + return reason.strip() if reason else "" + + def has_neg_signal(reason_str: str) -> bool: + if not reason_str or not isinstance(reason_str, str): + return False + neg_match = re.search(r"NEG\s*\[([^\]]*)\]", reason_str) + if not neg_match: + return False + neg_content = neg_match.group(1) + try: + nums = [int(x.strip()) for x in neg_content.split(",") if x.strip()] + return any(x >= 1 for x in nums) + except Exception: + return False + + def build_context_window( + target_idx: int, known_codes_set: set[str], total_rows: int, half_size: int = 10 + ) -> dict[str, Any]: + min_start = max(0, target_idx - half_size) + min_end = min(total_rows - 1, target_idx + half_size) + + start_idx = min_start + end_idx = min_end + + found_known_above = False + found_known_below = False + known_positions: list[int] = [] + + for index in range(start_idx, target_idx): + reason = df.iloc[index].get("reason", "") + sig = extract_reason_signature(reason) + if sig in known_codes_set: + found_known_above = True + known_positions.append(index) + + for index in range(target_idx + 1, end_idx + 1): + reason = df.iloc[index].get("reason", "") + sig = extract_reason_signature(reason) + if sig in known_codes_set: + found_known_below = True + known_positions.append(index) + + if not found_known_above and min_start > 0: + search_idx = min_start - 1 + while search_idx >= 0: + reason = df.iloc[search_idx].get("reason", "") + sig = extract_reason_signature(reason) + if sig in known_codes_set: + found_known_above = True + known_positions.append(search_idx) + start_idx = search_idx + break + search_idx -= 1 + + if not found_known_below and min_end < total_rows - 1: + search_idx = min_end + 1 + while search_idx < total_rows: + reason = df.iloc[search_idx].get("reason", "") + sig = extract_reason_signature(reason) + if sig in known_codes_set: + found_known_below = True + known_positions.append(search_idx) + end_idx = search_idx + break + search_idx += 1 + + return { + "start": start_idx, + "end": end_idx, + "found_known": found_known_above or found_known_below, + "known_positions": known_positions, + } + + non_neg_mapping = build_non_neg_mapping(lvl_mapping) + known_codes = set(lvl_mapping.keys()) + + all_codes_in_full: dict[str, dict[str, Any]] = {} + for seg_idx, seg_df in enumerate(level_dfs): + for _, row in seg_df.iterrows(): + reason = row.get("reason", "") + sig = extract_reason_signature(reason) + if not sig or sig == PLACEHOLDER_REASON: + continue + if sig not in all_codes_in_full: + all_codes_in_full[sig] = { + "first_seg": seg_idx, + "first_id": row.get("id", 0), + "reason": reason, + } + + unseen_codes: dict[str, dict[str, Any]] = {} + unseen_neg_filtered: dict[str, dict[str, Any]] = {} + for sig, info in all_codes_in_full.items(): + if sig in known_codes: + continue + if has_neg_signal(info["reason"]): + unseen_neg_filtered[sig] = info + else: + unseen_codes[sig] = info + + logger.info( + f"Unseen codes total: {len(unseen_codes) + len(unseen_neg_filtered)}, " + f"NEG filtered: {len(unseen_neg_filtered)}, to process: {len(unseen_codes)}" + ) + + for sig in unseen_neg_filtered: + lvl_mapping[sig] = {"mapped_lvl": -1, "note": "NEG_FILTERED"} + + if unseen_codes: + if strategy == "double_mapping": + fallback_success = 0 + fallback_failed = 0 + failed_codes = [] + for sig in unseen_codes: + non_neg_code = extract_non_neg_code(sig) + if non_neg_code in non_neg_mapping: + mapped_level = non_neg_mapping[non_neg_code] + lvl_mapping[sig] = { + "mapped_lvl": mapped_level, + "note": f"NON_NEG_FALLBACK from '{non_neg_code}'", + } + fallback_success += 1 + else: + lvl_mapping[sig] = {"mapped_lvl": -1, "note": "NO_MATCH_FALLBACK"} + fallback_failed += 1 + failed_codes.append( + f"'{non_neg_code}' (from '{sig[:60]}...')" + if len(sig) > 60 + else f"'{non_neg_code}' (from '{sig}')" + ) + + logger.info( + f"Double mapping result: success={fallback_success}, failed={fallback_failed}" + ) + if failed_codes: + logger.warning( + f"Failed codes (non_neg not in mapping): {failed_codes[:5]}" + f"{'...' if len(failed_codes) > 5 else ''}" + ) + + elif strategy == "window_llm" and output_dir: + total_rows = len(df) + windows: list[dict[str, Any]] = [] + for sig, info in unseen_codes.items(): + first_id = info["first_id"] + first_seg = info["first_seg"] + df_indices = df.index[df["id"] == first_id].tolist() + if df_indices: + first_df_idx = df_indices[0] + window_info = build_context_window( + first_df_idx, known_codes, total_rows, window_half_size + ) + windows.append( + { + "code": sig, + "first_id": first_id, + "first_seg": first_seg, + "start": window_info["start"], + "end": window_info["end"], + "found_known": window_info["found_known"], + } + ) + + sorted_windows = sorted(windows, key=lambda window: window["start"]) + merged_windows: list[dict[str, Any]] = [] + current_window: dict[str, Any] | None = None + + for window in sorted_windows: + if current_window is None: + current_window = { + "start": window["start"], + "end": window["end"], + "codes": [window["code"]], + "segments": [window["first_seg"]], + } + elif window["start"] <= current_window["end"]: + current_window["end"] = max(current_window["end"], window["end"]) + current_window["codes"].append(window["code"]) + current_window["segments"].append(window["first_seg"]) + else: + merged_windows.append(current_window) + current_window = { + "start": window["start"], + "end": window["end"], + "codes": [window["code"]], + "segments": [window["first_seg"]], + } + + if current_window: + merged_windows.append(current_window) + + windows_dir = os.path.join(output_dir, "merged_windows") + os.makedirs(windows_dir, exist_ok=True) + + unseen_codes_set = set(unseen_codes.keys()) + unseen_neg_set = set(unseen_neg_filtered.keys()) + + for index, merged_window in enumerate(merged_windows): + window_df = df.iloc[ + merged_window["start"] : merged_window["end"] + 1 + ].copy() + + def get_code_status(row: pd.Series) -> str: + reason = row.get("reason", "") + sig = extract_reason_signature(reason) + if not sig: + return "" + if sig in unseen_codes_set: + return "UNSEEN_TARGET" + if sig in unseen_neg_set: + return "NEG_TO_NEGATIVE_ONE" + if sig in known_codes: + return "KNOWN" + return "" + + window_df["code_status"] = window_df.apply(get_code_status, axis=1) + window_path = os.path.join( + windows_dir, + f"window_{index + 1:02d}_rows_" + f"{merged_window['start']}-{merged_window['end']}.csv", + ) + window_df.to_csv(window_path, index=False, encoding="utf-8-sig") + + logger.debug( + f"Window LLM: {len(merged_windows)} windows created in {windows_dir}" + ) + + return lvl_mapping + + +def compact_for_llm(df: pd.DataFrame) -> pd.DataFrame: + """Collapse consecutive body rows into placeholder rows before LLM chunking.""" + if df is None or len(df) == 0: + return pd.DataFrame(columns=["id", "heading", "level", "reason"]) + + rows: list[dict[str, Any]] = [] + index = 0 + row_count = len(df) + while index < row_count: + lvl_raw = df.iloc[index]["level"] + try: + lvl_int = int(lvl_raw) + except (TypeError, ValueError): + lvl_int = None + + if lvl_int == -1: + end_index = index + while end_index < row_count: + try: + next_level = int(df.iloc[end_index]["level"]) + except (TypeError, ValueError): + break + if next_level != -1: + break + end_index += 1 + start_id = int(df.iloc[index]["id"]) + end_id = int(df.iloc[end_index - 1]["id"]) + run_length = end_index - index + rows.append( + { + "id": f"{start_id}-{end_id}", + "heading": f"[{run_length} BODY LINES]", + "level": "-", + "reason": PLACEHOLDER_REASON, + } + ) + index = end_index + else: + row = df.iloc[index] + rows.append( + { + "id": int(row["id"]), + "heading": str(row["heading"]), + "level": ( + int(lvl_int) + if lvl_int is not None and lvl_int != -2 + else "Not Sure" + ), + "reason": str(row.get("reason", "") or ""), + } + ) + index += 1 + + return pd.DataFrame(rows, columns=["id", "heading", "level", "reason"]) + + +def split_heading_table( + df: pd.DataFrame, threshold: int = 3000, max_start: int = 50, max_end: int = 10 +) -> tuple[list[pd.DataFrame], list[str]]: + raw_headings = df["heading"].tolist() + working_df = df.copy() + working_df["heading"] = working_df["heading"].apply( + lambda heading: truncate_text_by_tokens(heading, max_start, max_end) + ) + + sub_dfs: list[pd.DataFrame] = [] + current_rows: list[list[Any]] = [] + current_len = 0 + for _, row in working_df.iterrows(): + row_filtered = row.drop(labels=["reason"], errors="ignore") + row_len = sum(count_cn_en(str(value)) for value in row_filtered.values) + + if current_len + row_len > threshold and current_rows: + sub_dfs.append(pd.DataFrame(current_rows, columns=working_df.columns)) + current_rows = [row.tolist()] + current_len = row_len + else: + current_rows.append(row.tolist()) + current_len += row_len + + if current_rows: + sub_dfs.append(pd.DataFrame(current_rows, columns=working_df.columns)) + return sub_dfs, raw_headings + + +def execute_llm_heading_hierarchy( + raw_preds: pd.DataFrame, + prompt_limt: int, + hierarchy_judge: HierarchyJudge, + fallback_hierarchy: FallbackHierarchy, + save_intermediate_csv: SaveIntermediateCsv, + toc_hierarchies: Any | None = None, + max_len: int = 30, + max_depth: int = 6, + model_name: str | None = None, + output_dir: str | None = None, + csv_suffix: str = "", +) -> pd.DataFrame: + if len(raw_preds) == 0: + return pd.DataFrame(columns=["id", "heading", "level", "reason"]) + + compact_enabled = os.environ.get( + "KB_LAYOUT_LLM_COMPACT_INPUT", "true" + ).strip().lower() in ("true", "1", "yes", "on") + preds_for_llm = compact_for_llm(raw_preds) if compact_enabled else raw_preds.copy() + if compact_enabled: + placeholder_count = int(preds_for_llm["reason"].eq(PLACEHOLDER_REASON).sum()) + logger.info( + f"smart parse => compact input: {len(raw_preds)} -> {len(preds_for_llm)} rows " + f"({placeholder_count} placeholder groups)" + ) + + non_placeholder = ( + preds_for_llm[preds_for_llm["reason"].astype(str) != PLACEHOLDER_REASON] + if compact_enabled + else preds_for_llm + ) + if len(non_placeholder) == 0: + logger.info( + "smart parse => no heading candidates, skipping LLM hierarchy detection" + ) + fallback = raw_preds.copy()[["id", "heading", "level", "reason"]] + fallback["level"] = -1 + return fallback.sort_values("id").reset_index(drop=True) + + level_dfs, _raw_headings = split_heading_table( + preds_for_llm, threshold=prompt_limt, max_start=max_len, max_end=5 + ) + chunk_sizes = [len(dataframe) for dataframe in level_dfs] + logger.info( + f"smart parse => {len(level_dfs)} chunk(s) | rows per chunk: {chunk_sizes} | " + f"threshold={prompt_limt} | max_start={max_len}" + ) + + basic_idx = 0 + for idx, chunk in enumerate(level_dfs): + if (chunk["reason"].astype(str) != PLACEHOLDER_REASON).any(): + basic_idx = idx + break + basic_df = level_dfs[basic_idx] + if basic_idx != 0: + logger.info( + f"smart parse => promoted chunk {basic_idx} as basic_df " + f"(chunks 0..{basic_idx - 1} contain only placeholders)" + ) + + full_preds: pd.DataFrame | None = None + try: + with stage_timer( + "heading.hierarchy_llm", + chunk_count=len(level_dfs), + base_chunk_rows=len(basic_df), + compact_enabled=compact_enabled, + source_row_count=len(raw_preds), + model_name=model_name, + ): + logger.debug("smart parse => interpreting hierarchy patterns...") + df4llm = basic_df.drop(columns=["reason"]).copy() + df4llm["heading"] = df4llm["heading"].apply(clean_md_text_for_llm) + logger.debug(f"DataFrame transformation completed, rows: {len(df4llm)}") + + layout_res = hierarchy_judge( + df4llm, model_name, max_depth, toc_hierarchies, task="eval-headings" + ) + + layout_level_by_id: dict[Any, Any] = {} + if isinstance(layout_res, list): + for item in layout_res: + if isinstance(item, dict) and "id" in item and "level" in item: + layout_level_by_id[item["id"]] = item["level"] + + def level_for(row_id: Any) -> Any: + if row_id in layout_level_by_id: + return layout_level_by_id[row_id] + try: + return layout_level_by_id.get(int(row_id), -1) + except (TypeError, ValueError): + return -1 + + base_preds = ( + basic_df[["id", "heading", "reason"]].copy().reset_index(drop=True) + ) + base_preds.insert(2, "level", base_preds["id"].map(level_for)) + save_intermediate_csv( + base_preds, output_dir, f"preds_3_llm_base{csv_suffix}" + ) + + llm_levels: dict[int, Any] = {} + for _, row in base_preds.iterrows(): + row_id = row["id"] + if isinstance(row_id, bool): + continue + if isinstance(row_id, int): + llm_levels[row_id] = row["level"] + + if len(level_dfs) > 1: + placeholder_mask_base = base_preds["reason"].eq(PLACEHOLDER_REASON) + figure_mask_base = base_preds["heading"].eq("Figure/Image") + exclude_mask_base = placeholder_mask_base | figure_mask_base + base_preds_for_mapping = base_preds[~exclude_mask_base].copy() + base_origin_for_mapping = basic_df.loc[ + ~exclude_mask_base.values, "level" + ].tolist() + + base_preds_for_mapping, lvl_mapping = build_level_mapping( + base_preds_for_mapping, base_origin_for_mapping, mode="freq" + ) + logger.debug( + f"mapping development finished: {len(lvl_mapping)} rules " + f"(placeholders and Figure/Image excluded)" + ) + + logger.debug( + f"mapping dataframe to levels across {len(level_dfs)} chunks..." + ) + lvl_mapping = handle_unseen_codes( + preds_for_llm, level_dfs, lvl_mapping, output_dir + ) + + for level_df in level_dfs: + placeholder_mask_chunk = level_df["reason"].eq(PLACEHOLDER_REASON) + figure_mask_chunk = level_df["heading"].eq("Figure/Image") + exclude_mask_chunk = placeholder_mask_chunk | figure_mask_chunk + non_excluded = level_df[~exclude_mask_chunk].copy() + if not non_excluded.empty: + non_excluded = execute_level_mapping(non_excluded, lvl_mapping) + for _, row in non_excluded.iterrows(): + row_id = row["id"] + if isinstance(row_id, bool): + continue + if isinstance(row_id, int): + llm_levels[row_id] = row["level"] + logger.info( + f"multi-chunk mapping produced {len(llm_levels)} id->level entries" + ) + else: + logger.info( + "single chunk - skipping reason-code mapping, using LLM output directly" + ) + + full_preds = raw_preds.copy()[["id", "heading", "level", "reason"]] + + def resolve_level(row_id: Any) -> int: + try: + int_id = int(row_id) + except (TypeError, ValueError): + return -1 + level = llm_levels.get(int_id, -1) + try: + return int(level) + except (TypeError, ValueError): + return -1 + + full_preds["level"] = full_preds["id"].map(resolve_level).astype(int) + save_intermediate_csv( + full_preds, output_dir, f"preds_4_llm_final{csv_suffix}" + ) + + except Exception as exc: + logger.warning( + f"LLM-based parsing fails due to {exc}, using non-llm pipeline..." + ) + full_preds = fallback_hierarchy(raw_preds.copy()) + return full_preds diff --git a/apps/worker/app/services/document_parser/structure/heading_tree.py b/apps/worker/app/services/document_parser/structure/heading_tree.py new file mode 100644 index 000000000..3f5c39dcc --- /dev/null +++ b/apps/worker/app/services/document_parser/structure/heading_tree.py @@ -0,0 +1,172 @@ +from __future__ import annotations + +import pandas as pd +from loguru import logger + + +def build_tree_from_dataframe( + heading_preds: pd.DataFrame, +) -> tuple[dict[str, dict], dict[tuple[str, str], int], dict[int, dict]]: + headings = heading_preds[heading_preds["level"] > -1].copy() + + node_to_id: dict[tuple[str, str], int] = {} + id_to_row: dict[int, dict] = {} + root: dict[str, dict] = {} + stack: list[tuple[int, dict, str, str]] = [(0, root, "ROOT", "")] + + for _, row in headings.iterrows(): + heading_text = str(row["heading"]) + row_id = int(row["id"]) + level = int(row["level"]) + + id_to_row[row_id] = row.to_dict() + + while len(stack) > 1 and stack[-1][0] >= level: + stack.pop() + + _, parent_dict, _, parent_path = stack[-1] + tree_node_key = heading_text + if tree_node_key in parent_dict: + tree_node_key = f"{heading_text}#{row_id}" + + node_key = (tree_node_key, parent_path) + node_to_id[node_key] = row_id + + parent_dict[tree_node_key] = {} + current_path = ( + f"{parent_path}/{tree_node_key}" if parent_path else tree_node_key + ) + stack.append((level, parent_dict[tree_node_key], tree_node_key, current_path)) + + return root, node_to_id, id_to_row + + +def tree_to_dataframe( + tree: dict[str, dict], + node_to_id: dict[tuple[str, str], int], + original_df: pd.DataFrame, +) -> pd.DataFrame: + preserved_headings = _extract_headings_from_tree(tree, node_to_id) + preserved_ids = {heading["id"] for heading in preserved_headings} + + updated_df = original_df.copy() + removed_count = 0 + level_changed_count = 0 + + for index, row in original_df.iterrows(): + row_id = int(row["id"]) + old_level = int(row["level"]) if row["level"] not in [-2, "nan", -1] else -1 + + if old_level <= -1: + continue + + if row_id in preserved_ids: + new_level = next( + ( + heading["level"] + for heading in preserved_headings + if heading["id"] == row_id + ), + old_level, + ) + updated_df.at[index, "level"] = new_level + if new_level != old_level: + level_changed_count += 1 + else: + updated_df.at[index, "level"] = -1 + removed_count += 1 + + logger.debug( + f"Tree changed: removed headings={removed_count}, " + f"level changed={level_changed_count}, preserved headings={len(preserved_ids)}" + ) + return updated_df + + +def remove_isolated_nodes(tree: dict[str, dict]) -> dict[str, dict]: + return _remove_isolated_nodes_recursive(tree) + + +def cleanup_heading_tree(heading_preds: pd.DataFrame) -> pd.DataFrame: + if heading_preds.empty: + return heading_preds + + tree, node_to_id, _ = build_tree_from_dataframe(heading_preds) + processed_tree = remove_isolated_nodes(tree) + return tree_to_dataframe(processed_tree, node_to_id, heading_preds) + + +def _extract_headings_from_tree( + node_dict: dict[str, dict], + node_to_id: dict[tuple[str, str], int], + *, + current_level: int = 1, + parent_path: str = "", +) -> list[dict[str, object]]: + results: list[dict[str, object]] = [] + for tree_node_key, children in node_dict.items(): + node_key = (tree_node_key, parent_path) + row_id = node_to_id.get(node_key, -1) + + if row_id >= 0: + original_heading = ( + tree_node_key.split("#")[0] if "#" in tree_node_key else tree_node_key + ) + results.append( + { + "id": row_id, + "heading": original_heading, + "level": current_level, + "tree_key": tree_node_key, + "parent_path": parent_path, + } + ) + + if isinstance(children, dict) and children: + current_path = ( + f"{parent_path}/{tree_node_key}" if parent_path else tree_node_key + ) + results.extend( + _extract_headings_from_tree( + children, + node_to_id, + current_level=current_level + 1, + parent_path=current_path, + ) + ) + return results + + +def _remove_isolated_nodes_recursive( + node_dict: dict[str, dict], + *, + parent_path: str = "", +) -> dict[str, dict]: + result_dict: dict[str, dict] = {} + + for heading, children in node_dict.items(): + if isinstance(children, dict) and len(children) == 1: + child_heading = list(children.keys())[0] + grandchildren = children[child_heading] + + if not grandchildren or ( + isinstance(grandchildren, dict) and len(grandchildren) == 0 + ): + result_dict[heading] = {} + logger.debug( + f"remove isolated heading: {parent_path}/{heading}/{child_heading}" + ) + else: + result_dict[heading] = _remove_isolated_nodes_recursive( + children, + parent_path=f"{parent_path}/{heading}" if parent_path else heading, + ) + elif isinstance(children, dict) and children: + result_dict[heading] = _remove_isolated_nodes_recursive( + children, + parent_path=f"{parent_path}/{heading}" if parent_path else heading, + ) + else: + result_dict[heading] = children + + return result_dict diff --git a/apps/worker/app/services/document_parser/structure/layout_parser.py b/apps/worker/app/services/document_parser/structure/layout_parser.py new file mode 100755 index 000000000..47f0650af --- /dev/null +++ b/apps/worker/app/services/document_parser/structure/layout_parser.py @@ -0,0 +1,678 @@ +# pyright: reportArgumentType=false, reportAssignmentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportGeneralTypeIssues=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalSubscript=false +import os + +import gevent +import pandas as pd +from app.services.document_parser.structure.heading_candidates import ( + filter_document_headings, + filter_markdown_headings, + postprocess_headings, +) +from app.services.document_parser.structure.heading_llm_executor import ( + build_level_mapping, + execute_level_mapping, + execute_llm_heading_hierarchy, +) +from app.services.document_parser.structure.heading_tree import ( + build_tree_from_dataframe as build_heading_tree_from_dataframe, +) +from app.services.document_parser.structure.heading_tree import ( + remove_isolated_nodes as remove_isolated_heading_nodes, +) +from app.services.document_parser.structure.heading_tree import ( + tree_to_dataframe as heading_tree_to_dataframe, +) +from app.services.document_parser.support.stage_profiler import stage_timer +from app.services.document_parser.tables.table_text_parser import df2md +from gevent.pool import Pool as GeventPool + +from loguru import logger + +from shared.core.config import settings + +# TaskRedis dependency is removed, use Redis directly to track +from shared.services.ai.prompt_service import build_prompt +from shared.services.ai.response_process_service import eval_response + +# ARQ dependency is removed, use Celery instead +from shared.services.ai.openai_compatible_client_sync import get_openai_client + +# ==================== Helper Functions ==================== + + +def _resolve_hierarchy_model_name(model_name=None): + """Resolve the dedicated hierarchy LLM model with backward-compatible fallback.""" + return model_name or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL + + +def save_intermediate_csv(df: pd.DataFrame, output_dir: str, filename: str): + """ + save intermediate result to csv file, use utf-8-sig encoding to support Chinese and English + Only saves when LOCAL_DEBUG environment variable is set to 'true'. + + Args: + df: DataFrame to save + output_dir: output directory path + filename: filename (without extension) + """ + if os.environ.get("LOCAL_DEBUG", "").lower() not in ("true", "1"): + return + if output_dir is None or df is None or df.empty: + return + + try: + csv_path = os.path.join(output_dir, f"{filename}.csv") + df.to_csv(csv_path, index=False, encoding="utf-8-sig") + logger.debug(f"📊 Saved intermediate result to {csv_path}, rows={len(df)}") + except Exception as e: + logger.warning(f"Failed to save intermediate CSV {filename}: {e}") + + +# ==================== Tree Structure Functions (from sxjg) ==================== + + +def build_tree_from_dataframe(df): + return build_heading_tree_from_dataframe(df) + + +def tree_to_dataframe(tree, node_to_id, original_df): + return heading_tree_to_dataframe(tree, node_to_id, original_df) + + +def remove_isolated_nodes(tree): + return remove_isolated_heading_nodes(tree) + + +def format_toc_context_for_llm(toc_context) -> str: + """Convert TOC hierarchy or structured payloads into compact LLM-friendly plain text.""" + if not toc_context: + return "" + + if isinstance(toc_context, str): + return toc_context + + toc_items = toc_context if isinstance(toc_context, list) else [toc_context] + formatted_blocks = [] + + for toc_idx, toc_item in enumerate(toc_items, start=1): + if not isinstance(toc_item, dict): + formatted_blocks.append(str(toc_item)) + continue + + toc_range = toc_item.get("toc_range") + toc_entries = toc_item.get("toc_with_level") or [] + + if toc_range and len(toc_range) == 2: + formatted_blocks.append( + f"TOC {toc_idx} (source rows {toc_range[0]}-{toc_range[1]}):" + ) + else: + formatted_blocks.append(f"TOC {toc_idx}:") + + if not toc_entries: + formatted_blocks.append("- No TOC entries available") + continue + + if isinstance(toc_entries, str): + toc_entries = toc_entries.strip() + if toc_entries: + formatted_blocks.append(toc_entries) + else: + formatted_blocks.append("- No TOC entries available") + continue + + for entry in toc_entries: + if not isinstance(entry, dict): + continue + + heading = str(entry.get("heading", "")).strip().replace("\n", " ") + if not heading: + continue + + level = entry.get("level") + line_id = entry.get("id") + if isinstance(level, int): + formatted_blocks.append(f"- level {level} | id {line_id} | {heading}") + else: + formatted_blocks.append(f"- id {line_id} | {heading}") + + return "\n".join(formatted_blocks) + + +def hiearchy_llm( + df, + model_name=None, + max_depth=6, + toc_context=None, + max_len=8192, + task="eval-headings", +): + """Apply LLM to analyze the hierarchy of headings + + Args: + df: DataFrame with id, heading columns + model_name: LLM model name (optional, uses default if None) + max_depth: Maximum hierarchy depth + max_len: Hard cap for LLM completion max_tokens (default 2048). + Actual value is derived from the number of heading candidates. + task: Prompt task type - "eval-headings" for general document, "eval-toc-headings" for TOC + toc_context: Optional formatted TOC context string for guiding level assignment + + Returns: + List of dicts with id and level, one per row in ``df`` (missing IDs -> level=-1). + """ + + model_name = _resolve_hierarchy_model_name(model_name) + level_md = df2md(df) + + # Completion budget is driven by the number of heading candidates, not the + # markdown input length. Each JSON entry is `{"id":X,"level":Y}` ≈ 25 tokens; + # add 200 tokens overhead for brackets/whitespace and leave a 512 floor for + # tiny inputs. Non-int ids (placeholders like "10-12" or "-") are excluded. + def _is_candidate_id(val): + if isinstance(val, bool): + return False + if isinstance(val, int): + return True + try: + int(val) + return True + except (TypeError, ValueError): + return False + + n_candidates = int(df["id"].apply(_is_candidate_id).sum()) if len(df) > 0 else 0 + ot_limit = max(512, n_candidates * 25 + 200) + ot_limit = min(ot_limit, max_len) + formatted_toc_context = format_toc_context_for_llm(toc_context) + + paras = { + "max_tokens": ot_limit, + "max_depth": max_depth, + "toc_context": formatted_toc_context, + } + prompt, temperature, top_p, max_tokens = build_prompt( + task=task, texts=level_md, query="", paras=paras + ) + messages = [ + {"role": "system", "content": "you are a document auditing expert"}, + {"role": "user", "content": prompt}, + ] + + try: + with stage_timer( + "heading.hierarchy_llm_call", + model_name=model_name, + row_count=len(df), + task=task, + candidate_count=n_candidates, + max_tokens=max_tokens, + ): + answer = get_openai_client(model=model_name).chat_completion( + messages=messages, + model=model_name, + max_tokens=max_tokens, + temperature=temperature, + ) + layout_res = eval_response(answer) + + # Validate eval_response result — it can return a raw string when JSON parsing fails + if not isinstance(layout_res, list): + raise ValueError( + f"LLM returned non-list response (type={type(layout_res).__name__}), " + f"raw content: {str(layout_res)[:200]}" + ) + + # Validate each item is a dict with required keys + for i, item in enumerate(layout_res): + if not isinstance(item, dict) or "id" not in item or "level" not in item: + raise ValueError(f"LLM response item[{i}] is malformed: {item!r}") + + # Drop items whose id is not a clean integer. This includes placeholder + # rows ("10-12", "-") that the LLM may echo back despite the prompt telling + # it not to. + clean_res = [] + dropped = 0 + for item in layout_res: + raw_id = item["id"] + if isinstance(raw_id, bool): + dropped += 1 + continue + if isinstance(raw_id, int): + clean_res.append({"id": raw_id, "level": item["level"]}) + continue + try: + clean_res.append({"id": int(raw_id), "level": item["level"]}) + except (TypeError, ValueError): + dropped += 1 + if dropped: + logger.debug(f"filtered {dropped} non-integer-id entries from LLM response") + + # LLM only returns heading rows (level >= 1). Reconstruct full result so the + # returned list has one entry per row in ``df``, with missing ids defaulting + # to level=-1. Rows whose ``id`` is itself non-integer (placeholders) keep + # their id as-is and level=-1 so the caller can filter them out. + llm_levels = {item["id"]: item["level"] for item in clean_res} + full_result = [] + for row_id in df["id"].tolist(): + if _is_candidate_id(row_id): + try: + int_id = int(row_id) + except (TypeError, ValueError): + int_id = row_id + full_result.append({"id": int_id, "level": llm_levels.get(int_id, -1)}) + else: + full_result.append({"id": row_id, "level": -1}) + logger.debug( + f"LLM returned {len(clean_res)} heading levels out of {n_candidates} candidates " + f"({len(df)} total rows)" + ) + return full_result + except Exception as e: + logger.error(f"detect hierarchy by LLM failed: {e}") + raise + + +def _compute_zone_boundaries(toc_hierarchies, coordinate_mode="post_removal"): + """Compute content zone boundaries for documents with multiple TOC areas. + + When multiple TOCs exist, they divide the document into zones. Each zone + starts right after a TOC area and extends to just before the next TOC area + (or end of document). + + coordinate_mode: + - "post_removal": TOC ranges are in original coordinates, but heading IDs + are measured after TOC rows were removed (MD/PDF path). + - "original": heading IDs stay in original document coordinates, so zones + can be computed directly from TOC boundaries (DOCX path). + + Args: + toc_hierarchies: List of toc hierarchy dicts (sorted by toc_range start) + + Returns: + List of (zone_start_post, zone_end_post_or_None, toc_hierarchy_dict) + zone_end_post is None for the last zone (extends to end of document) + """ + if coordinate_mode not in {"post_removal", "original"}: + raise ValueError(f"Unsupported coordinate_mode: {coordinate_mode}") + + sorted_tocs = sorted(toc_hierarchies, key=lambda t: t["toc_range"][0]) + + zones = [] + cumulative_removed = 0 + + for i, toc in enumerate(sorted_tocs): + toc_start, toc_end = toc["toc_range"] + zone_start = toc_end + 1 + + if coordinate_mode == "post_removal": + toc_size = toc_end - toc_start + 1 + cumulative_removed += toc_size + zone_start -= cumulative_removed + + if i + 1 < len(sorted_tocs): + next_toc_start = sorted_tocs[i + 1]["toc_range"][0] + zone_end = next_toc_start - 1 + if coordinate_mode == "post_removal": + zone_end -= cumulative_removed + else: + zone_end = None # to end of document + + if zone_end is not None and zone_end < zone_start: + continue + zones.append((zone_start, zone_end, toc)) + + return zones + + +def _resolve_first_toc_boundary(toc_hierarchies=None, first_toc_ele_num=None): + """Resolve the earliest available first-TOC boundary across coordinate sources.""" + toc_range_start = None + if toc_hierarchies: + first_range = toc_hierarchies[0].get("toc_range") + if first_range and len(first_range) == 2: + toc_range_start = first_range[0] + + candidates = [ + value for value in (toc_range_start, first_toc_ele_num) if value is not None + ] + if not candidates: + return None + + resolved_start = min(candidates) + if ( + toc_range_start is not None + and first_toc_ele_num is not None + and toc_range_start != first_toc_ele_num + ): + logger.info( + f"📌 TOC boundary mismatch detected: toc_range start={toc_range_start}, " + f"first_toc_ele_num={first_toc_ele_num}, using earliest={resolved_start}" + ) + return resolved_start + + +def pred_titles( + infos, + doc_type, + toc_hierarchies=None, + prompt_limt=4000, + enable_regx=True, + smart_parse=False, + model_name=None, + output_dir=None, + layout_json_path=None, + first_toc_ele_num=None, +): + """ + predict title hierarchy + + Args: + infos: document information + doc_type: document type (pptx, md, docx) + toc_hierarchies: TOC hierarchy information (if any) + prompt_limt: prompt character limit + enable_regx: whether to enable regex matching + smart_parse: whether to use LLM intelligent parsing + model_name: LLM model name + output_dir: output directory for saving intermediate CSV results + layout_json_path: path to layout.json for META features (optional) + first_toc_ele_num: ele_num of the first TOC block in DOCX (for pre-TOC exclusion) + """ + model_name = _resolve_hierarchy_model_name(model_name) + logger.info( + f"Start to predict title hierarchy: doc_type={doc_type}, smart_parse={smart_parse}, candidate titles={len(infos)}" + ) + + if doc_type == "pptx": + raw_preds = filter_markdown_headings(infos) + elif doc_type == "md": + raw_preds = filter_markdown_headings(infos, layout_json_path=layout_json_path) + elif doc_type == "docx": + raw_preds = filter_document_headings(infos, enable_regex=enable_regx) + else: + raw_preds = pd.DataFrame(columns=["id", "heading", "level", "reason"]) + + # ── Exclude pre-TOC lines from heading prediction ── + # When TOC is detected, lines/blocks before the first TOC area are typically + # cover/metadata (company name, version, classification marks), not real + # headings. Remove them before LLM judging to avoid misjudgment, then + # splice back with level=-1 after all processing is done. + pre_toc_rows = None + first_toc_start = None + if doc_type == "md": + first_toc_start = _resolve_first_toc_boundary(toc_hierarchies=toc_hierarchies) + elif doc_type == "docx": + first_toc_start = _resolve_first_toc_boundary( + toc_hierarchies=toc_hierarchies, + first_toc_ele_num=first_toc_ele_num, + ) + + if first_toc_start is not None and first_toc_start > 0: + pre_toc_mask = raw_preds["id"] < first_toc_start + if pre_toc_mask.any(): + pre_toc_rows = raw_preds[pre_toc_mask].copy() + pre_toc_rows["level"] = -1 + raw_preds = raw_preds[~pre_toc_mask].reset_index(drop=True) + if doc_type == "docx": + logger.info( + f"📌 Excluded {len(pre_toc_rows)} pre-TOC blocks " + f"(id < {first_toc_start}) from heading prediction" + ) + else: + logger.info( + f"📌 Excluded {len(pre_toc_rows)} pre-TOC lines " + f"(id < {first_toc_start}) from heading prediction" + ) + + # 2. Zone-based prediction when multiple TOCs exist + if ( + toc_hierarchies + and len(toc_hierarchies) > 1 + and doc_type in {"md", "docx"} + and smart_parse + ): + # Multiple TOCs divide the document into independent zones. + # Each zone gets its own naive + LLM pipeline with zone-specific TOC context. + coordinate_mode = "post_removal" if doc_type == "md" else "original" + zones = _compute_zone_boundaries( + toc_hierarchies, coordinate_mode=coordinate_mode + ) + logger.info( + f"🗂️ Zone-based prediction: {len(zones)} zones from {len(toc_hierarchies)} TOCs" + ) + + def _process_single_zone(zone_idx, zone_start, zone_end, zone_toc): + """Process a single zone independently. Returns (zone_idx, zone_heading_df).""" + # Extract rows belonging to this zone + if zone_end is not None: + zone_mask = (raw_preds["id"] >= zone_start) & ( + raw_preds["id"] <= zone_end + ) + else: + zone_mask = raw_preds["id"] >= zone_start + zone_preds = raw_preds[zone_mask].copy().reset_index(drop=True) + + if zone_preds.empty: + logger.warning(f" Zone {zone_idx}: empty, skipping") + return zone_idx, None + + zone_range_str = f"[{zone_start}, {zone_end or 'end'}]" + logger.info( + f" Zone {zone_idx}: {len(zone_preds)} rows, post-removal range {zone_range_str}" + ) + + # Independent naive + LLM prediction for this zone + zone_heading = est_hierarchies_naive( + zone_preds, smart_parse, output_dir=output_dir + ) + zone_heading = est_hierarchies_llm( + zone_heading, + prompt_limt, + toc_hierarchies=[zone_toc], # Single TOC for this zone + model_name=model_name, + output_dir=output_dir, + csv_suffix=f"_zone_{zone_idx}", + ) + valid_count = ( + len(zone_heading[zone_heading["level"] > 0]) + if not zone_heading.empty + else 0 + ) + logger.info(f" Zone {zone_idx}: ✅ {valid_count} valid headings") + return zone_idx, zone_heading + + if len(zones) == 1: + # Single zone: no parallel overhead + zone_start, zone_end, zone_toc = zones[0] + _, zone_heading = _process_single_zone(0, zone_start, zone_end, zone_toc) + zone_results = [zone_heading] if zone_heading is not None else [] + else: + # Multiple zones: parallel hierarchy prediction via gevent + logger.info( + f"Parallelizing zone hierarchy prediction for {len(zones)} zones" + ) + pool = GeventPool(size=len(zones)) + greenlets = [ + pool.spawn( + _process_single_zone, zone_idx, zone_start, zone_end, zone_toc + ) + for zone_idx, (zone_start, zone_end, zone_toc) in enumerate(zones) + ] + gevent.joinall(greenlets) + + # Collect results sorted by zone index to maintain document order + results = sorted( + [g.value for g in greenlets if g.value is not None], key=lambda r: r[0] + ) + zone_results = [ + heading_df for _, heading_df in results if heading_df is not None + ] + + if zone_results: + heading_preds = ( + pd.concat(zone_results, ignore_index=True) + .sort_values("id") + .reset_index(drop=True) + ) + else: + heading_preds = pd.DataFrame(columns=["id", "heading", "level", "reason"]) + logger.info("✅ Zone-based LLM hierarchy parsing completed") + else: + # Single-zone: current behavior + heading_preds = est_hierarchies_naive( + raw_preds, smart_parse, output_dir=output_dir + ) + if smart_parse: + heading_preds = est_hierarchies_llm( + heading_preds, + prompt_limt, + toc_hierarchies, + model_name=model_name, + output_dir=output_dir, + ) + logger.info("✅ LLM hierarchy parsing completed") + + # 3. final polishing for certain types + if doc_type in ["docx"]: + heading_preds = postprocess_headings(heading_preds, task="merge_continuous") + heading_preds = postprocess_headings(heading_preds, task="merge_short") + heading_preds = postprocess_headings(heading_preds, task="judge_negs") + logger.debug("Docx hiearchy detection postprocessing completed") + + if heading_preds["level"].eq(-1).all(): # if non are estimated as headings + logger.warning("⚠️ No valid headings estimated") + heading_preds = pd.DataFrame() + else: + heading_preds["level"] = ( + pd.to_numeric(heading_preds["level"], errors="coerce") + .fillna(-1) + .astype(int) + ) + + # process isolated nodes + try: + tree, node_to_id, _ = build_tree_from_dataframe(heading_preds) + processed_tree = remove_isolated_nodes(tree) + heading_preds = tree_to_dataframe(processed_tree, node_to_id, heading_preds) + except Exception as e: + logger.warning(f"Tree structure optimization failed, skipping: {e}") + + logger.info( + f"✅ Heading parsing completed, final {len(heading_preds[heading_preds['level'] > 0])} valid headings" + ) + + # ── Splice pre-TOC rows back ── + if pre_toc_rows is not None and not heading_preds.empty: + heading_preds = ( + pd.concat( + [pre_toc_rows[["id", "heading", "level", "reason"]], heading_preds], + ignore_index=True, + ) + .sort_values("id") + .reset_index(drop=True) + ) + logger.debug( + f"📌 Spliced {len(pre_toc_rows)} pre-TOC lines back into predictions" + ) + + # Save heading_preds as preds_5 + save_intermediate_csv(heading_preds, output_dir, "preds_5_final_output") + return heading_preds + + +def est_hierarchies_naive(raw_preds, proceed_smart=True, output_dir=None): + """Detect hierarchies by non-LLM + + Args: + raw_preds: raw data + proceed_smart: whether to proceed with smart parsing + output_dir: output directory, used to save intermediate results CSV + """ + logger.debug("🚀 non-llm parsing => recursive processing") + save_preds = raw_preds.copy() + + heading_preds = postprocess_headings(raw_preds, task="collapse") + save_preds.insert( + save_preds.columns.get_loc("level") + 1, + "lvl_cola", + heading_preds["level"].tolist(), + ) + + heading_preds = postprocess_headings(heading_preds, task="judge_negs") + save_preds.insert( + save_preds.columns.get_loc("lvl_cola") + 1, + "lvl_neg", + heading_preds["level"].tolist(), + ) + save_preds["reason"] = heading_preds["reason"] + + # mapping based on freq + if not proceed_smart: + heading_preds["level"] = heading_preds["level"].map( + lambda x: -1 if x == -2 else x + ) + heading_preds, lvl_mapping = build_level_mapping( + heading_preds, heading_preds["level"].tolist(), mode="freq" + ) + heading_preds = execute_level_mapping(heading_preds, lvl_mapping) + heading_preds.drop("origin_level", axis=1, inplace=True) + save_preds.insert( + save_preds.columns.get_loc("lvl_neg") + 1, + "lvl_map", + heading_preds["level"].tolist(), + ) + + return heading_preds + + +def est_hierarchies_llm( + raw_preds, + prompt_limt, + toc_hierarchies=None, + max_len=30, + max_depth=6, + model_name=None, + output_dir=None, + csv_suffix="", +): + """LLM-based hierarchy detection — first chunk via LLM, remaining chunks via reason-code mapping. + + When ``KB_LAYOUT_LLM_COMPACT_INPUT`` is enabled (default), consecutive + ``level == -1`` rows in ``raw_preds`` are folded into a single placeholder + row (``[N BODY LINES]``) before chunking. This shrinks the prompt, makes + most documents fit into a single chunk (skipping the lossy reason-code + mapping), and preserves the positional signal for the LLM. + + Strategy: + 1. (Optional) Compact raw_preds so consecutive body rows become placeholders. + 2. Send only the first chunk to LLM for hierarchy prediction. + 3. Collect ``{id -> level}`` from the LLM response (int ids only). + 4. For multi-chunk docs, extend that mapping via reason-code mapping on + chunks 1..N (placeholders excluded). + 5. Expand the id->level mapping back onto the ORIGINAL ``raw_preds``; + any row not present in the mapping defaults to ``level = -1``. + + Args: + raw_preds: raw data + prompt_limt: prompt character limit + toc_hierarchies: TOC hierarchies + max_len: maximum heading length for executor chunk preparation + max_depth: maximum hierarchy depth + model_name: LLM model name + output_dir: output directory, used to save intermediate results CSV + csv_suffix: suffix for intermediate CSV filenames + """ + model_name = _resolve_hierarchy_model_name(model_name) + return execute_llm_heading_hierarchy( + raw_preds=raw_preds, + prompt_limt=prompt_limt, + hierarchy_judge=hiearchy_llm, + fallback_hierarchy=est_hierarchies_naive, + save_intermediate_csv=save_intermediate_csv, + toc_hierarchies=toc_hierarchies, + max_len=max_len, + max_depth=max_depth, + model_name=model_name, + output_dir=output_dir, + csv_suffix=csv_suffix, + ) diff --git a/apps/worker/app/services/document_parser/metadata_extractor.py b/apps/worker/app/services/document_parser/structure/metadata_extractor.py similarity index 97% rename from apps/worker/app/services/document_parser/metadata_extractor.py rename to apps/worker/app/services/document_parser/structure/metadata_extractor.py index 3322f15ac..52d595424 100644 --- a/apps/worker/app/services/document_parser/metadata_extractor.py +++ b/apps/worker/app/services/document_parser/structure/metadata_extractor.py @@ -124,8 +124,8 @@ def extract_md_headings(md_lines: List[str]) -> List[dict]: if text: text_key = normalize_content(text) occurrence_counter[text_key] += 1 - except Exception: - pass + except Exception as exc: + logger.debug(f"Failed to collect table metadata from Markdown: {exc}") continue if is_heading: @@ -212,8 +212,10 @@ def process_block(block, page_idx): "type": "table", } ) - except Exception: - pass + except Exception as exc: + logger.debug( + f"Failed to collect table metadata from layout span: {exc}" + ) continue # Regular text span diff --git a/apps/worker/app/services/document_parser/structure/toc_hierarchy.py b/apps/worker/app/services/document_parser/structure/toc_hierarchy.py new file mode 100644 index 000000000..439223b38 --- /dev/null +++ b/apps/worker/app/services/document_parser/structure/toc_hierarchy.py @@ -0,0 +1,149 @@ +from __future__ import annotations + +import pandas as pd +from app.services.document_parser.structure.layout_parser import hiearchy_llm +from app.services.document_parser.support.stage_profiler import stage_timer +from app.services.document_parser.tables.table_text_parser import df2md +from app.services.document_parser.support.text_helpers import normalize_md +from loguru import logger +from pandas import Index + +from shared.core.config import settings + + +def resolve_hierarchy_model_name(model_name: str | None = None) -> str: + return model_name or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL + + +def parse_toc_hierarchy( + toc_df: pd.DataFrame, max_depth: int = 6, model_name: str | None = None +) -> list[dict]: + resolved_model_name = resolve_hierarchy_model_name(model_name) + try: + with stage_timer( + "toc.parse_hierarchy_llm", + model_name=resolved_model_name, + heading_count=len(toc_df), + max_depth=max_depth, + ): + toc_hierarchy = hiearchy_llm( + toc_df, + model_name=resolved_model_name, + max_depth=max_depth, + task="eval-toc-headings", + ) + id_to_level = {item["id"]: item["level"] for item in toc_hierarchy} + + toc_with_level = [] + for _, row in toc_df.iterrows(): + line_id = row["id"] + heading = row["heading"] + level = id_to_level.get(line_id, 1) + toc_with_level.append({"id": line_id, "heading": heading, "level": level}) + return toc_with_level + + except Exception as exc: + logger.error(f"LLM hierarchy analysis failed: {exc}") + return [] + + +def build_tree_tocs(toc_with_level: list[dict]) -> dict: + if not toc_with_level: + return {} + + positive_levels = [item["level"] for item in toc_with_level if item["level"] > 0] + level_for_minus_one = max(positive_levels) + 1 if positive_levels else 1 + + root = {} + stack = [(root, 0)] + + for item in toc_with_level: + heading = item["heading"] + original_level = item["level"] + normalized_level = ( + level_for_minus_one if original_level == -1 else original_level + ) + while len(stack) > 1 and stack[-1][1] >= normalized_level: + stack.pop() + + parent_dict = stack[-1][0] + parent_dict[heading] = {} + stack.append((parent_dict[heading], normalized_level)) + return root + + +def build_toc_hierarchy_payload( + toc_entries: list[dict], + toc_range: tuple | None = None, + scan_range: tuple | None = None, +) -> dict | None: + valid_entries = [] + for entry in toc_entries: + heading = str(entry.get("heading", "")).strip() + level = entry.get("level") + if not heading or not isinstance(level, int) or level <= 0: + continue + + valid_entries.append( + { + "id": entry.get("id"), + "heading": heading, + "level": level, + } + ) + + if not valid_entries: + return None + + toc_df = pd.DataFrame(valid_entries, columns=Index(["id", "heading", "level"])) + payload = { + "toc_range": toc_range or (valid_entries[0]["id"], valid_entries[-1]["id"]), + "toc_with_level": df2md(toc_df, index=False), + "toc_tree": build_tree_tocs(valid_entries), + } + if scan_range is not None: + payload["scan_range"] = scan_range + return payload + + +def eval_toc_levels( + toc_lines: list[str], model_name: str | None = None, max_depth: int = 6 +) -> tuple[str, dict]: + toc_title_keywords = {"目录", "目次", "tableofcontents", "contents"} + valid_data = [] + + for index, line in enumerate(toc_lines): + heading = line.strip() + if not heading: + continue + if normalize_md(heading) in toc_title_keywords: + logger.debug( + f"eval_toc_levels: skipping TOC keyword title line id={index}: {heading[:60]}" + ) + continue + + valid_data.append({"id": index, "heading": heading, "level": "Not Sure"}) + + toc_df = pd.DataFrame(valid_data) + + if toc_df.empty: + logger.info("No valid TOC content, skip hierarchy analysis") + return "", {} + + llm_result = parse_toc_hierarchy(toc_df, max_depth, model_name) + id_to_level = {item["id"]: item["level"] for item in llm_result} + + valid_items_for_tree = [] + for data in valid_data: + line_id = data["id"] + heading = data["heading"] + level = id_to_level.get(line_id, -1) + if level > 0: + valid_items_for_tree.append( + {"id": line_id, "heading": heading, "level": level} + ) + + payload = build_toc_hierarchy_payload(valid_items_for_tree) + if not payload: + return "", {} + return payload["toc_with_level"], payload["toc_tree"] diff --git a/apps/worker/app/services/document_parser/toc_parser.py b/apps/worker/app/services/document_parser/structure/toc_parser.py similarity index 51% rename from apps/worker/app/services/document_parser/toc_parser.py rename to apps/worker/app/services/document_parser/structure/toc_parser.py index e192bea40..9380884f3 100644 --- a/apps/worker/app/services/document_parser/toc_parser.py +++ b/apps/worker/app/services/document_parser/structure/toc_parser.py @@ -5,309 +5,23 @@ Provides functionality for: - Detecting TOC (Table of Contents) candidates in markdown documents -- Detecting TOC in DOCX documents (SDT containers, styles, field codes) - Using LLM to determine precise TOC boundaries -- Analyzing TOC hierarchy structure -- Building nested tree structures from TOC """ import re import gevent import pandas as pd -from app.services.common.kb_utils import ( - normalize_md, - truncate_text_by_tokens, -) -from app.services.document_parser.layout_parser import ( - hiearchy_llm, - judge_by_conditions, - remove_by_conditions, -) -from app.services.document_parser.stage_profiler import stage_timer -from app.services.document_parser.table_parser import df2md +from app.services.document_parser.structure.toc_hierarchy import eval_toc_levels +from app.services.document_parser.support.text_helpers import normalize_md, truncate_text_by_tokens +from app.services.document_parser.support.stage_profiler import stage_timer +from app.services.document_parser.tables.table_text_parser import df2md from gevent.pool import Pool as GeventPool from loguru import logger -from lxml import etree -from shared.core.config import settings from shared.services.ai.prompt_service import build_prompt from shared.services.ai.response_process_service import eval_response -from shared.utils.OpenAICompatibleClientSync import get_openai_client - - -def _resolve_hierarchy_model_name(model_name: str | None = None) -> str: - """Resolve dedicated hierarchy model, falling back to the normal model.""" - return model_name or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL - - -# ==================== DOCX TOC Detection Functions ==================== - -TOC_TITLE_KEYWORDS = {"目录", "目次", "contents", "table of contents"} - - -def _parse_w_int_attr(elem, ns, attr_names): - """Parse the first integer-valued OOXML attribute from a list of names.""" - if elem is None: - return None - - for attr_name in attr_names: - raw_val = elem.get("{%s}%s" % (ns["w"], attr_name)) - if raw_val is None: - continue - try: - return int(raw_val) - except (TypeError, ValueError): - continue - return None - - -def get_docx_toc_layout_hints(elem, ns): - """Extract outline/indent hints for TOC paragraphs without numeric TOC styles.""" - ppr = elem.find("./w:pPr", namespaces=ns) - if ppr is None: - ppr = elem.find(".//w:pPr", namespaces=ns) - - if ppr is None: - return { - "outline_level": None, - "left_indent": None, - } - - outline_elem = ppr.find("./w:outlineLvl", namespaces=ns) - outline_level = None - if outline_elem is not None: - outline_val = _parse_w_int_attr(outline_elem, ns, ["val"]) - if outline_val is not None: - outline_level = outline_val + 1 - - indent_elem = ppr.find("./w:ind", namespaces=ns) - left_indent = _parse_w_int_attr( - indent_elem, ns, ["left", "start", "leftChars", "startChars"] - ) - - return { - "outline_level": outline_level, - "left_indent": left_indent, - } - - -def infer_toc_level_from_text(text: str): - """Fallback TOC level inference from numbering patterns in TOC text.""" - text_clean = str(text).strip() - if not text_clean: - return None - - normalized = re.sub(r"\s+", " ", text_clean).lower() - if normalized in TOC_TITLE_KEYWORDS: - return None - - pos_code = judge_by_conditions(text_clean) - neg_code = remove_by_conditions(text_clean) - if any(x > 0 for x in neg_code) or not any(x > 0 for x in pos_code): - return None - - return max(int(x) for x in pos_code) - - -def is_toc_title_text(text: str) -> bool: - """Return True when the line is likely the standalone TOC heading itself.""" - normalized = re.sub(r"\s+", " ", str(text).strip()).lower() - return normalized in TOC_TITLE_KEYWORDS - - -def infer_toc_levels_from_indentation(entries: list) -> None: - """Populate missing TOC levels by ranking paragraph indentation within one TOC area.""" - indent_values = sorted( - { - entry["left_indent"] - for entry in entries - if entry.get("level") is None - and entry.get("left_indent") is not None - and not is_toc_title_text(entry.get("heading", "")) - } - ) - - if not indent_values: - return - - indent_to_level = {indent: idx + 1 for idx, indent in enumerate(indent_values)} - for entry in entries: - if entry.get("level") is not None: - continue - if is_toc_title_text(entry.get("heading", "")): - continue - left_indent = entry.get("left_indent") - if left_indent is None: - continue - entry["level"] = indent_to_level.get(left_indent) - - -def get_docx_toc_style_info(elem, ns): - """ - Parse TOC style metadata from a DOCX paragraph element. - - Returns: - dict: { - 'is_toc_style': bool, - 'toc_level': Optional[int], - 'style_name': Optional[str] - } - """ - style = elem.find(".//w:pPr/w:pStyle", namespaces=ns) - if style is None: - return { - "is_toc_style": False, - "toc_level": None, - "style_name": None, - } - - val = style.get("{%s}val" % ns["w"]) - if not val: - return { - "is_toc_style": False, - "toc_level": None, - "style_name": None, - } - - val_lower = val.lower().strip() - if "toc" not in val_lower and "目录" not in val: - return { - "is_toc_style": False, - "toc_level": None, - "style_name": val, - } - - level = None - match = re.search(r"(?:toc|目录)\s*[_-]?(\d+)$", val_lower) - if match: - level = int(match.group(1)) - - layout_hints = get_docx_toc_layout_hints(elem, ns) - if level is None: - level = layout_hints["outline_level"] - - return { - "is_toc_style": True, - "toc_level": level, - "style_name": val, - "outline_level": layout_hints["outline_level"], - "left_indent": layout_hints["left_indent"], - } - - -def get_toc_level(elem, ns): - """ - Detect whether a paragraph uses a TOC style. - - Args: - elem: XML paragraph element. - ns: XML namespace map. - - Returns: - bool: True when the paragraph uses a TOC style. - """ - style_info = get_docx_toc_style_info(elem, ns) - if not style_info["is_toc_style"]: - return False - - if style_info["toc_level"] is not None: - return style_info["toc_level"] - return True - - -def detect_sdt_toc(elem, ns): - """ - Detect an SDT (Structured Document Tag) TOC container. - Word-generated TOCs are often wrapped in ``sdt`` elements. - - Args: - elem: SDT element. - ns: XML namespace map. - - Returns: - dict: { - 'is_toc_sdt': bool - whether the element is a TOC SDT, - 'gallery_type': str - docPartGallery type - } - """ - tag = etree.QName(elem.tag).localname if isinstance(elem.tag, str) else None - - if tag != "sdt": - return {"is_toc_sdt": False, "gallery_type": None} - - is_toc_sdt = False - gallery_type = None - - sdt_pr = elem.find(".//w:sdtPr", namespaces=ns) - if sdt_pr is not None: - doc_part_obj = sdt_pr.find(".//w:docPartObj", namespaces=ns) - if doc_part_obj is not None: - doc_part_gallery = doc_part_obj.find(".//w:docPartGallery", namespaces=ns) - if doc_part_gallery is not None: - gallery_type = doc_part_gallery.get("{%s}val" % ns["w"]) - if gallery_type and "table of contents" in gallery_type.lower(): - is_toc_sdt = True - - return {"is_toc_sdt": is_toc_sdt, "gallery_type": gallery_type} - - -def detect_doc_tocs(elem, ns): - """ - Detect TOC regions using two strategies: - 1. paragraph style detection (TOC styles) - 2. field code detection (instrText) - - Note: SDT container detection is handled by ``detect_sdt_toc``. - - Args: - elem: XML paragraph element. - ns: XML namespace map. - - Returns: - dict: { - 'is_style': bool - whether the paragraph uses a TOC style, - 'is_field_start': bool - whether this starts a TOC field, - 'is_field_end': bool - whether this ends a field - } - """ - style_info = get_docx_toc_style_info(elem, ns) - is_style = style_info["is_toc_style"] - is_field_start = False - - instrs = elem.findall(".//w:instrText", namespaces=ns) - for instr in instrs: - if instr.text: - instr_text_stripped = instr.text.strip() - instr_text_lower = instr_text_stripped.lower() - # Match standalone TOC field commands, NOT "PAGEREF _TocXXXX" - # TOC fields start with "TOC" as the command word - if ( - instr_text_lower.startswith("toc") - or "table of contents" in instr_text_lower - or "目录" in instr_text_stripped - ): - is_field_start = True - break - - is_field_end = False - # Always check for fldChar end, even on TOC-styled paragraphs, - # so that the outer TOC field boundary can be properly closed. - fldchars = elem.findall(".//w:fldChar", namespaces=ns) - for fld in fldchars: - if fld.get("{%s}fldCharType" % ns["w"]) == "end": - is_field_end = True - break - - return { - "is_style": is_style, - "toc_level": style_info["toc_level"], - "style_name": style_info["style_name"], - "outline_level": style_info.get("outline_level"), - "left_indent": style_info.get("left_indent"), - "is_field_start": is_field_start, - "is_field_end": is_field_end, - } +from shared.services.ai.openai_compatible_client_sync import get_openai_client # ==================== Markdown TOC Detection Functions ==================== @@ -658,275 +372,6 @@ def _judge_single_area(idx, lines_, invalid_ids, area_start, area_end): return toc_ranges -def parse_toc_hierarchy(toc_df, max_depth: int = 6, model_name: str = None) -> list: - """ - Parse TOC hierarchy using LLM - - Args: - toc_df: DataFrame with id, heading columns - max_depth: max depth of hierarchy - model_name: model name (optional) - - Returns: - List of dicts with id, heading, level - """ - resolved_model_name = _resolve_hierarchy_model_name(model_name) - try: - with stage_timer( - "toc.parse_hierarchy_llm", - model_name=resolved_model_name, - heading_count=len(toc_df), - max_depth=max_depth, - ): - toc_hierarchy = hiearchy_llm( - toc_df, - model_name=resolved_model_name, - max_depth=max_depth, - task="eval-toc-headings", - ) - id_to_level = {item["id"]: item["level"] for item in toc_hierarchy} - - toc_with_level = [] - for _, row in toc_df.iterrows(): - line_id = row["id"] - heading = row["heading"] - level = id_to_level.get(line_id, 1) - toc_with_level.append({"id": line_id, "heading": heading, "level": level}) - return toc_with_level - - except Exception as e: - logger.error(f"LLM hierarchy analysis failed: {e}") - return [] - - -def build_tree_tocs(toc_with_level: list) -> dict: - """ - Build nested JSON from TOC with level - - Args: - toc_with_level: [{"id": line index, "heading": content, "level": level, "reason": ...}, ...] - level: 1 for h1, 2 for h2..., -1 will be treated as the lowest level title - - Returns: - nested JSON structure - - Notes: - in the TOC scenario, all lines are treated as titles: - - normal levels (1,2,3...) are treated as is - - -1 is treated as a level deeper than all normal levels - """ - if not toc_with_level: - return {} - - # Step 1: collect all levels (exclude -1) - positive_levels = [item["level"] for item in toc_with_level if item["level"] > 0] - - # Step 2: determine the level that -1 should be mapped to - if positive_levels: - # if there are normal levels, -1 is mapped to max + 1 - max_positive_level = max(positive_levels) - level_for_minus_one = max_positive_level + 1 - else: - level_for_minus_one = 1 - - # Step 3: build nested structure - root = {} - stack = [(root, 0)] - - for item in toc_with_level: - heading = item["heading"] - original_level = item["level"] - - # normalize level: -1 -> level_for_minus_one - normalized_level = ( - level_for_minus_one if original_level == -1 else original_level - ) - while len(stack) > 1 and stack[-1][1] >= normalized_level: - stack.pop() - - parent_dict = stack[-1][0] - parent_dict[heading] = {} - stack.append((parent_dict[heading], normalized_level)) - return root - - -def build_toc_hierarchy_payload( - toc_entries: list, - toc_range: tuple | None = None, - scan_range: tuple | None = None, -) -> dict | None: - """ - Build a toc_hierarchies-compatible payload from structured TOC entries. - """ - valid_entries = [] - for entry in toc_entries: - heading = str(entry.get("heading", "")).strip() - level = entry.get("level") - if not heading or not isinstance(level, int) or level <= 0: - continue - - normalized_entry = { - "id": entry.get("id"), - "heading": heading, - "level": level, - } - valid_entries.append(normalized_entry) - - if not valid_entries: - return None - - result_df = pd.DataFrame(valid_entries) - payload = { - "toc_range": toc_range or (valid_entries[0]["id"], valid_entries[-1]["id"]), - "toc_with_level": df2md(result_df[["id", "heading", "level"]], index=False), - "toc_tree": build_tree_tocs(valid_entries), - } - if scan_range is not None: - payload["scan_range"] = scan_range - return payload - - -def build_docx_toc_hierarchies(block_tuples: list) -> list: - """ - Convert DOCX TOC blocks into the same toc_hierarchies structure used by MD/PDF. - """ - toc_areas = [] - current_area = [] - - for ele_num, block, label, meta in block_tuples: - if "TOC" in label: - current_area.append((ele_num, block, meta or {})) - continue - - if current_area: - toc_areas.append(current_area) - current_area = [] - - if current_area: - toc_areas.append(current_area) - - toc_hierarchies = [] - for area in toc_areas: - toc_entries = [] - for ele_num, block, meta in area: - toc_level = meta.get("toc_level") - try: - toc_level = int(toc_level) if toc_level is not None else None - except (TypeError, ValueError): - toc_level = None - - text = getattr(block, "text", str(block)).strip() - if not text: - continue - - if toc_level is None: - outline_level = meta.get("toc_outline_level") - try: - toc_level = ( - int(outline_level) if outline_level is not None else None - ) - except (TypeError, ValueError): - toc_level = None - - if toc_level is None: - toc_level = infer_toc_level_from_text(text) - - left_indent = meta.get("toc_left_indent") - try: - left_indent = int(left_indent) if left_indent is not None else None - except (TypeError, ValueError): - left_indent = None - - toc_entries.append( - { - "id": ele_num, - "heading": text, - "level": toc_level if toc_level and toc_level > 0 else None, - "left_indent": left_indent, - } - ) - - infer_toc_levels_from_indentation(toc_entries) - payload = build_toc_hierarchy_payload( - toc_entries, - toc_range=(area[0][0], area[-1][0]), - scan_range=(area[0][0], area[-1][0]), - ) - if payload: - toc_hierarchies.append(payload) - - return toc_hierarchies - - -def eval_toc_levels( - toc_lines: list, model_name: str = None, max_depth: int = 6 -) -> tuple: - """ - Analyze TOC hierarchy and generate nested JSON - - Args: - toc_lines: list of pre-filtered valid TOC lines (invalid content already removed) - model_name: model name (optional) - max_depth: max depth of hierarchy - - Returns: - (toc_with_level, toc_tree) - - toc_with_level: list with level information - Format: [{"id": int, "heading": str, "level": int, "reason": str}, ...] - - toc_tree: nested JSON structure - """ - # Build data for LLM judgment (all lines are valid, pre-filtered) - # TOC title trigger lines are excluded from - # the LLM input: they stay within toc_range so they are stripped from md_lines, - # but they must not be sent to the hierarchy LLM to avoid a spurious Level=1 entry. - _toc_title_keywords = {"目录", "目次", "tableofcontents", "contents"} - valid_data = [] - - for i, line in enumerate(toc_lines): - heading = line.strip() - if not heading: - continue - if normalize_md(heading) in _toc_title_keywords: - logger.debug( - f"eval_toc_levels: skipping TOC keyword title line id={i}: {heading[:60]}" - ) - continue - - valid_data.append({"id": i, "heading": heading, "level": "Not Sure"}) - - toc_df = pd.DataFrame(valid_data) - - if toc_df.empty: - logger.info("No valid TOC content, skip hierarchy analysis") - return "", {} - - # Evaluate TOC hierarchy with LLM - llm_result = parse_toc_hierarchy(toc_df, max_depth, model_name) - - # Build id -> level mapping from LLM result - id_to_level = {item["id"]: item["level"] for item in llm_result} - - # Build final result - result_data = [] - valid_items_for_tree = [] - - for data in valid_data: - line_id = data["id"] - heading = data["heading"] - level = id_to_level.get(line_id, -1) - - if level > 0: - result_data.append({"id": line_id, "heading": heading, "level": level}) - valid_items_for_tree.append( - {"id": line_id, "heading": heading, "level": level} - ) - - payload = build_toc_hierarchy_payload(valid_items_for_tree) - if not payload: - return "", {} - return payload["toc_with_level"], payload["toc_tree"] - - def detect_tocs_in_texts( md_lines: list, model_name: str = None, diff --git a/apps/worker/app/services/document_parser/support/identifiers.py b/apps/worker/app/services/document_parser/support/identifiers.py new file mode 100644 index 000000000..c8d177c3a --- /dev/null +++ b/apps/worker/app/services/document_parser/support/identifiers.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +import uuid +from datetime import datetime + + +def gen_str_codes(input_string: str) -> str: + """Generate a UUID5 code from a string.""" + return str(uuid.uuid5(uuid.NAMESPACE_DNS, input_string)) + + +def get_str_time() -> str: + """Get the current time as a string.""" + return datetime.now().strftime("%Y-%m-%d %H:%M:%S") diff --git a/apps/worker/app/services/document_parser/internal_parse_name.py b/apps/worker/app/services/document_parser/support/internal_parse_name.py similarity index 97% rename from apps/worker/app/services/document_parser/internal_parse_name.py rename to apps/worker/app/services/document_parser/support/internal_parse_name.py index 785fe6ff5..568dc13c6 100644 --- a/apps/worker/app/services/document_parser/internal_parse_name.py +++ b/apps/worker/app/services/document_parser/support/internal_parse_name.py @@ -4,7 +4,7 @@ import os from dataclasses import dataclass -from shared.utils.file_utils import path_handle +from app.services.common.file_utils import path_handle @dataclass(frozen=True) diff --git a/apps/worker/app/services/document_parser/parser_log_utils.py b/apps/worker/app/services/document_parser/support/parser_log_utils.py similarity index 100% rename from apps/worker/app/services/document_parser/parser_log_utils.py rename to apps/worker/app/services/document_parser/support/parser_log_utils.py diff --git a/apps/worker/app/services/document_parser/support/parser_rows.py b/apps/worker/app/services/document_parser/support/parser_rows.py new file mode 100644 index 000000000..839ddd9f6 --- /dev/null +++ b/apps/worker/app/services/document_parser/support/parser_rows.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import pandas as pd +from pandas import Index + +from shared.core.config import settings + +PARSER_ROW_COLUMNS: tuple[str, ...] = tuple(settings.ALL_DF_COLS.split(",")) + + +@dataclass(frozen=True) +class ParsedRow: + content: str + path: str + type: str + know_id: str + addtime: str + keywords: str = "" + summary: str = "" + tokens: str = "" + connectto: str = "" + page_nums: str = "" + length: int | None = None + + def to_list(self) -> list[object]: + content_length = self.length if self.length is not None else len(self.content) + return [ + self.content, + self.path, + self.type, + content_length, + self.keywords, + self.summary, + self.know_id, + self.tokens, + self.connectto, + self.addtime, + self.page_nums, + ] + + def to_dict(self) -> dict[str, object]: + return dict(zip(PARSER_ROW_COLUMNS, self.to_list())) + + +class ParsedRowsBuilder: + def __init__(self) -> None: + self._rows: list[ParsedRow] = [] + + def append(self, row: ParsedRow) -> None: + self._rows.append(row) + + def extend(self, rows: list[ParsedRow]) -> None: + self._rows.extend(rows) + + def to_dataframe(self) -> pd.DataFrame: + return pd.DataFrame( + [row.to_list() for row in self._rows], + columns=Index(PARSER_ROW_COLUMNS), + ) diff --git a/apps/worker/app/services/document_parser/support/path_helpers.py b/apps/worker/app/services/document_parser/support/path_helpers.py new file mode 100644 index 000000000..4119ff53e --- /dev/null +++ b/apps/worker/app/services/document_parser/support/path_helpers.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import os +import re +from typing import Any + +from bs4 import BeautifulSoup +from shared.utils.chunk_refs import extract_chunk_refs +from app.services.common.file_utils import path_handle + +SUMMARY_PATH_MARKERS: tuple[str, ...] = ("summary", "\u6458\u8981\u603b\u7ed3") + + +def find_images(folder_path: str) -> list[str]: + """Find image files inside a folder tree.""" + image_extensions = {".png", ".jpg", ".jpeg"} + image_files: list[str] = [] + + for _, _, files in os.walk(folder_path): + files.sort() + for file in files: + if os.path.splitext(file)[1].lower() in image_extensions: + image_files.append(file) + return image_files + + +def find_matches_parsing(content: str, path: str) -> str: + """Parse table and image markers from content.""" + matches = extract_chunk_refs(content) + match_type = "PTXT" if len(matches) == 0 else "\n".join((["PTXT"] + matches)) + + split_char = os.getenv("SPLIT_CHAR", "/") + if any( + f"{split_char}{summary_marker}" in path + for summary_marker in SUMMARY_PATH_MARKERS + ): + parent_path = path.split(split_char)[-2] + match_type = "SUMMARY_" + parent_path + "_SUMMARY" + return match_type + + +def flatten_dic2paths( + d: dict[str, Any], + current_path: list[str] | None = None, + result: list[str] | None = None, +) -> list[str]: + """Flatten a nested dict into path strings.""" + if result is None: + result = [] + if current_path is None: + current_path = [] + + for key, value in d.items(): + if not isinstance(key, str): + continue + new_path = current_path + [key] + if isinstance(value, dict) and value: + flatten_dic2paths(value, new_path, result) + else: + split_char = os.getenv("SPLIT_CHAR", "/") + result.append(split_char.join(new_path)) + return result + + +def process_path_texts(path_: str, last: int = 50) -> str: + """Normalize path text for downstream use.""" + temp_path = path_handle(path_, mode="sanitize") + if not isinstance(temp_path, str) or temp_path == "": + return "" + return "_".join(temp_path.split(os.sep))[:last] + + +def remove_spaces(text: str, handle_punctuation: bool = False) -> str: + """Remove spaces between Chinese chars while keeping English word spacing.""" + if handle_punctuation: + punctuation = ( + r"""!"#$%&'()*+,-./:;<=>?@[\]^_`{|}~，。、【】《》？；：''""（）…—-！""" + ) + res_text = re.sub(f"[{re.escape(punctuation)}]", "", text) + else: + pattern = re.compile(r"([\u4e00-\u9fff])\s+|(?<=\s)([\u4e00-\u9fff])") + + def replacer(match: re.Match[str]) -> str: + return match.group(1) or match.group(2) + + res_text = pattern.sub(replacer, text) + + res_text = re.sub(r"\s+", " ", res_text) + return res_text.strip() + + +def traverse_dict(d: dict[str, Any], parent: str | None = None) -> list[str]: + """Traverse a dictionary and generate description text.""" + dic_texts: list[str] = [] + for key, value in d.items(): + if value: + child_keys = ", ".join(value.keys()) + text = f"'{key}' includes {child_keys}" + dic_texts.append(text) + dic_texts.extend(traverse_dict(value, key)) + return dic_texts + + +def restore_graph_by_paths(paths: list[str]) -> tuple[dict[str, Any], list[str]]: + """Rebuild a graph structure from path strings.""" + root_dict: dict[str, Any] = {} + split_char = os.getenv("SPLIT_CHAR", "/") + for path in paths: + nodes = path.split(split_char) + current_dict = root_dict + for node in nodes: + if node not in current_dict: + current_dict[node] = {} + current_dict = current_dict[node] + dic_texts = traverse_dict(root_dict) + return root_dict, dic_texts + + +def html2txt(html_text: str) -> str: + """Convert HTML into plain text.""" + soup = BeautifulSoup(html_text, "html.parser") + return soup.get_text() diff --git a/apps/worker/app/services/document_parser/stage_profiler.py b/apps/worker/app/services/document_parser/support/stage_profiler.py similarity index 100% rename from apps/worker/app/services/document_parser/stage_profiler.py rename to apps/worker/app/services/document_parser/support/stage_profiler.py diff --git a/apps/worker/app/services/document_parser/support/text_helpers.py b/apps/worker/app/services/document_parser/support/text_helpers.py new file mode 100644 index 000000000..b11fdc772 --- /dev/null +++ b/apps/worker/app/services/document_parser/support/text_helpers.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import re + +from shared.utils.text_utils import _CN_EN_NUM_RE + +_CN_CHAR_RE = re.compile(r"[\u4e00-\u9fff]") +EN_START_LIMIT = 15 +CN_RATIO_THRESHOLD = 0.3 + + +def normalize_md(text: str) -> str: + """Normalize markdown string for comparison.""" + text = re.sub(r"^\s*#+\s*", "", text) + text = re.sub(r"\s+", "", text) + return text.lower() + + +def truncate_text(text: str, start_limit: int, end_limit: int) -> str: + """Truncate text by raw character count, keeping start and end parts.""" + text = str(text) + total_limit = start_limit + end_limit + if len(text) <= total_limit: + return text + start_part = text[:start_limit] + end_part = text[-end_limit:] if end_limit > 0 else "" + return f"{start_part}...{end_part}" + + +def detect_primary_lang(text: str) -> str: + """Detect whether text is primarily Chinese or English/other.""" + if not text: + return "en" + tokens = _CN_EN_NUM_RE.findall(text) + if not tokens: + return "en" + cn_count = sum(1 for token in tokens if _CN_CHAR_RE.fullmatch(token)) + return "zh" if (cn_count / len(tokens)) >= CN_RATIO_THRESHOLD else "en" + + +def count_cn_en(text: str) -> int: + """Count semantic Chinese/English/number tokens in a string.""" + return len(_CN_EN_NUM_RE.findall(str(text))) + + +def truncate_text_by_tokens( + text: str, + start_limit: int, + end_limit: int, + lang_aware: bool = True, +) -> str: + """Truncate text by semantic token count, preserving whole words.""" + text = str(text) + matches = list(_CN_EN_NUM_RE.finditer(text)) + total = len(matches) + + if lang_aware and total > 0 and detect_primary_lang(text) == "en": + start_limit = min(start_limit, EN_START_LIMIT) + + if total <= start_limit + end_limit: + return text + + cut_start = matches[start_limit - 1].end() if start_limit > 0 else 0 + cut_end = matches[total - end_limit].start() if end_limit > 0 else len(text) + if cut_start >= cut_end: + return text + return text[:cut_start] + "..." + text[cut_end:] diff --git a/apps/worker/app/services/document_parser/table_parser.py b/apps/worker/app/services/document_parser/table_parser.py deleted file mode 100755 index 74860b867..000000000 --- a/apps/worker/app/services/document_parser/table_parser.py +++ /dev/null @@ -1,1633 +0,0 @@ -# pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportOptionalOperand=false, reportOptionalSubscript=false, reportReturnType=false -import datetime -import io -import os -import re -import threading -import uuid -from collections import OrderedDict -from typing import Dict, List, Optional, Tuple, Union - -import numpy as np -import openpyxl -import pandas as pd -from app.services.common.kb_utils import ( - flatten_dic2paths, - gen_str_codes, - get_str_time, - process_dup_paths_df, - remove_spaces, -) -from app.services.document_parser.html_parser import df2html -from bs4 import BeautifulSoup -from loguru import logger - -from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import TableParsingException -from shared.core.exceptions.knowhere_exception import KnowhereException -from shared.services.ai.prompt_service import build_prompt -from shared.services.ai.response_process_service import eval_response -from shared.utils.chunk_refs import build_chunk_ref -from shared.utils.CommonHelperSync import load_file_bytes -from shared.utils.file_utils import path_handle -from shared.utils.OpenAICompatibleClientSync import get_openai_client -from shared.utils.text_utils import remove_duplicates_orderkept, tokenize2stw_remove - -# ── Table filename sanitizer ──────────────────────────── -# Max byte-safe filename length. Most filesystems cap at 255 bytes; we leave -# room for the "table-N " prefix (~10 chars) and ".html" suffix (5 chars). -_MAX_TABLE_NAME_CHARS = 80 - - -def sanitize_table_name_from_header(raw_header_text: str) -> str: - """Build a concise, filesystem-safe table name from raw first-row header text. - - Pipeline: - 1. Split by common delimiters (' | ', '_br_'/'__br_', '\\n') - 2. Strip whitespace, deduplicate (preserve order) - 3. Drop trivial single-character tokens (single CJK char, single digit, - single letter) — reuses ``_is_meaningful_token`` from shared text_utils - 4. Rejoin with spaces and cap at ``_MAX_TABLE_NAME_CHARS`` - - Args: - raw_header_text: The raw first-row text, often pipe-separated. - - Returns: - A cleaned string suitable for use in a filename (may be empty if all - fields were trivial). - """ - from shared.utils.text_utils import _is_meaningful_token - - if not raw_header_text: - return "" - - # 1. Split on common header delimiters - parts = re.split(r"\s*\|\s*|_+br_|\n", raw_header_text) - - # 2. Strip + deduplicate (order-preserved) - seen: set[str] = set() - unique: list[str] = [] - for p in parts: - p = p.strip() - if not p or p in seen: - continue - seen.add(p) - unique.append(p) - - # 3. Keep only meaningful fields (drop single-char noise) - meaningful = [f for f in unique if _is_meaningful_token(f)] - - # 4. Join and enforce length cap - result = " ".join(meaningful) - if len(result) > _MAX_TABLE_NAME_CHARS: - result = result[:_MAX_TABLE_NAME_CHARS].rstrip() - return result - - -g_tbl_lock = threading.Lock() - -# ============================================================================ -# PRECISION MODE: Excel Header Detection with Merge Cell Metadata -# ============================================================================ - - -def _get_merged_cell_value(ws, row: int, col: int, merged_ranges: list): - """ - Get the value of a cell, accounting for merged cell regions. - For merged cells, returns the value from the top-left corner of the merge range. - - Args: - ws: openpyxl worksheet - row: 1-indexed row number - col: 1-indexed column number - merged_ranges: list of merged cell ranges from ws.merged_cells.ranges - - Returns: - The cell value (from merge origin if applicable) - """ - for mr in merged_ranges: - if mr.min_row <= row <= mr.max_row and mr.min_col <= col <= mr.max_col: - # This cell is part of a merged region, get value from top-left - return ws.cell(mr.min_row, mr.min_col).value - # Not a merged cell, return direct value - return ws.cell(row, col).value - - -# ============================================================================ -# NEW: Enhanced Header Detection with Row/Column MultiIndex Support -# ============================================================================ - -# Data types that indicate a cell is data, not header (parameterized for future extension) -DATA_TYPES_TO_EXCLUDE = (int, float, datetime.datetime) - - -def _get_unique_cells_in_row( - ws, row: int, col_range: Tuple[int, int], merged_ranges: list -) -> List[dict]: - """Get all unique cells in a row, treating merged cells as single cells. - - Returns: List of {col_start, col_end, value, is_merged} - """ - c_start, c_end = col_range - cells = [] - visited_cols = set() - - for col in range(c_start, c_end + 1): - if col in visited_cols: - continue - - in_merge = False - for mr in merged_ranges: - if mr.min_row <= row <= mr.max_row and mr.min_col <= col <= mr.max_col: - val = ws.cell(mr.min_row, mr.min_col).value - merge_col_end = min(mr.max_col, c_end) - - for mc in range(mr.min_col, merge_col_end + 1): - visited_cols.add(mc) - - cells.append( - { - "col_start": mr.min_col, - "col_end": merge_col_end, - "value": val, - "is_merged": True, - } - ) - in_merge = True - break - - if not in_merge: - val = ws.cell(row, col).value - cells.append( - {"col_start": col, "col_end": col, "value": val, "is_merged": False} - ) - visited_cols.add(col) - - return cells - - -def _get_unique_cells_in_col( - ws, col: int, row_range: Tuple[int, int], merged_ranges: list -) -> List[dict]: - """Get all unique cells in a column, treating merged cells as single cells.""" - r_start, r_end = row_range - cells = [] - visited_rows = set() - - for row in range(r_start, r_end + 1): - if row in visited_rows: - continue - - in_merge = False - for mr in merged_ranges: - if mr.min_row <= row <= mr.max_row and mr.min_col <= col <= mr.max_col: - val = ws.cell(mr.min_row, mr.min_col).value - merge_row_end = min(mr.max_row, r_end) - - for mr_row in range(mr.min_row, merge_row_end + 1): - visited_rows.add(mr_row) - - cells.append( - { - "row_start": mr.min_row, - "row_end": merge_row_end, - "value": val, - "is_merged": True, - } - ) - in_merge = True - break - - if not in_merge: - val = ws.cell(row, col).value - cells.append( - {"row_start": row, "row_end": row, "value": val, "is_merged": False} - ) - visited_rows.add(row) - - return cells - - -def _is_candidate_header_row( - ws, - row: int, - col_range: Tuple[int, int], - merged_ranges: list, - exclude_types: tuple = DATA_TYPES_TO_EXCLUDE, -) -> bool: - """Check if a row is a candidate header row. - - Logic: Row is a candidate if all cells are text (no numbers/dates). - Merged cells are treated as single cells. - """ - cells = _get_unique_cells_in_row(ws, row, col_range, merged_ranges) - - has_any_value = False - for cell in cells: - val = cell["value"] - if val is None: - continue - has_any_value = True - - if isinstance(val, bool): - continue - if isinstance(val, exclude_types): - return False - - return has_any_value - - -def _is_candidate_header_col( - ws, - col: int, - row_range: Tuple[int, int], - merged_ranges: list, - exclude_types: tuple = DATA_TYPES_TO_EXCLUDE, -) -> bool: - """Check if a column is a candidate header column (for row index).""" - cells = _get_unique_cells_in_col(ws, col, row_range, merged_ranges) - - has_any_value = False - for cell in cells: - val = cell["value"] - if val is None: - continue - has_any_value = True - - if isinstance(val, bool): - continue - if isinstance(val, exclude_types): - return False - - return has_any_value - - -def _detect_header_regions( - ws, row_range: Tuple[int, int], col_range: Tuple[int, int], merged_ranges: list -) -> Tuple[List[int], List[int]]: - """Detect header rows and columns. - - Scans rows first, then scans columns only in the data region (excluding header rows). - This prevents header row content from influencing column header detection. - - Returns: - header_rows: List of candidate header row numbers (1-indexed) - header_cols: List of candidate header column numbers (1-indexed) - """ - r_start, r_end = row_range - c_start, c_end = col_range - - # Scan for candidate header rows (top to bottom) - header_rows = [] - for row in range(r_start, r_end + 1): - if _is_candidate_header_row(ws, row, col_range, merged_ranges): - header_rows.append(row) - else: - break - - # Determine data region (excluding header rows) - data_row_start = header_rows[-1] + 1 if header_rows else r_start - - # Skip column scanning if no data rows remain - if data_row_start > r_end: - return header_rows, [] - - # Scan for candidate header columns (left to right) - only in data region - header_cols = [] - data_row_range = (data_row_start, r_end) - for col in range(c_start, c_end + 1): - if _is_candidate_header_col(ws, col, data_row_range, merged_ranges): - header_cols.append(col) - else: - break - - return header_rows, header_cols - - -def _build_column_multiindex( - ws, header_rows: List[int], col_range: Tuple[int, int], merged_ranges: list -) -> Union[pd.Index, pd.MultiIndex]: - """Build column MultiIndex from header rows.""" - c_start, c_end = col_range - levels = [] - - for row in header_rows: - row_values = [] - for col in range(c_start, c_end + 1): - val = _get_merged_cell_value(ws, row, col, merged_ranges) - row_values.append(str(val).strip() if val else "") - levels.append(row_values) - - # Forward fill for merged cells - for idx, level in enumerate(levels): - filled = [] - last = "" - for val in level: - if val: - last = val - filled.append(last if last else val) - levels[idx] = filled - - if len(levels) == 1: - return pd.Index(levels[0]) - return pd.MultiIndex.from_arrays(levels) - - -def _build_row_multiindex( - ws, - header_cols: List[int], - row_range: Tuple[int, int], - merged_ranges: list, - header_rows: List[int] = None, -) -> Union[pd.Index, pd.MultiIndex]: - """Build row MultiIndex from header columns. - - Args: - header_rows: If provided, use the last header row's values as index names - """ - r_start, r_end = row_range - levels = [] - names = [] - - for col in header_cols: - col_values = [] - for row in range(r_start, r_end + 1): - val = _get_merged_cell_value(ws, row, col, merged_ranges) - col_values.append(str(val).strip() if val else "") - levels.append(col_values) - - # Get the column name from the last header row - if header_rows: - name_row = header_rows[-1] - name_val = _get_merged_cell_value(ws, name_row, col, merged_ranges) - names.append(str(name_val).strip() if name_val else None) - else: - names.append(None) - - # Forward fill for merged cells - for idx, level in enumerate(levels): - filled = [] - last = "" - for val in level: - if val: - last = val - filled.append(last if last else val) - levels[idx] = filled - - if len(levels) == 1: - idx = pd.Index(levels[0]) - idx.name = names[0] if names else None - return idx - return pd.MultiIndex.from_arrays(levels, names=names) - - -def _parse_subtable( - ws, row_range: Tuple[int, int], col_range: Tuple[int, int], merged_ranges: list -) -> dict: - """Parse a subtable with new header detection logic. - - Returns: - dict with keys: df, header_rows, header_cols, fallback_col_header, fallback_row_header - """ - r_start, r_end = row_range - c_start, c_end = col_range - - header_rows, header_cols = _detect_header_regions( - ws, row_range, col_range, merged_ranges - ) - - total_rows = r_end - r_start + 1 - total_cols = c_end - c_start + 1 - - # Fall-back check: if all rows/cols are headers, treat as no-header - fallback_col_header = len(header_rows) == total_rows - fallback_row_header = len(header_cols) == total_cols - - # Determine data region - if fallback_col_header: - data_row_start = r_start - columns = None - else: - data_row_start = header_rows[-1] + 1 if header_rows else r_start - columns = ( - _build_column_multiindex(ws, header_rows, col_range, merged_ranges) - if header_rows - else None - ) - - if fallback_row_header: - data_col_start = c_start - row_index = None - else: - data_col_start = header_cols[-1] + 1 if header_cols else c_start - row_index = ( - _build_row_multiindex( - ws, header_cols, (data_row_start, r_end), merged_ranges, header_rows - ) - if header_cols - else None - ) - - # Read data - data = [] - for row in range(data_row_start, r_end + 1): - row_data = [] - for col in range(data_col_start, c_end + 1): - val = _get_merged_cell_value(ws, row, col, merged_ranges) - row_data.append(val) - data.append(row_data) - - # Adjust column index if there are row index columns - if columns is not None and header_cols and not fallback_row_header: - if isinstance(columns, pd.MultiIndex): - columns = columns[len(header_cols) :] - else: - columns = columns[len(header_cols) :] - - df = pd.DataFrame(data, columns=columns, index=row_index) - - # Append original Excel row numbers as the last column for cross-referencing - excel_row_numbers = list(range(data_row_start, r_end + 1)) - if isinstance(df.columns, pd.MultiIndex): - n_levels = df.columns.nlevels - src_row_key = tuple(["_src_row"] + [""] * (n_levels - 1)) - df[src_row_key] = excel_row_numbers - else: - df["_src_row"] = excel_row_numbers - - return { - "df": df, - "header_rows": header_rows if not fallback_col_header else [], - "header_cols": header_cols if not fallback_row_header else [], - "fallback_col_header": fallback_col_header, - "fallback_row_header": fallback_row_header, - } - - -# ============================================================================ -# Sheet Splitting: Detect true separators and split into subtables -# ============================================================================ - - -def _find_effective_range( - ws, row_range: Tuple[int, int], col_range: Tuple[int, int] -) -> Tuple[Tuple[int, int], Tuple[int, int]]: - """Find the effective (non-empty) row and column ranges within a region.""" - r_start, r_end = row_range - c_start, c_end = col_range - - eff_r_start, eff_r_end = None, None - eff_c_start, eff_c_end = None, None - - for row in range(r_start, r_end + 1): - for col in range(c_start, c_end + 1): - if ws.cell(row, col).value is not None: - if eff_r_start is None: - eff_r_start = row - eff_r_end = row - if eff_c_start is None or col < eff_c_start: - eff_c_start = col - if eff_c_end is None or col > eff_c_end: - eff_c_end = col - - if eff_r_start is None: - return ((r_start, r_start), (c_start, c_start)) - - return ((eff_r_start, eff_r_end), (eff_c_start, eff_c_end)) - - -def _is_true_separator_row( - ws, row: int, effective_col_range: Tuple[int, int], merged_ranges: list = None -) -> bool: - """Check if a row is a true separator (all empty within effective column range). - - Considers merged cells - a cell is not empty if it's part of any merged range. - """ - c_start, c_end = effective_col_range - merged_ranges = merged_ranges or [] - - for col in range(c_start, c_end + 1): - # Check if cell has a value - if ws.cell(row, col).value is not None: - return False - # Check if cell is part of a merged range - for mr in merged_ranges: - if mr.min_row <= row <= mr.max_row and mr.min_col <= col <= mr.max_col: - return False # Part of a merge, not truly empty - return True - - -def _is_true_separator_col( - ws, col: int, effective_row_range: Tuple[int, int], merged_ranges: list = None -) -> bool: - """Check if a column is a true separator (all empty within effective row range). - - Considers merged cells - a cell is not empty if it's part of any merged range. - """ - r_start, r_end = effective_row_range - merged_ranges = merged_ranges or [] - - for row in range(r_start, r_end + 1): - # Check if cell has a value - if ws.cell(row, col).value is not None: - return False - # Check if cell is part of a merged range - for mr in merged_ranges: - if mr.min_row <= row <= mr.max_row and mr.min_col <= col <= mr.max_col: - return False # Part of a merge, not truly empty - return True - - -def _find_separator_groups(items: List[int]) -> List[List[int]]: - """Group consecutive separator items together.""" - if not items: - return [] - - groups = [] - current_group = [items[0]] - - for i in range(1, len(items)): - if items[i] == items[i - 1] + 1: - current_group.append(items[i]) - else: - groups.append(current_group) - current_group = [items[i]] - - groups.append(current_group) - return groups - - -def _split_sheet_recursive( - ws, - row_range: Tuple[int, int], - col_range: Tuple[int, int], - merged_ranges: list = None, - min_rows: int = 2, - min_cols: int = 2, -) -> List[Tuple[Tuple[int, int], Tuple[int, int]]]: - """ - Recursively split a sheet region into subtables based on true separators. - - Args: - merged_ranges: List of merged cell ranges to consider when detecting separators - - Returns list of (row_range, col_range) tuples for each subtable. - """ - r_start, r_end = row_range - c_start, c_end = col_range - merged_ranges = merged_ranges or [] - - # Find effective range (trim empty edges) - (eff_r_start, eff_r_end), (eff_c_start, eff_c_end) = _find_effective_range( - ws, row_range, col_range - ) - - # If region is too small or empty, return as-is or empty - if eff_r_end - eff_r_start + 1 < min_rows or eff_c_end - eff_c_start + 1 < min_cols: - if eff_r_start is not None: - return [((eff_r_start, eff_r_end), (eff_c_start, eff_c_end))] - return [] - - # Find true separator rows (considering merged cells) - separator_rows = [] - for row in range(eff_r_start + 1, eff_r_end): - if _is_true_separator_row(ws, row, (eff_c_start, eff_c_end), merged_ranges): - separator_rows.append(row) - - # Find true separator columns (considering merged cells) - separator_cols = [] - for col in range(eff_c_start + 1, eff_c_end): - if _is_true_separator_col(ws, col, (eff_r_start, eff_r_end), merged_ranges): - separator_cols.append(col) - - # Group consecutive separators - row_groups = _find_separator_groups(separator_rows) - col_groups = _find_separator_groups(separator_cols) - - # Choose split direction - do_row_split = len(row_groups) > 0 and ( - len(col_groups) == 0 or len(row_groups) <= len(col_groups) - ) - do_col_split = len(col_groups) > 0 and not do_row_split - - if do_row_split: - subtables = [] - prev_end = eff_r_start - for group in row_groups: - if group[0] > prev_end: - sub_result = _split_sheet_recursive( - ws, - (prev_end, group[0] - 1), - (eff_c_start, eff_c_end), - merged_ranges, - min_rows, - min_cols, - ) - subtables.extend(sub_result) - prev_end = group[-1] + 1 - if prev_end <= eff_r_end: - sub_result = _split_sheet_recursive( - ws, - (prev_end, eff_r_end), - (eff_c_start, eff_c_end), - merged_ranges, - min_rows, - min_cols, - ) - subtables.extend(sub_result) - return subtables - - elif do_col_split: - subtables = [] - prev_end = eff_c_start - for group in col_groups: - if group[0] > prev_end: - sub_result = _split_sheet_recursive( - ws, - (eff_r_start, eff_r_end), - (prev_end, group[0] - 1), - merged_ranges, - min_rows, - min_cols, - ) - subtables.extend(sub_result) - prev_end = group[-1] + 1 - if prev_end <= eff_c_end: - sub_result = _split_sheet_recursive( - ws, - (eff_r_start, eff_r_end), - (prev_end, eff_c_end), - merged_ranges, - min_rows, - min_cols, - ) - subtables.extend(sub_result) - return subtables - - else: - return [((eff_r_start, eff_r_end), (eff_c_start, eff_c_end))] - - -# ============================================================================ -# Post-split Merge: Absorb small fragments into nearest neighbor -# ============================================================================ - - -def _count_non_empty_cells( - ws, row_range: Tuple[int, int], col_range: Tuple[int, int] -) -> int: - """Count non-empty cells in a region.""" - count = 0 - for r in range(row_range[0], row_range[1] + 1): - for c in range(col_range[0], col_range[1] + 1): - if ws.cell(r, c).value is not None: - count += 1 - return count - - -def _merge_small_subtables( - ws, subtables: List[Tuple[Tuple[int, int], Tuple[int, int]]], min_cells: int = 4 -) -> List[Tuple[Tuple[int, int], Tuple[int, int]]]: - """Merge subtables that have too few non-empty cells into their nearest neighbor. - - This is a post-processing step after _split_sheet_recursive to prevent - over-fragmentation. Fragments with fewer than min_cells non-empty cells - are iteratively absorbed into the nearest neighbor subtable (by bounding-box - distance), expanding the neighbor's bounding box to encompass both regions. - - Args: - ws: openpyxl worksheet - subtables: list of (row_range, col_range) tuples from _split_sheet_recursive - min_cells: minimum non-empty cells for a subtable to be kept standalone - - Returns: - Merged list of (row_range, col_range) tuples - """ - if len(subtables) <= 1: - return subtables - - # Build working list with cell counts - items = [] - for rr, cr in subtables: - count = _count_non_empty_cells(ws, rr, cr) - items.append({"rr": rr, "cr": cr, "cells": count}) - - # Iteratively merge the smallest sub-threshold fragment - changed = True - while changed and len(items) > 1: - changed = False - - # Find the smallest fragment below threshold - min_idx = None - for i, item in enumerate(items): - if item["cells"] < min_cells: - if min_idx is None or item["cells"] < items[min_idx]["cells"]: - min_idx = i - - if min_idx is None: - break # All subtables are above threshold - - # Find nearest neighbor by bounding-box gap distance - src = items[min_idx] - best_j = None - best_dist = float("inf") - for j, tgt in enumerate(items): - if j == min_idx: - continue - row_gap = max( - 0, tgt["rr"][0] - src["rr"][1] - 1, src["rr"][0] - tgt["rr"][1] - 1 - ) - col_gap = max( - 0, tgt["cr"][0] - src["cr"][1] - 1, src["cr"][0] - tgt["cr"][1] - 1 - ) - dist = row_gap + col_gap - if dist < best_dist or ( - dist == best_dist and tgt["cells"] > items[best_j]["cells"] - ): - best_dist = dist - best_j = j - - if best_j is None: - break # Should not happen when len(items) > 1 - - # Merge: expand neighbor's bounding box to encompass both - tgt = items[best_j] - merged_rr = (min(src["rr"][0], tgt["rr"][0]), max(src["rr"][1], tgt["rr"][1])) - merged_cr = (min(src["cr"][0], tgt["cr"][0]), max(src["cr"][1], tgt["cr"][1])) - items[best_j] = { - "rr": merged_rr, - "cr": merged_cr, - "cells": src["cells"] + tgt["cells"], - } - - logger.debug( - f"Merged small fragment (rows={src['rr']}, cols={src['cr']}, " - f"cells={src['cells']}) into neighbor (rows={tgt['rr']}, cols={tgt['cr']})" - ) - - del items[min_idx] - changed = True - - return [(item["rr"], item["cr"]) for item in items] - - -def parse_headers_from_excel( - file_source: Union[str, io.BytesIO], - sheet_name: Optional[str] = None, - split_subtables: bool = True, - include_hidden_sheets: bool = False, -) -> Dict[str, pd.DataFrame]: - """ - Parse Excel file using openpyxl to accurately detect headers via merged cell metadata. - - This is the PRECISION MODE for Excel parsing - it uses the actual merge cell - information from the Excel file to build correct MultiIndex headers without - relying on LLM or heuristics. - - Args: - file_source: Path to Excel file or BytesIO stream - sheet_name: Specific sheet to parse (None = all sheets) - split_subtables: If True, split sheets into subtables based on empty row/column separators (default: True) - include_hidden_sheets: If True, parse hidden/very-hidden sheets. Default False (skip them). - - Returns: - Dictionary mapping sheet/subtable names to DataFrames with correctly set headers - When split_subtables=True, keys are like 'SheetName', 'SheetName_2', 'SheetName_3' etc. - """ - try: - # Load workbook with data_only=True to get calculated values - if isinstance(file_source, str): - wb = openpyxl.load_workbook(file_source, data_only=True) - else: - # BytesIO stream - file_source.seek(0) # Ensure we're at the start - wb = openpyxl.load_workbook(file_source, data_only=True) - - results = {} - sheets_to_parse = [sheet_name] if sheet_name else wb.sheetnames - - for sn in sheets_to_parse: - if sn not in wb.sheetnames: - logger.warning(f"Sheet '{sn}' not found in workbook, skipping") - continue - - ws = wb[sn] - - # Skip hidden sheets unless explicitly included - if not include_hidden_sheets and ws.sheet_state != "visible": - logger.info( - f"Sheet '{sn}' is hidden (state={ws.sheet_state}), skipping" - ) - continue - - # Skip empty sheets - if ws.max_row is None or ws.max_row == 0: - logger.debug(f"Sheet '{sn}' is empty, skipping") - continue - - # Get merged cell ranges - merged_ranges = list(ws.merged_cells.ranges) - logger.debug(f"Sheet '{sn}': found {len(merged_ranges)} merged cell ranges") - - if split_subtables: - # Split sheet into subtables (considers merged cells) - subtable_regions = _split_sheet_recursive( - ws, (1, ws.max_row), (1, ws.max_column or 1), merged_ranges - ) - # Merge back small fragments to prevent over-fragmentation - before_count = len(subtable_regions) - subtable_regions = _merge_small_subtables(ws, subtable_regions) - if len(subtable_regions) != before_count: - logger.info( - f"Sheet '{sn}': merged {before_count} subtables → {len(subtable_regions)} " - f"(absorbed {before_count - len(subtable_regions)} small fragments)" - ) - logger.debug( - f"Sheet '{sn}': {len(subtable_regions)} subtables after merge" - ) - - for idx, (row_range, col_range) in enumerate(subtable_regions): - result = _parse_subtable(ws, row_range, col_range, merged_ranges) - df = result["df"] - - # Store header_cols count in DataFrame attrs for later use in HTML rendering - df.attrs["row_header_cols"] = len(result["header_cols"]) - - # Generate unique key for each subtable - if idx == 0: - key = sn - else: - key = f"{sn}_{idx + 1}" - - logger.debug( - f"Subtable '{key}': rows={row_range}, cols={col_range}, " - f"header_rows={result['header_rows']}, header_cols={result['header_cols']}" - ) - - results[key] = df - else: - # Treat entire sheet as one subtable - row_range = (1, ws.max_row) - col_range = (1, ws.max_column or 1) - - result = _parse_subtable(ws, row_range, col_range, merged_ranges) - df = result["df"] - - # Store header_cols count in DataFrame attrs for later use in HTML rendering - df.attrs["row_header_cols"] = len(result["header_cols"]) - - logger.debug( - f"Sheet '{sn}': header_rows={result['header_rows']}, " - f"header_cols={result['header_cols']}, " - f"fallback_col={result['fallback_col_header']}, " - f"fallback_row={result['fallback_row_header']}" - ) - - results[sn] = df - - wb.close() - return results - - except Exception as e: - logger.error(f"Error parsing Excel with precision mode: {e}") - raise TableParsingException( - user_message="Failed to parse Excel file headers", - reason="EXCEL_PRECISION_PARSE_FAILED", - internal_message=str(e), - original_exception=e, - ) - - -def identify_tables(line): - """Identify if a line contains a table. - - Note: For HTML tables, use merge_html_tables() from html_parser.py - to preprocess multi-line tables before calling this function. - """ - # HTML table: complete
...
in one line - html_tb_pattern = r".*?" - tables = re.findall(html_tb_pattern, line, re.DOTALL) - if bool(tables): - return True, "html", tables - - # MD table: lines starting and ending with | - if line.startswith("|") and line.endswith("|"): - return True, "md", [] - - return False, None, None - - -def df2md(tb_df: pd.DataFrame, *, index: bool = False, na_rep: str = "—") -> str: - """Convert DataFrame to Markdown table format with dynamic column widths. - - Note: Truncation should be done externally using truncate_text before calling this function. - - Args: - tb_df: Input DataFrame - index: Whether to include index column - na_rep: String to represent NA values - - Returns: - Markdown table string - """ - import unicodedata - - def get_display_width(text: str) -> int: - """eval width for both ASCII and Chinese""" - width = 0 - for char in text: - if unicodedata.east_asian_width(char) in ("F", "W"): - width += 2 - else: - width += 1 - return width - - def pad_to_width(text: str, target_width: int) -> str: - current_width = get_display_width(text) - padding = target_width - current_width - return text + " " * max(0, padding) - - df = tb_df.copy() - - # Handle index - if index: - df = df.reset_index() - - # Replace NA values - df = df.fillna(na_rep) - - # Convert all values to string - df = df.astype(str) - - # Calculate column widths based on actual display width (no truncation) - col_widths = {} - for col in df.columns: - header_width = get_display_width(str(col)) - max_content_width = max(df[col].apply(get_display_width)) if len(df) > 0 else 0 - col_widths[col] = max(header_width, max_content_width) - - # Build header row - header_cells = [pad_to_width(str(col), col_widths[col]) for col in df.columns] - header_line = "| " + " | ".join(header_cells) + " |" - - # Build separator row - separator_cells = ["-" * col_widths[col] for col in df.columns] - separator_line = "|-" + "-|-".join(separator_cells) + "-|" - - # Build data rows - data_lines = [] - for _, row in df.iterrows(): - cells = [pad_to_width(str(row[col]), col_widths[col]) for col in df.columns] - data_lines.append("| " + " | ".join(cells) + " |") - - # Combine all parts - lines = [header_line, separator_line] + data_lines - return "\n".join(lines) - - -def clean_html_tb(html: str) -> str: - soup = BeautifulSoup(html, "html.parser") - for row in soup.find_all("tr"): - seen = set() - unique_cells = [] - for cell in row.find_all("td", recursive=False): - content = cell.encode_contents() - if content not in seen: - seen.add(content) - unique_cells.append(cell) - row.clear() - for cell in unique_cells: - row.append(cell) - return soup.prettify() - - -def extract_tables_by_forms(tb_txt, form): - if form == "html": - return tb_txt - elif form == "md": - tb_df = pd.read_table( - pd.io.common.StringIO(tb_txt), sep="|", engine="python", on_bad_lines="skip" - ) - tb_df = tb_df.drop(columns=tb_df.columns[0]) # Drop extra leading column - tb_df = tb_df.drop(columns=tb_df.columns[-1]) # Drop extra trailing column - tb_df.columns = tb_df.columns.str.strip() # Clean up headers - # Filter out MD separator lines (e.g. "---", ":---:", "---:") - separator_pattern = r"^[\s\-:]+$" - tb_df = tb_df[ - ~tb_df.apply( - lambda row: row.astype(str).str.match(separator_pattern).all(), axis=1 - ) - ] - tb_strs = tb_df.to_html(index=False) - else: - tb_strs = None # UNDER DEVELOPMENT other forms of tables... - return tb_strs - - -def parse_headers(df_temp, paras=None, header_window=5, smart_headers=True): - def parse_headers_nonsmart(df_): - non_na_row = df_[df_.notna().any(axis=1)].head(1) - header_id = non_na_row.index[-1] if not non_na_row.empty else None - header_rows = list(range(header_id + 1)) - return header_rows - - if not pd.isna( - df_temp.columns - ).all(): # If columns are not all NaN, no need to add extra row - df_temp.loc[-1] = df_temp.columns - df_temp.index = df_temp.index + 1 - df_temp = df_temp.sort_index() - df_temp.columns = [np.nan] * df_temp.shape[1] - - if paras["summary_table"] and smart_headers: - try: - tb_small = df_temp.head(header_window) - tb_small_str = df2html(tb_small) - prompt, temperature, top_p, max_tokens = build_prompt( - task="detect-table-headers", texts=tb_small_str, query="", paras=paras - ) - - messages = [ - {"role": "system", "content": "You are a helpful assistant"}, - {"role": "user", "content": prompt}, - ] - - ctx_task_id = gen_str_codes((str(uuid.uuid4()) + tb_small_str)) - - # Track task status via Redis (skip in LOCAL_DEBUG mode) - import os - - if os.getenv("LOCAL_DEBUG", "0") != "1": - from shared.services.redis.redis_sync_service import ( - SyncRedisServiceFactory, - ) - - redis_service = SyncRedisServiceFactory.get_service() - redis_service.set(f"task:{ctx_task_id}:status", "processing", ttl=7200) - - # Use unified AI service - header_res = get_openai_client().chat_completion( - messages=messages, timeout=60 - ) - header_res = eval_response(header_res) - # Extract answer field - if isinstance(header_res, dict): - answer = header_res.get("answer", []) - else: - answer = header_res if isinstance(header_res, list) else [] - - # Check if answer is empty list - if not answer or len(answer) == 0: - logger.warning( - "AI returned empty list, cannot identify headers, falling back to traditional mode..." - ) - header_rows = parse_headers_nonsmart(df_temp) - else: - try: - header_id = answer[-1] - header_rows = list(range(header_id + 1)) - except Exception as e: - logger.warning( - f"Failed to parse header row number: {e}, falling back to traditional mode..." - ) - header_rows = parse_headers_nonsmart(df_temp) - - except Exception as e: - logger.warning( - f"Smart header parsing failed: {e}, falling back to traditional mode..." - ) - header_rows = parse_headers_nonsmart(df_temp) - else: - header_rows = parse_headers_nonsmart(df_temp) - - # improve table structure based on header rows - if len(header_rows) == 0 or (all(h is None for h in header_rows)): - logger.warning("No valid headers detected, fallback to using row 0 as header") - new_header = df_temp.iloc[0].ffill().bfill().tolist() - df_temp.columns = new_header - df_temp = df_temp.iloc[1:].reset_index(drop=True) - return df_temp - elif len(header_rows) > 1: - head_lst = [] - for i in range(0, len(header_rows)): - temp_lst = df_temp.iloc[i].ffill().bfill().tolist() - head_lst.append(temp_lst) - new_header = pd.MultiIndex.from_arrays(np.array(head_lst)) - else: - new_header = df_temp.iloc[header_rows[-1]].ffill().bfill().tolist() - - df_temp.columns = new_header - df_temp = df_temp.iloc[(header_rows[-1]) + 1 :] - df_temp = df_temp.reset_index(drop=True) - return df_temp - - -def parse_tb_keywords( - tb_df, kw_spit=">>>" -): # Extract keywords from headers (can also add LLM extraction) - def parse_single_level_(cols, keywords): - cols = [str(c) for c in cols] - for col in cols: - if kw_spit in col: - tmp_kw = col.split(">>>")[0] - else: # May be first occurrence - tmp_kw = col - if tmp_kw not in keywords: - keywords.append(col) - keywords_a_level = list(set([k.strip() for k in keywords])) - return keywords_a_level - - tb_keywords = [] - if isinstance(tb_df.columns, pd.MultiIndex): - multi_cols = tb_df.columns - cols_df = pd.DataFrame( - multi_cols.tolist(), - columns=[f"level_{i}" for i in range(multi_cols.nlevels)], - ) - for i in range(multi_cols.nlevels): # Extract each level as list - level_kws = [] - level_kws = parse_single_level_(cols_df[f"level_{i}"].tolist(), level_kws) - tb_keywords.extend(level_kws) - else: - tb_keywords = parse_single_level_(tb_df.columns, tb_keywords) - - # Remove duplicates while preserving column order - tb_keywords = remove_duplicates_orderkept(tb_keywords) - tb_keywords = [ - k - for k in tb_keywords - if isinstance(k, str) - and k.strip() - and k.strip() != "nan" - and "Unnamed" not in k - ] - return ";".join(tb_keywords) - - -def parse_tb_contents( - df_temp, parent_dic=None, file_name="", sheet_name="", row_header_cols=0 -): - """Parse table contents and generate HTML. - - Args: - row_header_cols: Number of leftmost columns that are row headers (will be rendered as ) - """ - if parent_dic is None: - parent_dic = {} - - tb_res = df_temp.fillna("").infer_objects(copy=False) - tb_strs = df2html(tb_res, row_header_cols=row_header_cols) - - tb_tree = tb_columns_to_tree(df_temp, parent_dic, file_name, sheet_name) - tb_paths = flatten_dic2paths(tb_tree) - return tb_paths, tb_strs - - -def tb_columns_to_tree(df, parent_dic, file_name, sheet_name): - if isinstance(df.columns, pd.MultiIndex): - # Convert MultiIndex columns to a nested dictionary (tree-like structure) - columns = pd.DataFrame(df.columns.tolist()) - for level in range(columns.shape[1]): - columns[level] = process_duplicate_cols(columns[level]) - - new_columns = pd.MultiIndex.from_frame(columns) - tree_structure = multiindex_to_tree(new_columns) - else: - # If columns are not MultiIndex, convert them to a dictionary with empty dictionaries as values - new_columns = process_duplicate_cols(df.columns) - tree_structure = {col: {} for col in new_columns} - - df.columns = new_columns - if (not file_name == "") and (not sheet_name == ""): - parent_dic[file_name][sheet_name] = tree_structure - elif not sheet_name == "": - parent_dic[sheet_name] = tree_structure - elif not file_name == "": - parent_dic[file_name] = tree_structure - else: - parent_dic = tree_structure - return parent_dic - - -def multiindex_to_tree(multiindex): - """Convert a MultiIndex to a tree-like nested dictionary structure.""" - - def tree(): - return OrderedDict() - - root = tree() - for keys in multiindex: - current_level = root - for key in keys: - if key not in current_level: - current_level[key] = tree() - current_level = current_level[key] - - def convert_to_dict(d): - if isinstance(d, OrderedDict): - d = {k: convert_to_dict(v) for k, v in d.items()} - return d - - return convert_to_dict(root) - - -def postprocess_tb(df, drop=False): - if drop: - # Track if index was originally a simple RangeIndex (no semantic meaning) - # dropna(how='all') can turn RangeIndex into Int64Index by introducing gaps, - # which would incorrectly trigger the "preserve row index" logic below. - was_range_index = isinstance(df.index, pd.RangeIndex) - - # Drop rows where all data columns are empty (exclude _src_row from the check) - # _src_row is always non-null, so including it would prevent any row from being dropped. - src_row_cols = [ - c - for c in df.columns - if (isinstance(c, tuple) and c[0] == "_src_row") or c == "_src_row" - ] - if src_row_cols: - data_cols = [c for c in df.columns if c not in src_row_cols] - mask = df[data_cols].isna().all(axis=1) - df = df[~mask] - else: - df = df.dropna(how="all") - - # If index was originally RangeIndex, re-number it to avoid gaps - if was_range_index: - df = df.reset_index(drop=True) - - # Drop columns that are all empty AND have no meaningful header - # A column with a valid header should be preserved even if data is empty - cols_to_drop = [] - for col_idx, col in enumerate(df.columns): - # Check if all data values are NaN - # Use iloc to avoid ambiguity when MultiIndex has duplicate tuple keys - if df.iloc[:, col_idx].isna().all(): - # Check if the column header is meaningful - # For MultiIndex: check if any level has a non-empty meaningful value - # For simple index: check if the header is not None/empty - has_meaningful_header = False - if isinstance(col, tuple): - # MultiIndex column - check if any level has meaningful content - for level in col: - if ( - level - and str(level).strip() - and str(level).strip() not in ["None", "nan", "NaN"] - ): - has_meaningful_header = True - break - else: - # Simple column name - if ( - col - and str(col).strip() - and str(col).strip() not in ["None", "nan", "NaN"] - ): - has_meaningful_header = True - - # Only drop if header is not meaningful - if not has_meaningful_header: - cols_to_drop.append(col_idx) - - if cols_to_drop: - # Use positional indices to drop columns safely (avoids duplicate MultiIndex key issues) - cols_to_keep = [i for i in range(len(df.columns)) if i not in cols_to_drop] - df = df.iloc[:, cols_to_keep] - - logger.debug(f"Dropped {len(cols_to_drop)} empty columns") - - # Preserve meaningful row index (header columns) as regular columns - # Only drop=True if it's a simple RangeIndex (no semantic meaning) - if not isinstance(df.index, pd.RangeIndex): - # Remember if columns were MultiIndex before reset - was_multiindex = isinstance(df.columns, pd.MultiIndex) - n_levels = df.columns.nlevels if was_multiindex else 1 - - # Avoid name collision before reset_index(). - # Two collision sources: - # A) An index level name, when padded into a tuple by pandas, - # matches an existing column. - # B) Multiple index levels share the same name → pandas tries - # to insert duplicate columns (e.g. five levels all named - # one merged header repeated across five padded columns. - # Strategy: de-duplicate index.names so every level gets a unique - # column name during reset_index, then clean up afterwards. - existing_col_set = set(df.columns) - - def _make_padded(name): - """Simulate the column name pandas would create for this index level.""" - if was_multiindex: - return (name,) + ("",) * (n_levels - 1) - return name - - if isinstance(df.index, pd.MultiIndex): - seen_counts = {} # name → how many times seen so far - deduped = [] - for n in df.index.names: - if n is None: - deduped.append(None) - continue - padded = _make_padded(n) - # Collision with existing column OR with a previously-seen index name - if padded in existing_col_set or n in seen_counts: - deduped.append( - None - ) # let pandas auto-name it (level_0, level_1 …) - else: - deduped.append(n) - seen_counts[n] = seen_counts.get(n, 0) + 1 - df.index.names = deduped - elif hasattr(df.index, "name") and df.index.name is not None: - padded = _make_padded(df.index.name) - if padded in existing_col_set: - df.index.name = None - - df = df.reset_index() # Converts index to columns - - # Clean up auto-generated column names like 'index', 'level_0', 'level_1' - # For MultiIndex columns, we need to preserve the structure - if was_multiindex: - # Build new column tuples for the index columns - new_cols = [] - for col in df.columns: - if isinstance(col, str) and ( - col.startswith("level_") or col == "index" - ): - # Create a tuple with empty strings to match MultiIndex levels - new_cols.append(tuple([""] * n_levels)) - else: - new_cols.append(col) - df.columns = pd.MultiIndex.from_tuples(new_cols) - else: - # For simple columns - new_cols = [] - for col in df.columns: - if isinstance(col, str) and ( - col.startswith("level_") or col == "index" - ): - new_cols.append("") - else: - new_cols.append(col) - df.columns = new_cols - else: - df.reset_index(drop=True, inplace=True) - - # Clean column names - preserve MultiIndex structure if present - if isinstance(df.columns, pd.MultiIndex): - # For MultiIndex, clean each level's values while preserving structure - new_levels = [] - for level_idx in range(df.columns.nlevels): - level_vals = df.columns.get_level_values(level_idx) - cleaned = [ - str(v).replace("\n", "") if v is not None else "" for v in level_vals - ] - new_levels.append(cleaned) - df.columns = pd.MultiIndex.from_arrays(new_levels, names=df.columns.names) - # Also handle 'Unnamed' in MultiIndex - new_levels = [] - for level_idx in range(df.columns.nlevels): - level_vals = df.columns.get_level_values(level_idx) - cleaned = [np.nan if "Unnamed" in str(v) else v for v in level_vals] - new_levels.append(cleaned) - df.columns = pd.MultiIndex.from_arrays(new_levels, names=df.columns.names) - else: - df.columns = [ - str(col).replace("\n", "") for col in df.columns - ] # Replace '\n' in column headers - df.columns = [ - np.nan if "Unnamed" in str(col) else col for col in df.columns - ] # Replace Unnamed with nan - df = df.map( - lambda x: x.replace("\n", "") if isinstance(x, str) else x - ) # Replace '\n' in each cell - df = process_datetime_cells(df) - return df - - -def process_datetime_cells(df): - df = df.copy() - - def convert(x): - if isinstance(x, (pd.Timestamp, datetime.datetime)): - return x.strftime("%Y-%m-%d %H:%M:%S") - return x - - return df.apply(lambda col: col.map(convert)) - - -def process_duplicate_cols(columns): - col_count = {} - new_columns = [] - for col in columns: - if col in col_count: - new_columns.append(f"{col}>>>{col_count[col]}") - col_count[col] += 1 - else: - new_columns.append(col) - col_count[col] = 1 - return new_columns - - -def format_tb_scope(df, num): - if len(df) > int(num * 3 + 1): - # Get head and tail rows - head_df = df.head(num) - tail_df = df.tail(num) - # Middle portion excluding head and tail - middle_df = df.iloc[num : len(df) - num] - - if len(middle_df) >= num: - mid_sample_df = middle_df.sample(n=num, random_state=42) - else: # If middle has less than num rows, take all - mid_sample_df = middle_df - scope_df = pd.concat(objs=[head_df, mid_sample_df, tail_df], ignore_index=True) - else: - scope_df = df - scope_df = scope_df.applymap(lambda x: str(x).strip() if pd.notnull(x) else x) - scope_str = df2html(scope_df) - return scope_str - - -def parse_xlsx( - file_path, - file_name, - output_dir, - baseurl, - base_llm_paras=None, - window_h=10, - relative_root=None, - use_precision_mode=True, - include_hidden_sheets=False, -): - """ - Parse Excel file and extract table content. - - Args: - file_path: Path or URL to the Excel file - file_name: Display name for the file - output_dir: Directory to save extracted tables - baseurl: Base URL for file loading - base_llm_paras: LLM parameters for summarization - window_h: Window size for table scope - relative_root: Root path for relative paths - use_precision_mode: If True, use openpyxl merged cell metadata for accurate - header detection. If False, use LLM/heuristic mode. - Default is True for better accuracy. - include_hidden_sheets: If True, parse hidden/very-hidden sheets. Default False. - - Returns: - DataFrame with parsed table information - """ - time_stamp = get_str_time() - df_list = [] - - table_data = load_file_bytes(file_path, file_url=baseurl) - table_stream = io.BytesIO(table_data) - - tb_dir = os.path.join(output_dir, "tables") - os.makedirs(tb_dir, exist_ok=True) - all_tb_paths = [] - exist_sheets = [] - - if use_precision_mode: - # PRECISION MODE: Use openpyxl metadata for accurate header detection - logger.info("Using precision mode for Excel header detection") - try: - sheets_dict = parse_headers_from_excel( - table_stream, include_hidden_sheets=include_hidden_sheets - ) - precision_mode_active = True - except Exception as e: - logger.warning(f"Precision mode failed, falling back to legacy mode: {e}") - table_stream.seek(0) # Reset stream position - sheets_dict = pd.read_excel(table_stream, sheet_name=None) - precision_mode_active = False - else: - # LEGACY MODE: Use pandas read_excel + LLM/heuristic header detection - sheets_dict = pd.read_excel(table_stream, sheet_name=None) - precision_mode_active = False - - all_sheets = sheets_dict.items() - - for sheet_name, sheet_content in all_sheets: - sheet_name = sheet_name.strip() - if sheet_name in exist_sheets: - sheet_name = sheet_name + str(len(exist_sheets)) - else: - exist_sheets.append(sheet_name) - - sheet_tbs = [sheet_content] - for tb in sheet_tbs: - try: - tb = postprocess_tb(tb, drop=True) - if len(tb) == 0 or tb.empty or tb.isna().all().all(): - continue - - # In precision mode, headers are already correctly set by parse_headers_from_excel - # In legacy mode, use LLM/heuristic header parsing - if not precision_mode_active: - tb = parse_headers(tb, paras=base_llm_paras) - - # Drop _src_row column before converting to HTML/keywords - # (_src_row is a debug column added by _parse_subtable for cross-referencing) - src_row_cols = [ - c - for c in tb.columns - if (isinstance(c, tuple) and c[0] == "_src_row") or c == "_src_row" - ] - if src_row_cols: - tb = tb.drop(columns=src_row_cols) - - # Get row header column count from DataFrame attrs (set in parse_headers_from_excel) - row_header_cols = tb.attrs.get("row_header_cols", 0) - - tb_paths, tb_strs = parse_tb_contents( - tb, - parent_dic={file_name: {sheet_name: {}}}, - file_name=file_name, - sheet_name=sheet_name, - row_header_cols=row_header_cols, - ) - - # Unified LLM extraction: title + keywords + summary in one call - # (consistent with doc_parser.py and md_parser.py) - llm_title = None - llm_summary = None - tb_keywords = "" - if base_llm_paras["summary_table"]: - from app.services.document_parser.txt_parser import ( - extract_title_keywords_summary, - ) - - llm_title, tb_keywords, llm_summary = ( - extract_title_keywords_summary(tb_strs, max_keywords=3) - ) - - # Build tb_summary: table index + optional LLM summary - table_index = f"table-{sheet_name}" - if llm_summary: - tb_summary = f"{table_index}\n{llm_summary}" - else: - # Fallback: use mechanical column keywords when LLM is off - tb_keywords_fallback = parse_tb_keywords(tb) - tb_summary = table_index - tb_keywords = tb_keywords if tb_keywords else tb_keywords_fallback - - # Use a filesystem-safe filename so LLM titles like "A/B" do not - # accidentally create nested paths under tables/. - effective_name = llm_title if llm_title else sheet_name - tb_name = ( - path_handle( - remove_spaces("table-" + effective_name), mode="clean_single" - ) - + ".html" - ) - tb_path = os.path.join(tb_dir, tb_name) - soup = BeautifulSoup(tb_strs, features="html.parser") - tb_html_str = soup.prettify() - with open(tb_path, "w", encoding="utf-8") as f: - f.write(tb_html_str) - - # Use same temp_uid for both marker and know_id (aligned with doc_parser/md_parser) - temp_uid = gen_str_codes(tb_strs + str(sheet_name)) - relative_tb_path = f"tables/{tb_name}" - tb_ref = build_chunk_ref(relative_tb_path) - tb_bottom_content = f"{tb_ref}\nTable summary:\n{tb_summary}\nMain columns:\n{tb_keywords}" - - bottom_tokens = tokenize2stw_remove( - [tb_bottom_content], base_llm_paras["stopwords"] - ) - - all_tb_paths.extend(tb_paths) - # Use relative path for tables: "tables/xxx.html" - df_list.append( - [ - tb_bottom_content, - relative_tb_path, - "table", - len(tb_strs), - tb_keywords, - tb_summary, - temp_uid, - bottom_tokens, - "", - time_stamp, - "", - ] - ) - - except KnowhereException: - raise - except Exception as e: - logger.error(f"Table parsing failed: {e}") - raise TableParsingException( - user_message="Failed to parse Excel table content", - reason="TABLE_PROCESSING_FAILED", - internal_message=str(e), - original_exception=e, - ) - - table_df = pd.DataFrame(df_list, columns=settings.ALL_DF_COLS.split(",")) - table_df = process_dup_paths_df(table_df) - return table_df diff --git a/apps/worker/app/services/document_parser/tables/dataframe_helpers.py b/apps/worker/app/services/document_parser/tables/dataframe_helpers.py new file mode 100644 index 000000000..cef9b6446 --- /dev/null +++ b/apps/worker/app/services/document_parser/tables/dataframe_helpers.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +import os +from typing import Any + +import pandas as pd + +from app.services.document_parser.support.identifiers import gen_str_codes + + +def flatten_list(nested_list: list[Any]) -> list[Any]: + """Flatten a nested list.""" + flat_list: list[Any] = [] + for item in nested_list: + if isinstance(item, list): + flat_list.extend(flatten_list(item)) + else: + flat_list.append(item) + return flat_list + + +def merge_df(input_df: pd.DataFrame) -> pd.DataFrame: + """Merge DataFrame rows that share the same path.""" + dfs_by_path = list(input_df.groupby("path", sort=False)) + processed_dfs: list[pd.DataFrame] = [] + + for key, df in dfs_by_path: + content_to_merge: list[str] = [] + types_to_merge: list[str] = [] + total_length = 0 + + for _, row in df.iterrows(): + content_to_merge.append(str(row["content"])) + types_to_merge.extend(str(row["type"]).split("\n")) + total_length += len(str(row["content"])) + + processed_dfs.append( + pd.DataFrame( + [ + { + "content": "\n".join(content_to_merge), + "type": "\n".join(sorted(set(types_to_merge))), + "path": key, + "length": total_length, + "know_id": gen_str_codes("\n".join(content_to_merge)), + } + ] + ) + ) + + return pd.concat(processed_dfs, axis=0, ignore_index=True) + + +def process_dup_paths_df(df: pd.DataFrame) -> pd.DataFrame: + """De-duplicate KB DataFrame paths for final output.""" + if "path" not in df.columns: + return df + + split_char = os.getenv("SPLIT_CHAR", "/") + dup_mask = df["path"].duplicated(keep=False) + if not dup_mask.any(): + return df + + path_occurrences: dict[str, list[int]] = {} + for idx, path in enumerate(df["path"]): + path_occurrences.setdefault(str(path), []).append(idx) + + path_renames: dict[int, str] = {} + parent_rename_map: dict[str, dict[int, str]] = {} + + for path, indices in path_occurrences.items(): + if len(indices) > 1: + parent_rename_map[path] = {} + for occurrence, idx in enumerate(indices): + if occurrence == 0: + path_renames[idx] = path + else: + new_path = f"{path}_{occurrence + 1}" + path_renames[idx] = new_path + parent_rename_map[path][idx] = new_path + + new_paths: list[str] = [] + for idx, row in df.iterrows(): + row_index = int(str(idx)) + path = str(row["path"]) + new_path = path_renames.get(row_index, path) + path_parts = new_path.split(split_char) + + for parent_path, rename_info in parent_rename_map.items(): + parent_parts = parent_path.split(split_char) + if len(path_parts) > len(parent_parts) and path_parts[: len(parent_parts)] == parent_parts: + matching_parent_idx = None + for parent_idx in sorted(rename_info.keys(), reverse=True): + if parent_idx < row_index: + matching_parent_idx = parent_idx + break + if matching_parent_idx is not None: + renamed_parent = rename_info[matching_parent_idx] + renamed_parent_parts = renamed_parent.split(split_char) + new_path = split_char.join( + renamed_parent_parts + path_parts[len(parent_parts) :] + ) + break + + new_paths.append(new_path) + + df = df.copy() + df["path"] = new_paths + return df diff --git a/apps/worker/app/services/document_parser/tables/dataframe_html_renderer.py b/apps/worker/app/services/document_parser/tables/dataframe_html_renderer.py new file mode 100644 index 000000000..97da217ce --- /dev/null +++ b/apps/worker/app/services/document_parser/tables/dataframe_html_renderer.py @@ -0,0 +1,294 @@ +from __future__ import annotations + +from typing import List, Union + +import pandas as pd + + +def render_multiindex_thead(columns: pd.MultiIndex, escape: bool = False) -> str: + """Convert MultiIndex columns to an HTML thead with colspan and rowspan.""" + import html as html_lib + + level_count = columns.nlevels + column_count = len(columns) + + grid = [] + for level in range(level_count): + row = [columns.get_level_values(level)[col] for col in range(column_count)] + grid.append(row) + + colspan = [[1] * column_count for _ in range(level_count)] + + for level in range(level_count): + col = 0 + while col < column_count: + span = 1 + while col + span < column_count and grid[level][col] == grid[level][col + span]: + parent_match = True + for parent_level in range(level): + if grid[parent_level][col] != grid[parent_level][col + span]: + parent_match = False + break + if parent_match: + span += 1 + else: + break + colspan[level][col] = span + col += span + + rowspan = [[1] * column_count for _ in range(level_count)] + + for col in range(column_count): + level = 0 + while level < level_count: + span = 1 + while level + span < level_count: + if ( + grid[level][col] == grid[level + span][col] + and colspan[level][col] == colspan[level + span][col] + ): + span += 1 + else: + break + rowspan[level][col] = span + level += span + + covered = [[False] * column_count for _ in range(level_count)] + html_parts = [""] + + for level in range(level_count): + html_parts.append('') + col = 0 + while col < column_count: + if covered[level][col]: + col += 1 + continue + + val = grid[level][col] + val_str = str(val) if val is not None else "" + if escape: + val_str = html_lib.escape(val_str) + + column_span = colspan[level][col] + row_span = rowspan[level][col] + + for row_offset in range(row_span): + for column_offset in range(column_span): + if row_offset > 0 or column_offset > 0: + if ( + level + row_offset < level_count + and col + column_offset < column_count + ): + covered[level + row_offset][col + column_offset] = True + + attrs = [] + if column_span > 1: + attrs.append(f'colspan="{column_span}"') + if row_span > 1: + attrs.append(f'rowspan="{row_span}"') + + attr_str = " " + " ".join(attrs) if attrs else "" + html_parts.append(f"{val_str}") + + col += column_span + + html_parts.append("") + + html_parts.append("") + return "".join(html_parts) + + +def render_tbody_with_row_headers( + tb_df: pd.DataFrame, + row_header_cols: int = 0, + na_rep: str = "—", + escape: bool = False, +) -> str: + """Render a DataFrame body with optional row-header cells and merged spans.""" + import html as html_lib + + if row_header_cols <= 0: + html_parts = [""] + for _, row in tb_df.iterrows(): + html_parts.append("") + for val in row: + val_str = na_rep if pd.isna(val) else str(val) + if escape: + val_str = html_lib.escape(val_str) + html_parts.append(f"{val_str}") + html_parts.append("") + html_parts.append("") + return "".join(html_parts) + + row_count = len(tb_df) + column_count = len(tb_df.columns) + + if row_count == 0: + return "" + + grid = [] + for row_idx in range(row_count): + row_values = [] + for col_idx in range(row_header_cols): + val = tb_df.iloc[row_idx, col_idx] + val = na_rep if pd.isna(val) else str(val) + row_values.append(val) + grid.append(row_values) + + colspan = [[1] * row_header_cols for _ in range(row_count)] + + for row_idx in range(row_count): + col_idx = 0 + while col_idx < row_header_cols: + span = 1 + while ( + col_idx + span < row_header_cols + and grid[row_idx][col_idx] == grid[row_idx][col_idx + span] + ): + span += 1 + colspan[row_idx][col_idx] = span + col_idx += span + + rowspan = [[1] * row_header_cols for _ in range(row_count)] + + col_idx = 0 + while col_idx < row_header_cols: + row_idx = 0 + while row_idx < row_count: + if col_idx > 0 and grid[row_idx][col_idx] == grid[row_idx][col_idx - 1]: + row_idx += 1 + continue + + current_colspan = colspan[row_idx][col_idx] + span = 1 + + while row_idx + span < row_count: + if grid[row_idx][col_idx] != grid[row_idx + span][col_idx]: + break + if colspan[row_idx + span][col_idx] != current_colspan: + break + + parent_match = True + for parent_col in range(col_idx): + if grid[row_idx][parent_col] != grid[row_idx + span][parent_col]: + parent_match = False + break + if parent_match: + span += 1 + else: + break + + rowspan[row_idx][col_idx] = span + row_idx += span + col_idx += 1 + + covered = [[False] * row_header_cols for _ in range(row_count)] + + for row_idx in range(row_count): + col_idx = 0 + while col_idx < row_header_cols: + column_span = colspan[row_idx][col_idx] + for offset in range(1, column_span): + if col_idx + offset < row_header_cols: + covered[row_idx][col_idx + offset] = True + col_idx += column_span + + for row_idx in range(row_count): + for col_idx in range(row_header_cols): + if covered[row_idx][col_idx]: + continue + row_span = rowspan[row_idx][col_idx] + for offset in range(1, row_span): + if row_idx + offset < row_count: + column_span = colspan[row_idx][col_idx] + for column_offset in range(column_span): + if col_idx + column_offset < row_header_cols: + covered[row_idx + offset][col_idx + column_offset] = True + + html_parts = [""] + + for row_idx in range(row_count): + html_parts.append("") + + for col_idx in range(row_header_cols): + if covered[row_idx][col_idx]: + continue + + val_str = grid[row_idx][col_idx] + if escape: + val_str = html_lib.escape(val_str) + + row_span = rowspan[row_idx][col_idx] + column_span = colspan[row_idx][col_idx] + + attrs = [] + if row_span > 1: + attrs.append(f'rowspan="{row_span}"') + if column_span > 1: + attrs.append(f'colspan="{column_span}"') + + attr_str = " " + " ".join(attrs) if attrs else "" + html_parts.append(f'{val_str}') + + for col_idx in range(row_header_cols, column_count): + val = tb_df.iloc[row_idx, col_idx] + val_str = na_rep if pd.isna(val) else str(val) + if escape: + val_str = html_lib.escape(val_str) + html_parts.append(f"{val_str}") + + html_parts.append("") + + html_parts.append("") + return "".join(html_parts) + + +def df2html( + tb_df: pd.DataFrame, + *, + index: bool = False, + classes: Union[str, List[str], None] = "table table-striped", + na_rep: str = "—", + escape: bool = False, + row_header_cols: int = 0, +) -> str: + """Convert a DataFrame to an HTML table.""" + class_str = ( + classes if isinstance(classes, str) else " ".join(classes) if classes else "" + ) + + if isinstance(tb_df.columns, pd.MultiIndex): + thead_html = render_multiindex_thead(tb_df.columns, escape=escape) + tbody_html = render_tbody_with_row_headers( + tb_df, row_header_cols, na_rep, escape + ) + return f'{thead_html}{tbody_html}
' + + if row_header_cols <= 0: + table_html = tb_df.to_html( + index=index, + na_rep=na_rep, + classes=classes, + escape=escape, + border=0, + justify="center", + ) + return table_html.replace("\n", "") + + html_parts = [f''] + html_parts.append("") + html_parts.append('') + for col in tb_df.columns: + col_str = str(col) if col is not None else "" + if escape: + import html as html_lib + + col_str = html_lib.escape(col_str) + html_parts.append(f"") + html_parts.append("") + html_parts.append("") + + tbody_html = render_tbody_with_row_headers(tb_df, row_header_cols, na_rep, escape) + html_parts.append(tbody_html) + html_parts.append("
{col_str}
") + return "".join(html_parts) diff --git a/apps/worker/app/services/document_parser/tables/table_asset_writer.py b/apps/worker/app/services/document_parser/tables/table_asset_writer.py new file mode 100644 index 000000000..679ee028e --- /dev/null +++ b/apps/worker/app/services/document_parser/tables/table_asset_writer.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +import os +from dataclasses import dataclass + +from app.services.document_parser.support.parser_rows import ParsedRow + + +@dataclass(frozen=True) +class TableAssetInput: + html: str + output_dir: str + table_name: str + summary: str + keywords: str + know_id: str + addtime: str + page_nums: str = "" + content: str | None = None + tokens: str = "" + length: int | None = None + + +def write_table_asset(table_input: TableAssetInput) -> ParsedRow: + table_dir = os.path.join(table_input.output_dir, "tables") + os.makedirs(table_dir, exist_ok=True) + table_filename = _ensure_html_extension(table_input.table_name) + table_path = os.path.join(table_dir, table_filename) + with open(table_path, "w", encoding="utf-8") as table_file: + table_file.write(table_input.html) + row_content = table_input.content if table_input.content is not None else table_input.html + return ParsedRow( + content=row_content, + path=f"tables/{table_filename}", + type="table", + keywords=table_input.keywords, + summary=table_input.summary, + know_id=table_input.know_id, + tokens=table_input.tokens, + connectto="", + addtime=table_input.addtime, + page_nums=table_input.page_nums, + length=table_input.length, + ) + + +def _ensure_html_extension(table_name: str) -> str: + return table_name if table_name.endswith(".html") else f"{table_name}.html" diff --git a/apps/worker/app/services/document_parser/tables/table_frame_parser.py b/apps/worker/app/services/document_parser/tables/table_frame_parser.py new file mode 100644 index 000000000..76d7965dd --- /dev/null +++ b/apps/worker/app/services/document_parser/tables/table_frame_parser.py @@ -0,0 +1,422 @@ +# pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportOptionalOperand=false, reportOptionalSubscript=false, reportReturnType=false, reportOperatorIssue=false, reportIndexIssue=false, reportAssignmentType=false, reportGeneralTypeIssues=false +from __future__ import annotations + +import datetime +import os +import uuid +from collections import OrderedDict + +import numpy as np +import pandas as pd +from app.services.document_parser.tables.dataframe_html_renderer import df2html +from app.services.document_parser.support.identifiers import gen_str_codes +from app.services.document_parser.support.path_helpers import flatten_dic2paths +from loguru import logger + +from shared.services.ai.prompt_service import build_prompt +from shared.services.ai.response_process_service import eval_response +from shared.services.ai.openai_compatible_client_sync import get_openai_client +from shared.utils.text_utils import remove_duplicates_orderkept + + +def parse_headers( + table_frame: pd.DataFrame, + paras: dict[str, object] | None = None, + header_window: int = 5, + smart_headers: bool = True, +) -> pd.DataFrame: + llm_parameters = paras or {"summary_table": False} + + def parse_headers_nonsmart(candidate_frame: pd.DataFrame) -> list[int]: + non_na_row = candidate_frame[candidate_frame.notna().any(axis=1)].head(1) + header_id = non_na_row.index[-1] if not non_na_row.empty else None + return list(range(header_id + 1)) + + if not pd.isna(table_frame.columns).all(): + table_frame.loc[-1] = table_frame.columns + table_frame.index = table_frame.index + 1 + table_frame = table_frame.sort_index() + table_frame.columns = [np.nan] * table_frame.shape[1] + + if llm_parameters["summary_table"] and smart_headers: + try: + table_sample = table_frame.head(header_window) + table_sample_html = df2html(table_sample) + prompt, _temperature, _top_p, _max_tokens = build_prompt( + task="detect-table-headers", + texts=table_sample_html, + query="", + paras=llm_parameters, + ) + + messages = [ + {"role": "system", "content": "You are a helpful assistant"}, + {"role": "user", "content": prompt}, + ] + + context_task_id = gen_str_codes((str(uuid.uuid4()) + table_sample_html)) + + if os.getenv("LOCAL_DEBUG", "0") != "1": + from shared.services.redis.redis_sync_service import ( + SyncRedisServiceFactory, + ) + + redis_service = SyncRedisServiceFactory.get_service() + redis_service.set( + f"task:{context_task_id}:status", + "processing", + ttl=7200, + ) + + header_response = get_openai_client().chat_completion( + messages=messages, + timeout=60, + ) + parsed_response = eval_response(header_response) + if isinstance(parsed_response, dict): + answer = parsed_response.get("answer", []) + else: + answer = parsed_response if isinstance(parsed_response, list) else [] + + if not answer or len(answer) == 0: + logger.warning( + "AI returned empty list, cannot identify headers, falling back to traditional mode..." + ) + header_rows = parse_headers_nonsmart(table_frame) + else: + try: + header_id = answer[-1] + header_rows = list(range(header_id + 1)) + except Exception as exc: + logger.warning( + f"Failed to parse header row number: {exc}, falling back to traditional mode..." + ) + header_rows = parse_headers_nonsmart(table_frame) + + except Exception as exc: + logger.warning( + f"Smart header parsing failed: {exc}, falling back to traditional mode..." + ) + header_rows = parse_headers_nonsmart(table_frame) + else: + header_rows = parse_headers_nonsmart(table_frame) + + if len(header_rows) == 0 or (all(header is None for header in header_rows)): + logger.warning("No valid headers detected, fallback to using row 0 as header") + new_header = table_frame.iloc[0].ffill().bfill().tolist() + table_frame.columns = new_header + return table_frame.iloc[1:].reset_index(drop=True) + + if len(header_rows) > 1: + header_levels = [] + for header_index in range(0, len(header_rows)): + header_level = table_frame.iloc[header_index].ffill().bfill().tolist() + header_levels.append(header_level) + new_header = pd.MultiIndex.from_arrays(np.array(header_levels)) + else: + new_header = table_frame.iloc[header_rows[-1]].ffill().bfill().tolist() + + table_frame.columns = new_header + table_frame = table_frame.iloc[(header_rows[-1]) + 1 :] + return table_frame.reset_index(drop=True) + + +def parse_tb_keywords(table_frame: pd.DataFrame, kw_spit: str = ">>>") -> str: + def parse_single_level(columns: list[object], keywords: list[str]) -> list[str]: + column_texts = [str(column) for column in columns] + for column_text in column_texts: + if kw_spit in column_text: + keyword = column_text.split(">>>")[0] + else: + keyword = column_text + if keyword not in keywords: + keywords.append(column_text) + return list({keyword.strip() for keyword in keywords}) + + table_keywords: list[str] = [] + if isinstance(table_frame.columns, pd.MultiIndex): + multi_columns = table_frame.columns + columns_frame = pd.DataFrame( + multi_columns.tolist(), + columns=[f"level_{i}" for i in range(multi_columns.nlevels)], + ) + for level_index in range(multi_columns.nlevels): + level_keywords: list[str] = [] + level_keywords = parse_single_level( + columns_frame[f"level_{level_index}"].tolist(), + level_keywords, + ) + table_keywords.extend(level_keywords) + else: + table_keywords = parse_single_level(table_frame.columns.tolist(), table_keywords) + + table_keywords = remove_duplicates_orderkept(table_keywords) + table_keywords = [ + keyword + for keyword in table_keywords + if isinstance(keyword, str) + and keyword.strip() + and keyword.strip() != "nan" + and "Unnamed" not in keyword + ] + return ";".join(table_keywords) + + +def parse_tb_contents( + table_frame: pd.DataFrame, + parent_dic: dict[str, object] | None = None, + file_name: str = "", + sheet_name: str = "", + row_header_cols: int = 0, +) -> tuple[list[str], str]: + if parent_dic is None: + parent_dic = {} + + rendered_frame = table_frame.fillna("").infer_objects(copy=False) + table_html = df2html(rendered_frame, row_header_cols=row_header_cols) + + table_tree = tb_columns_to_tree(table_frame, parent_dic, file_name, sheet_name) + table_paths = flatten_dic2paths(table_tree) + return table_paths, table_html + + +def tb_columns_to_tree( + table_frame: pd.DataFrame, + parent_dic: dict[str, object], + file_name: str, + sheet_name: str, +) -> dict[str, object]: + if isinstance(table_frame.columns, pd.MultiIndex): + columns = pd.DataFrame(table_frame.columns.tolist()) + for level in range(columns.shape[1]): + columns[level] = process_duplicate_cols(columns[level]) + + new_columns = pd.MultiIndex.from_frame(columns) + tree_structure = multiindex_to_tree(new_columns) + else: + new_columns = process_duplicate_cols(table_frame.columns) + tree_structure = {column: {} for column in new_columns} + + table_frame.columns = new_columns + if (not file_name == "") and (not sheet_name == ""): + parent_dic[file_name][sheet_name] = tree_structure + elif not sheet_name == "": + parent_dic[sheet_name] = tree_structure + elif not file_name == "": + parent_dic[file_name] = tree_structure + else: + parent_dic = tree_structure + return parent_dic + + +def multiindex_to_tree(multiindex: pd.MultiIndex) -> dict[object, object]: + def tree() -> OrderedDict[object, object]: + return OrderedDict() + + root = tree() + for keys in multiindex: + current_level = root + for key in keys: + if key not in current_level: + current_level[key] = tree() + current_level = current_level[key] + + def convert_to_dict(value: object) -> object: + if isinstance(value, OrderedDict): + return {key: convert_to_dict(child) for key, child in value.items()} + return value + + return dict(convert_to_dict(root)) + + +def postprocess_tb(table_frame: pd.DataFrame, drop: bool = False) -> pd.DataFrame: + if drop: + was_range_index = isinstance(table_frame.index, pd.RangeIndex) + + source_row_columns = [ + column + for column in table_frame.columns + if (isinstance(column, tuple) and column[0] == "_src_row") + or column == "_src_row" + ] + if source_row_columns: + data_columns = [ + column for column in table_frame.columns if column not in source_row_columns + ] + mask = table_frame[data_columns].isna().all(axis=1) + table_frame = table_frame[~mask] + else: + table_frame = table_frame.dropna(how="all") + + if was_range_index: + table_frame = table_frame.reset_index(drop=True) + + cols_to_drop: list[int] = [] + for column_index, column in enumerate(table_frame.columns): + if table_frame.iloc[:, column_index].isna().all(): + has_meaningful_header = False + if isinstance(column, tuple): + for level in column: + if ( + level + and str(level).strip() + and str(level).strip() not in ["None", "nan", "NaN"] + ): + has_meaningful_header = True + break + elif ( + column + and str(column).strip() + and str(column).strip() not in ["None", "nan", "NaN"] + ): + has_meaningful_header = True + + if not has_meaningful_header: + cols_to_drop.append(column_index) + + if cols_to_drop: + cols_to_keep = [ + index + for index in range(len(table_frame.columns)) + if index not in cols_to_drop + ] + table_frame = table_frame.iloc[:, cols_to_keep] + + logger.debug(f"Dropped {len(cols_to_drop)} empty columns") + + if not isinstance(table_frame.index, pd.RangeIndex): + was_multiindex = isinstance(table_frame.columns, pd.MultiIndex) + column_level_count = table_frame.columns.nlevels if was_multiindex else 1 + existing_column_set = set(table_frame.columns) + + def make_padded(name: object) -> object: + if was_multiindex: + return (name,) + ("",) * (column_level_count - 1) + return name + + if isinstance(table_frame.index, pd.MultiIndex): + seen_counts: dict[object, int] = {} + deduped_names: list[object | None] = [] + for name in table_frame.index.names: + if name is None: + deduped_names.append(None) + continue + padded = make_padded(name) + if padded in existing_column_set or name in seen_counts: + deduped_names.append(None) + else: + deduped_names.append(name) + seen_counts[name] = seen_counts.get(name, 0) + 1 + table_frame.index.names = deduped_names + elif hasattr(table_frame.index, "name") and table_frame.index.name is not None: + padded = make_padded(table_frame.index.name) + if padded in existing_column_set: + table_frame.index.name = None + + table_frame = table_frame.reset_index() + + if was_multiindex: + new_columns = [] + for column in table_frame.columns: + if isinstance(column, str) and ( + column.startswith("level_") or column == "index" + ): + new_columns.append(tuple([""] * column_level_count)) + else: + new_columns.append(column) + table_frame.columns = pd.MultiIndex.from_tuples(new_columns) + else: + new_columns = [] + for column in table_frame.columns: + if isinstance(column, str) and ( + column.startswith("level_") or column == "index" + ): + new_columns.append("") + else: + new_columns.append(column) + table_frame.columns = new_columns + else: + table_frame.reset_index(drop=True, inplace=True) + + if isinstance(table_frame.columns, pd.MultiIndex): + new_levels = [] + for level_index in range(table_frame.columns.nlevels): + level_values = table_frame.columns.get_level_values(level_index) + cleaned = [ + str(value).replace("\n", "") if value is not None else "" + for value in level_values + ] + new_levels.append(cleaned) + table_frame.columns = pd.MultiIndex.from_arrays( + new_levels, + names=table_frame.columns.names, + ) + + new_levels = [] + for level_index in range(table_frame.columns.nlevels): + level_values = table_frame.columns.get_level_values(level_index) + cleaned = [np.nan if "Unnamed" in str(value) else value for value in level_values] + new_levels.append(cleaned) + table_frame.columns = pd.MultiIndex.from_arrays( + new_levels, + names=table_frame.columns.names, + ) + else: + table_frame.columns = [ + str(column).replace("\n", "") for column in table_frame.columns + ] + table_frame.columns = [ + np.nan if "Unnamed" in str(column) else column + for column in table_frame.columns + ] + + table_frame = table_frame.map( + lambda value: value.replace("\n", "") if isinstance(value, str) else value + ) + return process_datetime_cells(table_frame) + + +def process_datetime_cells(table_frame: pd.DataFrame) -> pd.DataFrame: + table_frame = table_frame.copy() + + def convert(value: object) -> object: + if isinstance(value, (pd.Timestamp, datetime.datetime)): + return value.strftime("%Y-%m-%d %H:%M:%S") + return value + + return table_frame.apply(lambda column: column.map(convert)) + + +def process_duplicate_cols(columns: object) -> list[object]: + column_counts: dict[object, int] = {} + new_columns: list[object] = [] + for column in columns: + if column in column_counts: + new_columns.append(f"{column}>>>{column_counts[column]}") + column_counts[column] += 1 + else: + new_columns.append(column) + column_counts[column] = 1 + return new_columns + + +def format_tb_scope(table_frame: pd.DataFrame, num: int) -> str: + if len(table_frame) > int(num * 3 + 1): + head_frame = table_frame.head(num) + tail_frame = table_frame.tail(num) + middle_frame = table_frame.iloc[num : len(table_frame) - num] + + if len(middle_frame) >= num: + mid_sample_frame = middle_frame.sample(n=num, random_state=42) + else: + mid_sample_frame = middle_frame + scope_frame = pd.concat( + objs=[head_frame, mid_sample_frame, tail_frame], + ignore_index=True, + ) + else: + scope_frame = table_frame + scope_frame = scope_frame.map( + lambda value: str(value).strip() if pd.notnull(value) else value + ) + return df2html(scope_frame) diff --git a/apps/worker/app/services/document_parser/tables/table_parser.py b/apps/worker/app/services/document_parser/tables/table_parser.py new file mode 100755 index 000000000..ef69e1bfc --- /dev/null +++ b/apps/worker/app/services/document_parser/tables/table_parser.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +import pandas as pd + + +def parse_xlsx( + file_path: str, + file_name: str, + output_dir: str, + baseurl: str, + base_llm_paras: dict[str, object] | None = None, + window_h: int = 10, + relative_root: str | None = None, + use_precision_mode: bool = True, + include_hidden_sheets: bool = False, +) -> pd.DataFrame: + from app.services.document_parser.formats.excel.table_parser import ( + parse_xlsx as parse_excel_xlsx, + ) + + return parse_excel_xlsx( + file_path=file_path, + file_name=file_name, + output_dir=output_dir, + baseurl=baseurl, + base_llm_paras=base_llm_paras, + window_h=window_h, + relative_root=relative_root, + use_precision_mode=use_precision_mode, + include_hidden_sheets=include_hidden_sheets, + ) diff --git a/apps/worker/app/services/document_parser/tables/table_text_parser.py b/apps/worker/app/services/document_parser/tables/table_text_parser.py new file mode 100644 index 000000000..261bbe653 --- /dev/null +++ b/apps/worker/app/services/document_parser/tables/table_text_parser.py @@ -0,0 +1,149 @@ +# pyright: reportArgumentType=false +from __future__ import annotations + +import io +import re +import unicodedata + +import pandas as pd +from bs4 import BeautifulSoup, Tag + +_MAX_TABLE_NAME_CHARS = 80 + + +def sanitize_table_name_from_header(raw_header_text: str) -> str: + """Build a concise, filesystem-safe table name from raw first-row header text.""" + from shared.utils.text_utils import _is_meaningful_token + + if not raw_header_text: + return "" + + parts = re.split(r"\s*\|\s*|_+br_|\n", raw_header_text) + + seen: set[str] = set() + unique: list[str] = [] + for part in parts: + part = part.strip() + if not part or part in seen: + continue + seen.add(part) + unique.append(part) + + meaningful = [field for field in unique if _is_meaningful_token(field)] + result = " ".join(meaningful) + if len(result) > _MAX_TABLE_NAME_CHARS: + result = result[:_MAX_TABLE_NAME_CHARS].rstrip() + return result + + +def identify_tables(line: str) -> tuple[bool, str | None, list[str] | None]: + """Identify whether one logical Markdown line contains a table.""" + html_table_pattern = r".*?" + tables = re.findall(html_table_pattern, line, re.DOTALL) + if bool(tables): + return True, "html", tables + + if line.startswith("|") and line.endswith("|"): + return True, "md", [] + + return False, None, None + + +def df2md(table_frame: pd.DataFrame, *, index: bool = False, na_rep: str = "—") -> str: + """Convert a DataFrame to a Markdown table while preserving display width.""" + + def get_display_width(text: str) -> int: + width = 0 + for character in text: + if unicodedata.east_asian_width(character) in ("F", "W"): + width += 2 + else: + width += 1 + return width + + def pad_to_width(text: str, target_width: int) -> str: + current_width = get_display_width(text) + padding = target_width - current_width + return text + " " * max(0, padding) + + table_frame = table_frame.copy() + + if index: + table_frame = table_frame.reset_index() + + table_frame = table_frame.fillna(na_rep).astype(str) + + column_widths: dict[object, int] = {} + for column in table_frame.columns: + header_width = get_display_width(str(column)) + max_content_width = ( + max(table_frame[column].apply(get_display_width)) + if len(table_frame) > 0 + else 0 + ) + column_widths[column] = max(header_width, max_content_width) + + header_cells = [ + pad_to_width(str(column), column_widths[column]) + for column in table_frame.columns + ] + header_line = "| " + " | ".join(header_cells) + " |" + + separator_cells = ["-" * column_widths[column] for column in table_frame.columns] + separator_line = "|-" + "-|-".join(separator_cells) + "-|" + + data_lines: list[str] = [] + for _, row in table_frame.iterrows(): + cells = [ + pad_to_width(str(row[column]), column_widths[column]) + for column in table_frame.columns + ] + data_lines.append("| " + " | ".join(cells) + " |") + + return "\n".join([header_line, separator_line, *data_lines]) + + +def clean_html_tb(html: str) -> str: + soup = BeautifulSoup(html, "html.parser") + for row in soup.find_all("tr"): + if not isinstance(row, Tag): + continue + seen: set[bytes] = set() + unique_cells: list[Tag] = [] + for cell in row.find_all("td", recursive=False): + if not isinstance(cell, Tag): + continue + content = cell.encode_contents() + if content not in seen: + seen.add(content) + unique_cells.append(cell) + row.clear() + for cell in unique_cells: + row.append(cell) + return str(soup.prettify()) + + +def extract_tables_by_forms(table_text: str, form: str) -> str | None: + if form == "html": + return table_text + + if form != "md": + return None + + table_frame = pd.read_table( + io.StringIO(table_text), + sep="|", + engine="python", + on_bad_lines="skip", + ) + table_frame = table_frame.iloc[:, 1:-1] + table_frame.columns = table_frame.columns.astype(str).str.strip() + + separator_pattern = r"^[\s\-:]+$" + table_frame = table_frame[ + ~table_frame.apply( + lambda row: row.astype(str).str.match(separator_pattern).all(), + axis=1, + ) + ] + return table_frame.to_html(index=False) diff --git a/apps/worker/app/services/storage/__init__.py b/apps/worker/app/services/storage/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/apps/worker/app/services/storage/sync_storage_service.py b/apps/worker/app/services/storage/sync_storage_service.py deleted file mode 100644 index 268c68b90..000000000 --- a/apps/worker/app/services/storage/sync_storage_service.py +++ /dev/null @@ -1,141 +0,0 @@ -""" -Sync storage operations for worker tasks. -Provides S3 file operations and HTTP file downloads using sync adapters -that yield cooperatively under gevent. -""" - -import os -import tempfile -from typing import Any, Dict, Optional - -from loguru import logger - -from shared.core.config import settings -from shared.core.config.storage import get_cached_storage_adapter -from shared.core.exceptions.domain_exceptions import StorageServiceException -from shared.utils.pinned_outbound_http import download_pinned_outbound_file -from shared.utils.url_security import validate_http_url_and_resolve_ip - - -def get_storage_adapter(): - """Get the storage adapter for direct sync S3 operations.""" - return get_cached_storage_adapter() - - -def verify_s3_file_exists(s3_key: str, bucket: Optional[str] = None) -> Dict[str, Any]: - """Verify S3 file exists using sync adapter calls.""" - adapter = get_storage_adapter() - bucket_name = bucket or settings.S3_BUCKET_NAME - try: - if not adapter.exists(s3_key, bucket_name): - return {"exists": False} - size = adapter.get_object_size(s3_key, bucket_name) - return {"exists": True, "size": size} - except Exception as e: - if "404" in str(e) or "not found" in str(e).lower(): - return {"exists": False} - raise StorageServiceException( - internal_message=f"S3 file verification failed: {e}", - operation="verify_s3_file_exists", - original_exception=e, - ) - - -def generate_download_url( - s3_key: str, bucket: Optional[str] = None, expires_in: int = 3600 -) -> Dict[str, Any]: - """Generate presigned download URL using sync adapter.""" - adapter = get_storage_adapter() - bucket_name = bucket or settings.S3_BUCKET_NAME - download_url = adapter.generate_presigned_url( - s3_key, expiration=expires_in, bucket=bucket_name, method="GET" - ) - return {"download_url": download_url, "expires_in": expires_in} - - -def upload_to_s3(local_file_path: str, s3_key: str, bucket: str): - """Upload file to S3 using sync adapter.""" - adapter = get_storage_adapter() - adapter.upload_file(local_file_path, s3_key, bucket) - - -def download_s3_object_to_temp( - s3_key: str, - suffix: str, - temp_dir: str, - bucket: Optional[str] = None, -) -> str: - """Download an object-storage file into a task-local temp file.""" - adapter = get_storage_adapter() - bucket_name = bucket or settings.S3_BUCKET_NAME - local_temp_path: str | None = None - - try: - os.makedirs(temp_dir, exist_ok=True) - with tempfile.NamedTemporaryFile( - delete=False, - suffix=suffix, - dir=temp_dir, - ) as temp_file: - local_temp_path = temp_file.name - adapter.download_file(s3_key, local_temp_path, bucket_name) - return local_temp_path - except Exception as e: - if local_temp_path and os.path.exists(local_temp_path): - os.remove(local_temp_path) - raise StorageServiceException( - internal_message=( - f"Failed to download object-storage file to temp path: " - f"s3_key={s3_key}, temp_dir={temp_dir}, error={e}" - ), - operation="download_s3_object_to_temp", - original_exception=e, - ) from e - - -def upload_zip_result(job_id: str, zip_file_path: str) -> str: - """Upload ZIP result file to S3 and cleanup temp file.""" - results_bucket = getattr(settings, "S3_RESULTS_BUCKET", settings.S3_BUCKET_NAME) - s3_key = f"results/{job_id}.zip" - upload_to_s3(zip_file_path, s3_key, results_bucket) - logger.info(f"Result ZIP uploaded: job_id={job_id}, key={s3_key}") - try: - if os.path.exists(zip_file_path): - os.remove(zip_file_path) - except Exception as e: - logger.warning(f"Failed to cleanup temp ZIP: {e}") - return s3_key - - -def download_file_from_url(file_url: str) -> str: - """Download a URL file through SSRF validation and IP pinning.""" - temp_file_path = "" - try: - validation = validate_http_url_and_resolve_ip(file_url) - if not validation.is_valid or not validation.validated_ip: - raise StorageServiceException( - internal_message=f"Invalid URL: {validation.error_message}", - operation="download_from_url", - ) - - temp_dir = getattr(settings, "TMP_PATH", "/tmp") - os.makedirs(temp_dir, exist_ok=True) - download_result = download_pinned_outbound_file( - url=validation.url, - pinned_ip=validation.validated_ip, - timeout_seconds=300, - user_agent="Knowhere-FileDownloader/1.0", - temp_dir=temp_dir, - ) - temp_file_path = download_result.temp_file_path - return temp_file_path - except StorageServiceException: - raise - except Exception as e: - if os.path.exists(temp_file_path): - os.remove(temp_file_path) - raise StorageServiceException( - internal_message=f"Failed to download file: {e}", - operation="download_from_url", - original_exception=e, - ) diff --git a/apps/worker/app/services/webhook/__init__.py b/apps/worker/app/services/webhook/__init__.py deleted file mode 100644 index 3040aede5..000000000 --- a/apps/worker/app/services/webhook/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Worker webhook services.""" diff --git a/apps/worker/app/services/workload/__init__.py b/apps/worker/app/services/workload/__init__.py index b7df5ee6b..3ad3226c0 100644 --- a/apps/worker/app/services/workload/__init__.py +++ b/apps/worker/app/services/workload/__init__.py @@ -1,7 +1 @@ -""" -Billing services for the worker. -""" - -from .page_estimator import PageEstimator - -__all__ = ["PageEstimator"] +"""Worker workload adapters.""" diff --git a/apps/worker/app/services/workload/page_estimator.py b/apps/worker/app/services/workload/page_estimator.py deleted file mode 100644 index a2dfb3799..000000000 --- a/apps/worker/app/services/workload/page_estimator.py +++ /dev/null @@ -1,213 +0,0 @@ -""" -Page Estimator Service - -Calculates page counts for billing based on: -- PDF: Physical page count from metadata -- PPTX: Slide count -- Text-based (DOC, DOCX, TXT, MD, JSON): Word-based estimation using count_cn_en -- Spreadsheet-based (XLS, XLSX): Row-based estimation - -Supported file types: -- .pdf, .doc, .docx, .pptx, .xls, .xlsx -- .txt, .md, .json, .fragment -- .png, .jpg, .jpeg -""" - -import math -import tempfile -from pathlib import Path - -from shared.core.logging import logger -from shared.utils.text_utils import count_cn_en - -# Constants for page estimation -WORDS_PER_PAGE = 500 # Chinese chars + English words + numbers per page -ROWS_PER_PAGE = 50 # For Excel files - - -class PageEstimator: - """ - Estimates page count for billing purposes. - - Billing Logic: - - PDF: Physical page count from metadata - - PPTX: Slide count - - DOC/DOCX/TXT/MD/JSON: Word-based estimation (count_cn_en / 500) - - XLS/XLSX: Row-based estimation (rows / 50) - - Images: 1 page per image - """ - - @classmethod - def estimate(cls, file_path: str) -> int: - """ - Estimate page count for a file. - - Args: - file_path: Path to the file - - Returns: - Estimated page count (minimum 1) - """ - path = Path(file_path) - suffix = path.suffix.lower() - - try: - if suffix == ".pdf": - return cls._estimate_pdf(file_path) - elif suffix == ".pptx": - return cls._estimate_pptx(file_path) - elif suffix == ".doc": - return cls._estimate_doc(file_path) - elif suffix == ".docx": - return cls._estimate_docx(file_path) - elif suffix == ".xls": - return cls._estimate_xls(file_path) - elif suffix == ".xlsx": - return cls._estimate_xlsx(file_path) - elif suffix in [".txt", ".md", ".json", ".fragment"]: - return cls._estimate_text(file_path) - elif suffix in [".png", ".jpg", ".jpeg"]: - return 1 # Image = 1 page - else: - logger.warning( - f"Unknown file type for billing: {suffix}, defaulting to 1 page" - ) - return 1 - except Exception as e: - logger.error(f"Error estimating pages for {file_path}: {e}") - return 1 # Fallback to minimum charge - - @classmethod - def _estimate_pdf(cls, file_path: str) -> int: - """ - Estimate pages for PDF using physical page count from metadata. - """ - try: - from pypdf import PdfReader - - reader = PdfReader(file_path) - return max(1, len(reader.pages)) - except ImportError: - logger.warning("pypdf not installed, defaulting to 1 page") - return 1 - except Exception as e: - logger.error(f"PDF estimation error: {e}") - return 1 - - @classmethod - def _estimate_pptx(cls, file_path: str) -> int: - """ - Estimate pages for PPTX using slide count. - """ - try: - from pptx import Presentation - - prs = Presentation(file_path) - return max(1, len(prs.slides)) - except ImportError: - logger.warning("python-pptx not installed, defaulting to 1 page") - return 1 - except Exception as e: - logger.error(f"PPTX estimation error: {e}") - return 1 - - @classmethod - def _estimate_docx(cls, file_path: str) -> int: - """ - Estimate pages for DOCX using word-based counting. - """ - try: - from docx import Document - - doc = Document(file_path) - total_text = "" - - # Collect text from paragraphs - for para in doc.paragraphs: - total_text += para.text + " " - - # Collect text from tables - for table in doc.tables: - for row in table.rows: - for cell in row.cells: - total_text += cell.text + " " - - word_count = count_cn_en(total_text) - return max(1, math.ceil(word_count / WORDS_PER_PAGE)) - - except ImportError: - logger.warning("python-docx not installed") - return 1 - except Exception as e: - logger.error(f"DOCX estimation error: {e}") - return 1 - - @classmethod - def _estimate_doc(cls, file_path: str) -> int: - """ - Estimate pages for DOC by converting it to DOCX first. - """ - try: - from app.services.document_parser.legacy_converter import doc_to_docx - - with tempfile.TemporaryDirectory(prefix="page-estimator-doc-") as temp_dir: - converted_path, _ = doc_to_docx(file_path, temp_dir) - return cls._estimate_docx(converted_path) - except Exception as e: - logger.error(f"DOC estimation error: {e}") - return 1 - - @classmethod - def _estimate_xlsx(cls, file_path: str) -> int: - """ - Estimate pages for XLSX using row count. - """ - try: - import pandas as pd - - xlsx = pd.ExcelFile(file_path) - total_rows = 0 - - for sheet_name in xlsx.sheet_names: - df = pd.read_excel(xlsx, sheet_name=sheet_name) - total_rows += len(df) - - return max(1, math.ceil(total_rows / ROWS_PER_PAGE)) - - except ImportError: - logger.warning("pandas not installed") - return 1 - except Exception as e: - logger.error(f"XLSX estimation error: {e}") - return 1 - - @classmethod - def _estimate_xls(cls, file_path: str) -> int: - """ - Estimate pages for XLS by converting it to XLSX first. - """ - try: - from app.services.document_parser.legacy_converter import xls_to_xlsx - - with tempfile.TemporaryDirectory(prefix="page-estimator-xls-") as temp_dir: - converted_path, _ = xls_to_xlsx(file_path, temp_dir) - return cls._estimate_xlsx(converted_path) - except Exception as e: - logger.error(f"XLS estimation error: {e}") - return 1 - - @classmethod - def _estimate_text(cls, file_path: str) -> int: - """ - Estimate pages for text files using word-based counting. - """ - try: - with open(file_path, "r", encoding="utf-8", errors="ignore") as f: - content = f.read() - - word_count = count_cn_en(content) - return max(1, math.ceil(word_count / WORDS_PER_PAGE)) - - except Exception as e: - logger.error(f"Text estimation error: {e}") - return 1 diff --git a/apps/worker/app/services/workload/url_upload_context.py b/apps/worker/app/services/workload/url_upload_context.py new file mode 100644 index 000000000..af7fb807b --- /dev/null +++ b/apps/worker/app/services/workload/url_upload_context.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from shared.core.exceptions.domain_exceptions import NotFoundException +from shared.services.redis.redis_sync_service import ( + SyncJobInfoRedisService, + SyncJobMetadataService, +) + + +@dataclass(frozen=True) +class UrlUploadContext: + s3_key: str + + +def load_url_upload_context(job_id: str, redis_service: Any) -> UrlUploadContext: + job_info_service = SyncJobInfoRedisService(redis_service) + job_info = job_info_service.get_job_info(job_id) + + if job_info: + raw_s3_key = job_info.get("s3_key") + else: + metadata_service = SyncJobMetadataService(redis_service) + job_metadata = metadata_service.get_metadata(job_id) + if not job_metadata: + raise NotFoundException( + resource="JobInfo", + resource_id=job_id, + internal_message="Job info not found in Redis or Metadata", + ) + raw_s3_key = job_metadata.get("s3_key") + + if not raw_s3_key: + raise NotFoundException( + resource="JobInfo", + resource_id="s3_key", + internal_message=f"Missing s3_key in Redis job info for job_id={job_id}", + ) + + return UrlUploadContext(s3_key=str(raw_s3_key)) diff --git a/apps/worker/app/services/workload/url_upload_service.py b/apps/worker/app/services/workload/url_upload_service.py new file mode 100644 index 000000000..27cf81616 --- /dev/null +++ b/apps/worker/app/services/workload/url_upload_service.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from typing import Any + +from loguru import logger + +from app.services.workload.url_upload_context import load_url_upload_context +from app.services.workload.url_upload_transfer import ( + assert_temp_file_within_size_limit, + cleanup_temp_file, + download_source_url_to_temp, + resolve_supported_url_extension, + upload_temp_file_to_source_storage, + verify_source_upload, +) +from shared.services.jobs.lifecycle.service import get_sync_job_lifecycle_service +from shared.services.redis.redis_sync_service import SyncRedisServiceFactory + + +def upload_url_file( + job_id: str, + source_url: str, + user_id: str | None, + job_type: str | None = None, +) -> dict[str, Any]: + del user_id, job_type + + lifecycle_service = get_sync_job_lifecycle_service() + redis_service = SyncRedisServiceFactory.get_service() + upload_context = load_url_upload_context(job_id, redis_service) + + lifecycle_service.update_progress( + job_id, progress=3, message="Validating URL file type..." + ) + file_extension = resolve_supported_url_extension(source_url) + + lifecycle_service.update_progress( + job_id, progress=10, message="Downloading file from URL..." + ) + temp_file_path = download_source_url_to_temp(source_url) + + try: + lifecycle_service.update_progress( + job_id, progress=30, message="Validating file size..." + ) + assert_temp_file_within_size_limit( + temp_file_path=temp_file_path, + file_extension=file_extension, + ) + + lifecycle_service.update_progress( + job_id, progress=50, message="Uploading file to S3..." + ) + upload_temp_file_to_source_storage( + temp_file_path=temp_file_path, + s3_key=upload_context.s3_key, + ) + + finally: + cleanup_temp_file(temp_file_path) + + lifecycle_service.update_progress( + job_id, progress=80, message="Verifying upload result..." + ) + file_info = verify_source_upload(upload_context.s3_key) + + lifecycle_service.update_progress( + job_id, + progress=100, + message="URL file upload complete, waiting for processing...", + ) + logger.info( + "URL file upload complete, waiting for S3 webhook: " + f"{job_id} -> {upload_context.s3_key}" + ) + + return { + "status": "success", + "job_id": job_id, + "s3_key": upload_context.s3_key, + "file_size": file_info.get("size"), + } diff --git a/apps/worker/app/services/workload/url_upload_transfer.py b/apps/worker/app/services/workload/url_upload_transfer.py new file mode 100644 index 000000000..c9da336a8 --- /dev/null +++ b/apps/worker/app/services/workload/url_upload_transfer.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +import os + +from loguru import logger + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ( + StorageServiceException, + ValidationException, +) +from shared.services.storage.job_file_storage import JobFileStorage +from shared.services.http.url_file_type import resolve_file_extension_sync + + +def resolve_supported_url_extension(source_url: str) -> str: + file_extension = resolve_file_extension_sync(source_url) + if file_extension: + return file_extension + + supported_formats = ", ".join(sorted(settings.get_supported_extensions())) + raise ValidationException( + user_message="Unsupported file type", + violations=[ + { + "field": "file_extension", + "description": f"Must be one of: {supported_formats}", + } + ], + ) + + +def download_source_url_to_temp(source_url: str) -> str: + storage = JobFileStorage() + try: + return storage.download_file_from_url( + source_url, + temp_dir=getattr(settings, "TMP_PATH", "/tmp"), + ) + except Exception as exc: + raise ValidationException( + user_message="Failed to download file from URL", + violations=[ + { + "field": "source_url", + "description": "Could not download file from the provided URL", + } + ], + internal_message=( + f"Failed to download file from URL: {source_url}, error: {exc}" + ), + ) + + +def assert_temp_file_within_size_limit( + *, + temp_file_path: str, + file_extension: str, +) -> int: + file_size = os.path.getsize(temp_file_path) + if file_size <= settings.MAX_FILE_SIZE: + return file_size + + limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024) + raise ValidationException( + user_message=f"File size exceeds limit (max {limit_mb}MB for {file_extension})", + violations=[ + { + "field": "file_size", + "description": ( + f"Size {file_size} bytes exceeds limit of " + f"{settings.MAX_FILE_SIZE} bytes" + ), + } + ], + ) + + +def upload_temp_file_to_source_storage( + *, + temp_file_path: str, + s3_key: str, +) -> None: + storage = JobFileStorage() + storage.upload_source_file(temp_file_path, s3_key) + logger.info(f"File uploaded to S3: {s3_key}") + + +def verify_source_upload(s3_key: str) -> dict[str, object]: + storage = JobFileStorage() + file_info = storage.verify_upload_exists(s3_key) + if file_info.get("exists"): + return dict(file_info) + + raise StorageServiceException( + user_message="We failed to verify your file upload", + internal_message=f"S3 file verification failed for {s3_key}", + ) + + +def cleanup_temp_file(temp_file_path: str | None) -> None: + if not temp_file_path: + return + if os.path.exists(temp_file_path): + os.remove(temp_file_path) + logger.debug(f"Temp file cleaned up: {temp_file_path}") diff --git a/apps/worker/tests/contract/README.md b/apps/worker/tests/contract/README.md index 725e982dc..d9f009a5c 100644 --- a/apps/worker/tests/contract/README.md +++ b/apps/worker/tests/contract/README.md @@ -18,9 +18,9 @@ These tests should avoid: - `app.core.tasks.stale_job_sweeper.expire_stale_jobs` verifies stale-job expiration, durable failure state, audit logging, and the Redis-backed duplicate-Beat lock -- `app.core.tasks.kb_tasks.upload_url_file_task` +- `app.core.tasks.document_ingestion_tasks.upload_url_file_task` verifies the cached storage target, Redis progress publication, and the stable `waiting-file` job state while mocking only outbound URL download and S3 boundaries -- `app.core.tasks.kb_tasks.parse_task` +- `app.core.tasks.document_ingestion_tasks.parse_task` verifies success publication, terminal skip handling, and failure cleanup/refund behavior while keeping billing, finalization, and retrieval publication real - `app.core.tasks.webhook_tasks.recover_orphaned_webhooks` verifies orphaned pending webhook recovery, durable QStash delivery-state persistence, and the Redis-backed duplicate-Beat lock diff --git a/apps/worker/tests/contract/test_excel_parser_contract.py b/apps/worker/tests/contract/test_excel_parser_contract.py new file mode 100644 index 000000000..6f4935cf2 --- /dev/null +++ b/apps/worker/tests/contract/test_excel_parser_contract.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +import os +from pathlib import Path + + +def _write_contract_workbook(workbook_path: Path) -> None: + import openpyxl + + workbook = openpyxl.Workbook() + visible_sheet = workbook.active + visible_sheet.title = "Visible" + visible_sheet["A1"] = "Region" + visible_sheet["B1"] = "Value" + visible_sheet["A2"] = "North" + visible_sheet["B2"] = 10 + + hidden_sheet = workbook.create_sheet("Hidden") + hidden_sheet.sheet_state = "hidden" + hidden_sheet["A1"] = "Secret" + hidden_sheet["B1"] = "Value" + hidden_sheet["A2"] = "Hidden" + hidden_sheet["B2"] = 99 + + workbook.save(workbook_path) + + +def test_xlsx_parser_contract_uses_stable_entrypoint_and_ignores_hidden_sheets( + worker_contract_environment: None, + tmp_path: Path, +) -> None: + from app.services.document_parser.parse_service import checkerboard_parse_output + + workbook_path = tmp_path / "budget.xlsx" + output_root = tmp_path / "parser-output" + _write_contract_workbook(workbook_path) + + parse_output = checkerboard_parse_output( + file_full_path=str(workbook_path), + filename="budget.xlsx", + output_dir=str(output_root), + internal_output_filename="budget.xlsx", + summary_image=False, + summary_table=False, + summary_txt=False, + smart_title_parse=False, + stopwords=[], + ) + + full_output_dir = parse_output.output_dir + parsed_df = parse_output.parsed_df + assert full_output_dir.endswith("budget.xlsx") + assert parsed_df is not None + assert parsed_df["type"].tolist() == ["table"] + assert parsed_df["path"].tolist() == ["tables/table-Visible.html"] + assert parsed_df["summary"].tolist() == ["table-Visible"] + assert "Region" in parsed_df["keywords"].iloc[0] + assert "Value" in parsed_df["keywords"].iloc[0] + + table_html = Path(full_output_dir) / "tables" / "table-Visible.html" + table_html_text = table_html.read_text(encoding="utf-8") + + assert "North" in table_html_text + assert "10" in table_html_text + assert "Secret" not in table_html_text + assert "Hidden" not in table_html_text + + +def test_parser_maps_document_name_to_task_local_path_segment( + worker_contract_environment: None, + tmp_path: Path, +) -> None: + from app.services.document_parser.parse_service import checkerboard_parse_output + + workbook_path = tmp_path / "budget.xlsx" + output_root = tmp_path / "parser-output" + _write_contract_workbook(workbook_path) + + parse_output = checkerboard_parse_output( + file_full_path=str(workbook_path), + filename="/tmp/../images.xlsx", + output_dir=str(output_root), + internal_output_filename="../../images.xlsx", + summary_image=False, + summary_table=False, + summary_txt=False, + smart_title_parse=False, + stopwords=[], + ) + + full_output_dir = parse_output.output_dir + parsed_df = parse_output.parsed_df + assert ( + os.path.commonpath([str(output_root.resolve()), full_output_dir]) + == str(output_root.resolve()) + ) + assert full_output_dir.endswith("images.xlsx") + assert parsed_df is not None + assert parsed_df["path"].tolist() == ["tables/table-Visible.html"] diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index 12f672267..1ca733f56 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -28,15 +28,15 @@ def _build_pending_file_job_metadata(source_file_name: str) -> dict[str, Any]: "namespace": "worker-contract", "source_type": "file", "source_file_name": source_file_name, - "kb_dir": "Default_Root", + "parsing_params": {"kb_dir": "legacy-ignored"}, } return job_metadata def _load_parse_task_modules() -> tuple[Any, Any, Any, Engine, Any, Any, Any]: - import app.core.tasks.kb_tasks as kb_tasks + import app.core.tasks.document_ingestion_tasks as document_ingestion_tasks + import app.services.document_ingestion.processing_run as parse_job_service import app.services.document_parser.parse_service as parse_service - import app.services.storage.sync_storage_service as sync_storage_service from shared.core.database_sync import get_sync_engine from shared.services.redis.redis_sync_service import ( SyncJobInfoRedisService, @@ -45,9 +45,9 @@ def _load_parse_task_modules() -> tuple[Any, Any, Any, Engine, Any, Any, Any]: ) return ( - kb_tasks, + document_ingestion_tasks, parse_service, - sync_storage_service, + parse_job_service, get_sync_engine(), SyncJobInfoRedisService, SyncJobMetadataService, @@ -55,6 +55,12 @@ def _load_parse_task_modules() -> tuple[Any, Any, Any, Engine, Any, Any, Any]: ) +def _load_worker_settings() -> Any: + from shared.core.config import settings + + return settings + + def _save_worker_task_cache( *, job_id: str, @@ -76,7 +82,7 @@ def _save_worker_task_cache( "s3_key": s3_key, "user_id": user_id, "webhook_enabled": False, - "job_type": "kb_management", + "job_type": "document_ingestion", "source_type": "file", }, ) @@ -84,25 +90,44 @@ def _save_worker_task_cache( return redis_service +def _patch_verify_upload_exists( + monkeypatch: MonkeyPatch, + file_info_for_storage_key: Any, +) -> None: + from shared.services.storage.job_file_storage import JobFileStorage + + monkeypatch.setattr( + JobFileStorage, + "verify_upload_exists", + lambda self, storage_key: file_info_for_storage_key(storage_key), + ) + + def _find_task_workspaces(root: Path, job_id: str) -> list[Path]: return sorted( path for path in root.iterdir() - if path.is_dir() and path.name.startswith(f"kb_task_{job_id}_") + if path.is_dir() and path.name.startswith(f"document_ingestion_task_{job_id}_") ) +def _build_fake_parse_output(*, output_dir: Path, rows: list[dict[str, Any]]) -> Any: + from app.services.document_parser.orchestration.parse_output import ParseOutput + + return ParseOutput(output_dir=str(output_dir), parsed_df=pd.DataFrame(rows)) + + def _bind_parse_task_to_current_module( monkeypatch: MonkeyPatch, *, - kb_tasks: Any, + document_ingestion_tasks: Any, ) -> None: monkeypatch.setitem( - kb_tasks.parse_task._orig_run.__globals__, + document_ingestion_tasks.parse_task._orig_run.__globals__, "_parse", - kb_tasks._parse, + document_ingestion_tasks._parse, ) - monkeypatch.setattr(kb_tasks.parse_task, "__trace__", None, raising=False) + monkeypatch.setattr(document_ingestion_tasks.parse_task, "__trace__", None, raising=False) @pytest.mark.parametrize( @@ -122,14 +147,15 @@ def test_should_parse_a_pending_file_job_and_persist_the_published_result_state( ) -> None: monkeypatch.setenv("BILLING_ENABLED", "true" if billing_enabled else "false") ( - kb_tasks, + document_ingestion_tasks, parse_service, - sync_storage_service, + parse_job_service, engine, sync_job_info_service_cls, sync_job_metadata_service_cls, sync_redis_service_factory, ) = _load_parse_task_modules() + settings = _load_worker_settings() user_id: str = f"worker-user-{uuid4().hex[:12]}" job_id: str = f"job_parse_success_{uuid4().hex[:12]}" @@ -165,9 +191,9 @@ def test_should_parse_a_pending_file_job_and_persist_the_published_result_state( sync_redis_service_factory=sync_redis_service_factory, ) - _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) - monkeypatch.setattr(kb_tasks.settings, "TMP_PATH", str(tmp_path)) - monkeypatch.setattr(kb_tasks.settings, "BILLING_ENABLED", billing_enabled) + _bind_parse_task_to_current_module(monkeypatch, document_ingestion_tasks=document_ingestion_tasks) + monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path)) + monkeypatch.setattr(settings, "BILLING_ENABLED", billing_enabled) def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: return { @@ -175,35 +201,21 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: "size": _SAMPLE_PDF_PATH.stat().st_size, } - def fake_generate_download_url(storage_key: str, bucket: str | None) -> dict[str, str]: - return {"download_url": f"https://example.test/{storage_key}"} - - monkeypatch.setattr(kb_tasks, "verify_s3_file_exists", fake_verify_s3_file_exists) - monkeypatch.setattr( - sync_storage_service, - "verify_s3_file_exists", - fake_verify_s3_file_exists, - ) - monkeypatch.setattr(kb_tasks, "generate_download_url", fake_generate_download_url) - monkeypatch.setattr( - sync_storage_service, - "generate_download_url", - fake_generate_download_url, - ) + _patch_verify_upload_exists(monkeypatch, fake_verify_s3_file_exists) def fake_download_s3_file_to_temp( - file_url: str, file_ext: str, temp_dir: str + storage_key: str, file_ext: str, temp_dir: str ) -> str: + assert storage_key == s3_key assert file_ext == ".pdf" downloaded_path = Path(temp_dir) / f"downloaded{file_ext}" shutil.copy2(_SAMPLE_PDF_PATH, downloaded_path) return str(downloaded_path) - def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: + def fake_checkerboard_parse_output(**kwargs: Any) -> Any: captured_artifacts["parse_kwargs"] = kwargs output_dir = ( Path(str(kwargs["output_dir"])) - / str(kwargs["kb_dir"]) / str(kwargs["internal_output_filename"]) ) images_dir = output_dir / "images" @@ -218,7 +230,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: parsed_rows: list[dict[str, Any]] = [ { "content": text_content_with_refs, - "path": f"Default_Root/{file_root}/公司研究/自主可控加强，寒武纪或迎来营收快速放量周期", + "path": f"{file_root}/公司研究/自主可控加强，寒武纪或迎来营收快速放量周期", "type": "text", "length": len(text_content_with_refs), "keywords": "", @@ -239,7 +251,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: }, { "content": "chunk-2", - "path": f"Default_Root/{file_root}/相关研报/要点", + "path": f"{file_root}/相关研报/要点", "type": "text", "length": 7, "keywords": "", @@ -252,7 +264,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: }, { "content": "image caption", - "path": f"Default_Root/{file_root}/images/page-1.png", + "path": f"{file_root}/images/page-1.png", "type": "image", "length": 13, "keywords": "", @@ -265,7 +277,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: }, { "content": "table content", - "path": f"Default_Root/{file_root}/tables/table-1.html", + "path": f"{file_root}/tables/table-1.html", "type": "table", "length": 13, "keywords": "", @@ -277,7 +289,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: "page_nums": "3", }, ] - return str(output_dir), pd.DataFrame(parsed_rows) + return _build_fake_parse_output(output_dir=output_dir, rows=parsed_rows) class FakeResultStorage: def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: @@ -309,11 +321,11 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: raw_files={}, ) - monkeypatch.setattr(kb_tasks, "download_s3_file_to_temp", fake_download_s3_file_to_temp) - monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse) - monkeypatch.setattr(kb_tasks, "get_result_storage", lambda: FakeResultStorage()) + monkeypatch.setattr(parse_job_service, "download_s3_file_to_temp", fake_download_s3_file_to_temp) + monkeypatch.setattr(parse_service, "checkerboard_parse_output", fake_checkerboard_parse_output) + monkeypatch.setattr(parse_job_service, "get_result_storage", lambda: FakeResultStorage()) - result = kb_tasks.parse_task.run(job_id, user_id, "kb_management") + result = document_ingestion_tasks.parse_task.run(job_id, user_id, "document_ingestion") expected_summary = "This document includes: 公司研究, 相关研报" expected_connect_to = [ @@ -338,8 +350,8 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: }, }, ] - expected_credits_charged = 3 * int(kb_tasks.settings.MICRO_DOLLARS_PER_PAGE) - expected_initial_balance = int(kb_tasks.settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 + expected_credits_charged = 3 * int(settings.MICRO_DOLLARS_PER_PAGE) + expected_initial_balance = int(settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 assert result == { "status": "success", @@ -354,7 +366,9 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: assert captured_artifacts["parse_kwargs"]["filename"] == source_file_name assert captured_artifacts["parse_kwargs"]["internal_output_filename"] == source_file_name assert Path(str(captured_artifacts["parse_kwargs"]["file_full_path"])).name == source_file_name - assert captured_artifacts["result_dir"].endswith("Default_Root/contract-parse.pdf") + assert "namespace" not in captured_artifacts["parse_kwargs"] + assert "kb_dir" not in captured_artifacts["parse_kwargs"] + assert captured_artifacts["result_dir"].endswith("contract-parse.pdf") assert captured_artifacts["doc_nav"]["file_name"] == source_file_name assert captured_artifacts["doc_nav"]["sections"][0]["title"] == "公司研究" assert captured_artifacts["manifest"]["HIERARCHY"] == { @@ -390,6 +404,8 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: metadata = sync_job_metadata_service_cls(redis_service).get_metadata(job_id) assert metadata is not None assert metadata["page_count"] == 3 + assert metadata["workload_estimate_method"] == "pdf_metadata" + assert "workload_estimate_fallback_reason" not in metadata assert metadata["billing_status"] == expected_billing_status if billing_enabled: assert metadata["billing_amount_micro_dollars"] == expected_credits_charged @@ -426,7 +442,7 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: connection.execute( text( """ - SELECT delivery_mode, result_s3_key, result_size, inline_payload + SELECT id, delivery_mode, result_s3_key, result_size, inline_payload FROM job_results WHERE job_id = :job_id """ @@ -436,6 +452,21 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: .mappings() .one() ) + job_chunks = list( + connection.execute( + text( + """ + SELECT chunk_type, text, path, sort_order + FROM job_chunks + WHERE job_result_id = :job_result_id + ORDER BY sort_order + """ + ), + {"job_result_id": job_result_row["id"]}, + ) + .mappings() + .all() + ) document_row = ( connection.execute( text( @@ -515,7 +546,7 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: SELECT transition_reason, to_state FROM job_state_audit_logs WHERE job_id = :job_id - ORDER BY created_at ASC + ORDER BY id ASC """ ), {"job_id": job_id}, @@ -539,6 +570,17 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: assert job_result_row["result_s3_key"] == f"results/{job_id}.zip" assert job_result_row["result_size"] > 0 assert dict(job_result_row["inline_payload"])["checksum"] + assert [chunk["chunk_type"] for chunk in job_chunks] == [ + "text", + "text", + "image", + "table", + ] + assert job_chunks[0]["text"] == text_content_with_refs + assert job_chunks[0]["path"].endswith( + "公司研究/自主可控加强，寒武纪或迎来营收快速放量周期" + ) + assert job_chunks[0]["sort_order"] == 0 assert document_row["namespace"] == "worker-contract" assert document_row["status"] == "active" assert document_row["source_file_name"] == source_file_name @@ -575,14 +617,15 @@ def test_should_export_full_result_when_publication_deduplicates_existing_chunks ) -> None: monkeypatch.setenv("BILLING_ENABLED", "false") ( - kb_tasks, + document_ingestion_tasks, parse_service, - sync_storage_service, + parse_job_service, engine, sync_job_info_service_cls, sync_job_metadata_service_cls, sync_redis_service_factory, ) = _load_parse_task_modules() + settings = _load_worker_settings() user_id: str = f"worker-user-{uuid4().hex[:12]}" existing_job_id: str = f"job_existing_{uuid4().hex[:12]}" @@ -709,7 +752,7 @@ def test_should_export_full_result_when_publication_deduplicates_existing_chunks :result_id, 'text', 'already published text', - 'Default_Root/existing.pdf/Section/Duplicate text', + 'existing.pdf/Section/Duplicate text', NULL, CAST(:text_metadata AS JSON), 0, @@ -724,7 +767,7 @@ def test_should_export_full_result_when_publication_deduplicates_existing_chunks :result_id, 'image', 'already published image', - 'Default_Root/existing.pdf/images/duplicate.png', + 'existing.pdf/images/duplicate.png', 'images/duplicate.png', CAST(:image_metadata AS JSON), 1, @@ -764,9 +807,9 @@ def test_should_export_full_result_when_publication_deduplicates_existing_chunks sync_redis_service_factory=sync_redis_service_factory, ) - _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) - monkeypatch.setattr(kb_tasks.settings, "TMP_PATH", str(tmp_path)) - monkeypatch.setattr(kb_tasks.settings, "BILLING_ENABLED", False) + _bind_parse_task_to_current_module(monkeypatch, document_ingestion_tasks=document_ingestion_tasks) + monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path)) + monkeypatch.setattr(settings, "BILLING_ENABLED", False) def fake_cleanup_task_workspace(workspace_dir: str | None) -> bool: captured_artifacts["workspace_dir"] = workspace_dir @@ -778,20 +821,17 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: "size": _SAMPLE_PDF_PATH.stat().st_size, } - def fake_generate_download_url(storage_key: str, bucket: str | None) -> dict[str, str]: - return {"download_url": f"https://example.test/{storage_key}"} - def fake_download_s3_file_to_temp( - file_url: str, file_ext: str, temp_dir: str + storage_key: str, file_ext: str, temp_dir: str ) -> str: + assert storage_key == s3_key downloaded_path = Path(temp_dir) / f"downloaded{file_ext}" shutil.copy2(_SAMPLE_PDF_PATH, downloaded_path) return str(downloaded_path) - def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: + def fake_checkerboard_parse_output(**kwargs: Any) -> Any: output_dir = ( Path(str(kwargs["output_dir"])) - / str(kwargs["kb_dir"]) / str(kwargs["internal_output_filename"]) ) images_dir = output_dir / "images" @@ -803,7 +843,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: parsed_rows: list[dict[str, Any]] = [ { "content": "duplicate text", - "path": f"Default_Root/{file_root}/Section/Duplicate text", + "path": f"{file_root}/Section/Duplicate text", "type": "text", "length": 14, "keywords": "", @@ -816,7 +856,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: }, { "content": "duplicate image", - "path": f"Default_Root/{file_root}/images/duplicate.png", + "path": f"{file_root}/images/duplicate.png", "type": "image", "length": 15, "keywords": "", @@ -829,7 +869,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: }, { "content": "new text", - "path": f"Default_Root/{file_root}/Section/New text", + "path": f"{file_root}/Section/New text", "type": "text", "length": 8, "keywords": "", @@ -841,7 +881,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: "page_nums": "3", }, ] - return str(output_dir), pd.DataFrame(parsed_rows) + return _build_fake_parse_output(output_dir=output_dir, rows=parsed_rows) class FakeResultStorage: def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: @@ -867,24 +907,13 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: raw_files={}, ) - monkeypatch.setattr(kb_tasks, "verify_s3_file_exists", fake_verify_s3_file_exists) - monkeypatch.setattr( - sync_storage_service, - "verify_s3_file_exists", - fake_verify_s3_file_exists, - ) - monkeypatch.setattr(kb_tasks, "generate_download_url", fake_generate_download_url) - monkeypatch.setattr( - sync_storage_service, - "generate_download_url", - fake_generate_download_url, - ) - monkeypatch.setattr(kb_tasks, "download_s3_file_to_temp", fake_download_s3_file_to_temp) - monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse) - monkeypatch.setattr(kb_tasks, "get_result_storage", lambda: FakeResultStorage()) - monkeypatch.setattr(kb_tasks, "cleanup_task_workspace", fake_cleanup_task_workspace) + _patch_verify_upload_exists(monkeypatch, fake_verify_s3_file_exists) + monkeypatch.setattr(parse_job_service, "download_s3_file_to_temp", fake_download_s3_file_to_temp) + monkeypatch.setattr(parse_service, "checkerboard_parse_output", fake_checkerboard_parse_output) + monkeypatch.setattr(parse_job_service, "get_result_storage", lambda: FakeResultStorage()) + monkeypatch.setattr(parse_job_service, "cleanup_task_workspace", fake_cleanup_task_workspace) - result = kb_tasks.parse_task.run(job_id, user_id, "kb_management") + result = document_ingestion_tasks.parse_task.run(job_id, user_id, "document_ingestion") assert result["contents_count"] == 3 assert "images/duplicate.png" in captured_artifacts["zip_entries"] @@ -963,14 +992,15 @@ def test_should_initialize_billing_once_for_concurrent_parse_tasks( tmp_path: Path, ) -> None: ( - kb_tasks, + document_ingestion_tasks, parse_service, - sync_storage_service, + parse_job_service, engine, sync_job_info_service_cls, sync_job_metadata_service_cls, sync_redis_service_factory, ) = _load_parse_task_modules() + settings = _load_worker_settings() user_id: str = f"worker-concurrent-user-{uuid4().hex[:12]}" job_ids: list[str] = [f"job_cb_{index}_{uuid4().hex[:12]}" for index in range(2)] @@ -1010,9 +1040,9 @@ def test_should_initialize_billing_once_for_concurrent_parse_tasks( sync_redis_service_factory=sync_redis_service_factory, ) - _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) - monkeypatch.setattr(kb_tasks.settings, "TMP_PATH", str(tmp_path)) - monkeypatch.setattr(kb_tasks.settings, "BILLING_ENABLED", True) + _bind_parse_task_to_current_module(monkeypatch, document_ingestion_tasks=document_ingestion_tasks) + monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path)) + monkeypatch.setattr(settings, "BILLING_ENABLED", True) def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: return { @@ -1020,12 +1050,10 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: "size": _SAMPLE_PDF_PATH.stat().st_size, } - def fake_generate_download_url(storage_key: str, bucket: str | None) -> dict[str, str]: - return {"download_url": f"https://example.test/{storage_key}"} - def fake_download_s3_file_to_temp( - file_url: str, file_ext: str, temp_dir: str + storage_key: str, file_ext: str, temp_dir: str ) -> str: + assert storage_key in s3_keys.values() assert file_ext == ".pdf" downloaded_path = Path(temp_dir) / f"downloaded{file_ext}" shutil.copy2(_SAMPLE_PDF_PATH, downloaded_path) @@ -1033,14 +1061,15 @@ def fake_download_s3_file_to_temp( billing_start_barrier = Barrier(len(job_ids)) - def fake_estimate_page_count(file_path: str) -> int: + def fake_estimate_workload(file_path: str) -> Any: + from app.services.document_ingestion.page_estimator import WorkloadEstimate + billing_start_barrier.wait(timeout=10) - return 1 + return WorkloadEstimate(page_count=1, method="contract_fake") - def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: + def fake_checkerboard_parse_output(**kwargs: Any) -> Any: output_dir = ( Path(str(kwargs["output_dir"])) - / str(kwargs["kb_dir"]) / str(kwargs["internal_output_filename"]) ) output_dir.mkdir(parents=True, exist_ok=True) @@ -1050,7 +1079,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: parsed_rows: list[dict[str, Any]] = [ { "content": "chunk body", - "path": f"Default_Root/{file_root}/Section/Point", + "path": f"{file_root}/Section/Point", "type": "text", "length": 10, "keywords": "", @@ -1062,7 +1091,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: "page_nums": "1", } ] - return str(output_dir), pd.DataFrame(parsed_rows) + return _build_fake_parse_output(output_dir=output_dir, rows=parsed_rows) class FakeResultStorage: def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: @@ -1072,32 +1101,25 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: raw_files={}, ) - monkeypatch.setattr(kb_tasks, "verify_s3_file_exists", fake_verify_s3_file_exists) + _patch_verify_upload_exists(monkeypatch, fake_verify_s3_file_exists) + monkeypatch.setattr(parse_job_service, "download_s3_file_to_temp", fake_download_s3_file_to_temp) monkeypatch.setattr( - sync_storage_service, - "verify_s3_file_exists", - fake_verify_s3_file_exists, + parse_job_service.PageEstimator, + "estimate_workload", + fake_estimate_workload, ) - monkeypatch.setattr(kb_tasks, "generate_download_url", fake_generate_download_url) - monkeypatch.setattr( - sync_storage_service, - "generate_download_url", - fake_generate_download_url, - ) - monkeypatch.setattr(kb_tasks, "download_s3_file_to_temp", fake_download_s3_file_to_temp) - monkeypatch.setattr(kb_tasks.PageEstimator, "estimate", fake_estimate_page_count) - monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse) - monkeypatch.setattr(kb_tasks, "get_result_storage", lambda: FakeResultStorage()) + monkeypatch.setattr(parse_service, "checkerboard_parse_output", fake_checkerboard_parse_output) + monkeypatch.setattr(parse_job_service, "get_result_storage", lambda: FakeResultStorage()) def run_parse_task(job_id: str) -> dict[str, Any]: - return dict(kb_tasks.parse_task.run(job_id, user_id, "kb_management")) + return dict(document_ingestion_tasks.parse_task.run(job_id, user_id, "document_ingestion")) with ThreadPoolExecutor(max_workers=len(job_ids)) as executor: results = list(executor.map(run_parse_task, job_ids)) - expected_credits_charged = int(kb_tasks.settings.MICRO_DOLLARS_PER_PAGE) + expected_credits_charged = int(settings.MICRO_DOLLARS_PER_PAGE) expected_initial_balance = ( - int(kb_tasks.settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 + int(settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 ) with engine.begin() as connection: @@ -1213,14 +1235,15 @@ def test_should_skip_parse_task_when_the_job_is_already_terminal( tmp_path: Path, ) -> None: ( - kb_tasks, + document_ingestion_tasks, parse_service, - sync_storage_service, + parse_job_service, engine, sync_job_info_service_cls, sync_job_metadata_service_cls, sync_redis_service_factory, ) = _load_parse_task_modules() + settings = _load_worker_settings() user_id: str = f"worker-user-{uuid4().hex[:12]}" job_id: str = f"job_parse_skipped_{uuid4().hex[:12]}" @@ -1252,33 +1275,21 @@ def test_should_skip_parse_task_when_the_job_is_already_terminal( sync_redis_service_factory=sync_redis_service_factory, ) - _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) - monkeypatch.setattr(kb_tasks.settings, "TMP_PATH", str(tmp_path)) + _bind_parse_task_to_current_module(monkeypatch, document_ingestion_tasks=document_ingestion_tasks) + monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path)) def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: return {"exists": storage_key == s3_key, "size": 1024} - monkeypatch.setattr(kb_tasks, "verify_s3_file_exists", fake_verify_s3_file_exists) - monkeypatch.setattr( - sync_storage_service, - "verify_s3_file_exists", - fake_verify_s3_file_exists, - ) - monkeypatch.setattr( - kb_tasks, - "generate_download_url", - lambda *_args, **_kwargs: (_ for _ in ()).throw( - AssertionError("terminal parse task should not request a download URL") - ), - ) + _patch_verify_upload_exists(monkeypatch, fake_verify_s3_file_exists) monkeypatch.setattr( parse_service, - "checkerboard_inject_parse", + "checkerboard_parse_output", lambda **_kwargs: (_ for _ in ()).throw( AssertionError("terminal parse task should not invoke the parser") ), ) - result = kb_tasks.parse_task.run(job_id, user_id, "kb_management") + result = document_ingestion_tasks.parse_task.run(job_id, user_id, "document_ingestion") assert result == { "status": "skipped", @@ -1314,14 +1325,15 @@ def test_should_mark_the_job_failed_and_cleanup_the_workspace_when_parse_executi tmp_path: Path, ) -> None: ( - kb_tasks, + document_ingestion_tasks, parse_service, - sync_storage_service, + parse_job_service, engine, sync_job_info_service_cls, sync_job_metadata_service_cls, sync_redis_service_factory, ) = _load_parse_task_modules() + settings = _load_worker_settings() user_id: str = f"worker-user-{uuid4().hex[:12]}" job_id: str = f"job_parse_failure_{uuid4().hex[:12]}" @@ -1353,61 +1365,48 @@ def test_should_mark_the_job_failed_and_cleanup_the_workspace_when_parse_executi sync_redis_service_factory=sync_redis_service_factory, ) - _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) - monkeypatch.setattr(kb_tasks.settings, "TMP_PATH", str(tmp_path)) + _bind_parse_task_to_current_module(monkeypatch, document_ingestion_tasks=document_ingestion_tasks) + monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path)) def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: return { "exists": storage_key == s3_key, "size": _SAMPLE_PDF_PATH.stat().st_size, } - def fake_generate_download_url(storage_key: str, bucket: str | None) -> dict[str, str]: - return {"download_url": f"https://example.test/{storage_key}"} - - monkeypatch.setattr(kb_tasks, "verify_s3_file_exists", fake_verify_s3_file_exists) - monkeypatch.setattr( - sync_storage_service, - "verify_s3_file_exists", - fake_verify_s3_file_exists, - ) - monkeypatch.setattr(kb_tasks, "generate_download_url", fake_generate_download_url) - monkeypatch.setattr( - sync_storage_service, - "generate_download_url", - fake_generate_download_url, - ) + _patch_verify_upload_exists(monkeypatch, fake_verify_s3_file_exists) def fake_download_s3_file_to_temp( - file_url: str, file_ext: str, temp_dir: str + storage_key: str, file_ext: str, temp_dir: str ) -> str: + assert storage_key == s3_key downloaded_path = Path(temp_dir) / f"downloaded{file_ext}" shutil.copy2(_SAMPLE_PDF_PATH, downloaded_path) return str(downloaded_path) - monkeypatch.setattr(kb_tasks, "download_s3_file_to_temp", fake_download_s3_file_to_temp) + monkeypatch.setattr(parse_job_service, "download_s3_file_to_temp", fake_download_s3_file_to_temp) monkeypatch.setattr( parse_service, - "checkerboard_inject_parse", + "checkerboard_parse_output", lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("parse failed")), ) monkeypatch.setattr( - kb_tasks, + parse_job_service, "get_result_storage", lambda: (_ for _ in ()).throw( AssertionError("result storage should not run after parser failure") ), ) - result = kb_tasks.parse_task.apply( - args=[job_id, user_id, "kb_management"], + result = document_ingestion_tasks.parse_task.apply( + args=[job_id, user_id, "document_ingestion"], throw=False, ) assert result.status == "FAILURE" assert _find_task_workspaces(tmp_path, job_id) == [] - expected_credits_charged = 3 * int(kb_tasks.settings.MICRO_DOLLARS_PER_PAGE) - expected_initial_balance = int(kb_tasks.settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 + expected_credits_charged = 3 * int(settings.MICRO_DOLLARS_PER_PAGE) + expected_initial_balance = int(settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 with engine.begin() as connection: job_row = ( @@ -1460,7 +1459,7 @@ def fake_download_s3_file_to_temp( SELECT transition_reason, to_state FROM job_state_audit_logs WHERE job_id = :job_id - ORDER BY created_at ASC + ORDER BY id ASC """ ), {"job_id": job_id}, diff --git a/apps/worker/tests/contract/test_url_upload_contract.py b/apps/worker/tests/contract/test_url_upload_contract.py index 8c3ade508..4765472cd 100644 --- a/apps/worker/tests/contract/test_url_upload_contract.py +++ b/apps/worker/tests/contract/test_url_upload_contract.py @@ -13,15 +13,40 @@ from support.contract_database import insert_contract_job, insert_contract_user -def _load_upload_task_modules() -> tuple[Any, Engine, Any, Any]: - import app.core.tasks.kb_tasks as kb_tasks +def _load_upload_task_modules() -> tuple[Any, Any, Engine, Any, Any]: + import app.core.tasks.document_ingestion_tasks as document_ingestion_tasks + import app.services.workload.url_upload_service as url_upload_service from shared.core.database_sync import get_sync_engine from shared.services.redis.redis_sync_service import ( SyncJobInfoRedisService, SyncRedisServiceFactory, ) - return kb_tasks, get_sync_engine(), SyncJobInfoRedisService, SyncRedisServiceFactory + return ( + document_ingestion_tasks, + url_upload_service, + get_sync_engine(), + SyncJobInfoRedisService, + SyncRedisServiceFactory, + ) + + +def _bind_upload_task_to_current_module( + monkeypatch: MonkeyPatch, + *, + document_ingestion_tasks: Any, +) -> None: + monkeypatch.setitem( + document_ingestion_tasks.upload_url_file_task._orig_run.__globals__, + "_upload_url_file", + document_ingestion_tasks._upload_url_file, + ) + monkeypatch.setattr( + document_ingestion_tasks.upload_url_file_task, + "__trace__", + None, + raising=False, + ) def test_should_upload_a_url_job_to_the_expected_storage_key_and_publish_progress( @@ -29,9 +54,14 @@ def test_should_upload_a_url_job_to_the_expected_storage_key_and_publish_progres monkeypatch: MonkeyPatch, tmp_path: Path, ) -> None: - kb_tasks, engine, sync_job_info_service_cls, sync_redis_service_factory = ( - _load_upload_task_modules() - ) + ( + document_ingestion_tasks, + url_upload_service, + engine, + sync_job_info_service_cls, + sync_redis_service_factory, + ) = _load_upload_task_modules() + from shared.core.config import settings user_id = f"worker-user-{uuid4().hex[:12]}" job_id = f"job_url_upload_{uuid4().hex[:12]}" @@ -43,27 +73,50 @@ def test_should_upload_a_url_job_to_the_expected_storage_key_and_publish_progres def resolve_public_address( host: str, port: int | None, + *args: object, + **kwargs: object, ) -> list[tuple[socket.AddressFamily, socket.SocketKind, int, str, tuple[str, int]]]: return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 0))] monkeypatch.setattr( - kb_tasks, - "download_file_from_url", + url_upload_service, + "download_source_url_to_temp", + lambda _source_url: str(downloaded_path), + ) + monkeypatch.setitem( + document_ingestion_tasks.upload_url_file.__globals__, + "download_source_url_to_temp", lambda _source_url: str(downloaded_path), ) monkeypatch.setattr( - kb_tasks, - "upload_to_s3", - lambda local_path, storage_key, bucket: uploaded_calls.append( - (local_path, storage_key, bucket) + url_upload_service, + "upload_temp_file_to_source_storage", + lambda *, temp_file_path, s3_key: uploaded_calls.append( + (temp_file_path, s3_key, settings.S3_BUCKET_NAME) + ), + ) + monkeypatch.setitem( + document_ingestion_tasks.upload_url_file.__globals__, + "upload_temp_file_to_source_storage", + lambda *, temp_file_path, s3_key: uploaded_calls.append( + (temp_file_path, s3_key, settings.S3_BUCKET_NAME) ), ) monkeypatch.setattr( - kb_tasks, - "verify_s3_file_exists", + url_upload_service, + "verify_source_upload", + lambda storage_key: {"exists": storage_key == s3_key, "size": 3}, + ) + monkeypatch.setitem( + document_ingestion_tasks.upload_url_file.__globals__, + "verify_source_upload", lambda storage_key: {"exists": storage_key == s3_key, "size": 3}, ) monkeypatch.setattr(socket, "getaddrinfo", resolve_public_address) + _bind_upload_task_to_current_module( + monkeypatch, + document_ingestion_tasks=document_ingestion_tasks, + ) downloaded_path.write_bytes(b"pdf") @@ -93,16 +146,16 @@ def resolve_public_address( "s3_key": s3_key, "user_id": user_id, "webhook_enabled": False, - "job_type": "kb_management", + "job_type": "document_ingestion", "source_type": "url", }, ) - result = kb_tasks.upload_url_file_task.run( + result = document_ingestion_tasks.upload_url_file_task.run( job_id, source_url, user_id, - "kb_management", + "document_ingestion", ) assert result == { @@ -112,7 +165,7 @@ def resolve_public_address( "file_size": 3, } assert uploaded_calls == [ - (str(downloaded_path), s3_key, kb_tasks.settings.S3_BUCKET_NAME), + (str(downloaded_path), s3_key, settings.S3_BUCKET_NAME), ] assert os.path.exists(downloaded_path) is False @@ -140,4 +193,3 @@ def resolve_public_address( assert job_row["status"] == "waiting-file" assert job_row["source_type"] == "url" assert job_row["s3_key"] == s3_key - diff --git a/apps/worker/tests/contract/test_webhook_recovery_contract.py b/apps/worker/tests/contract/test_webhook_recovery_contract.py index b30633b9f..606e25818 100644 --- a/apps/worker/tests/contract/test_webhook_recovery_contract.py +++ b/apps/worker/tests/contract/test_webhook_recovery_contract.py @@ -36,7 +36,9 @@ def _insert_webhook_event( created_at: datetime, updated_at: datetime | None = None, qstash_message_id: str | None = None, + payload: dict[str, Any] | None = None, ) -> None: + event_payload = payload or {"event": "job.failed", "job_id": job_id} connection.execute( text( """ @@ -69,7 +71,7 @@ def _insert_webhook_event( "id": event_id, "job_id": job_id, "target_url": target_url, - "payload": json.dumps({"event": "job.failed", "job_id": job_id}), + "payload": json.dumps(event_payload), "status": status, "attempts": attempts, "next_retry_at": None, @@ -80,6 +82,51 @@ def _insert_webhook_event( ) +def _insert_job_result( + connection: Connection, + *, + job_result_id: str, + job_id: str, + result_s3_key: str, + inline_payload: dict[str, Any], +) -> None: + timestamp = _utc_now() + connection.execute( + text( + """ + INSERT INTO job_results ( + id, + job_id, + delivery_mode, + inline_payload, + result_s3_key, + result_size, + created_at, + updated_at + ) VALUES ( + :id, + :job_id, + 'url', + CAST(:inline_payload AS JSON), + :result_s3_key, + :result_size, + :created_at, + :updated_at + ) + """ + ), + { + "id": job_result_id, + "job_id": job_id, + "inline_payload": json.dumps(inline_payload), + "result_s3_key": result_s3_key, + "result_size": 123, + "created_at": timestamp, + "updated_at": timestamp, + }, + ) + + def _load_worker_modules() -> tuple[Any, Any, Engine]: import app.core.tasks.webhook_tasks as webhook_tasks from shared.core.database_sync import get_sync_engine @@ -132,8 +179,8 @@ def publish(self, **kwargs: Any) -> SimpleNamespace: ), ) monkeypatch.setattr( - publisher, - "_get_client", + publisher._client_adapter, + "get_client", lambda: SimpleNamespace(message=FakeMessageClient()), ) @@ -271,6 +318,130 @@ def publish(self, **kwargs: Any) -> SimpleNamespace: assert secrets_count_row["secrets_count"] == 1 +def test_should_publish_completed_webhook_with_result_delivery_payload( + worker_contract_environment: None, + monkeypatch: MonkeyPatch, +) -> None: + _, qstash_publisher, engine = _load_worker_modules() + from shared.core.config import app_config + from shared.services.jobs.result_delivery import JobResultDeliveryResolver + from shared.services.storage.job_file_storage import JobFileStorage + + user_id = f"worker-user-{uuid4().hex[:12]}" + target_url = "https://hooks.contract.test/worker" + job_id = f"job_completed_{uuid4().hex[:12]}" + event_id = str(uuid4()) + result_s3_key = f"results/{job_id}.zip" + published_calls: list[dict[str, Any]] = [] + signed_url_calls: list[dict[str, Any]] = [] + + class FakeMessageClient: + def publish(self, **kwargs: Any) -> SimpleNamespace: + published_calls.append(kwargs) + return SimpleNamespace(message_id=f"msg_{event_id}") + + class FakeStorageAdapter: + def generate_presigned_url( + self, + key: str, + expiration: int = 3600, + bucket: str | None = None, + method: str = "GET", + headers: dict[str, str] | None = None, + ) -> str: + signed_url_calls.append( + { + "key": key, + "expiration": expiration, + "bucket": bucket, + "method": method, + "headers": headers, + } + ) + return f"signed://{bucket}/{key}?expires={expiration}" + + monkeypatch.setattr( + qstash_publisher, + "validate_http_url_and_resolve_ip", + lambda *args, **kwargs: SimpleNamespace( + is_valid=True, + error_message=None, + validated_ip="93.184.216.34", + hostname="hooks.contract.test", + ), + ) + monkeypatch.setattr( + JobResultDeliveryResolver, + "__init__", + lambda self: setattr( + self, + "_storage", + JobFileStorage(storage_adapter=FakeStorageAdapter()), + ), + ) + publisher = qstash_publisher.QStashWebhookPublisher() + monkeypatch.setattr( + publisher._client_adapter, + "get_client", + lambda: SimpleNamespace(message=FakeMessageClient()), + ) + + now = _utc_now() + with engine.begin() as connection: + insert_contract_user(connection, user_id=user_id) + insert_contract_job( + connection, + job_id=job_id, + user_id=user_id, + status="done", + source_type="file", + webhook_url=target_url, + webhook_enabled=True, + job_metadata=_build_file_job_metadata(), + billing_status="charged", + ) + _insert_job_result( + connection, + job_result_id=str(uuid4()), + job_id=job_id, + result_s3_key=result_s3_key, + inline_payload={"checksum": "contract-checksum"}, + ) + _insert_webhook_event( + connection, + event_id=event_id, + job_id=job_id, + target_url=target_url, + status="pending", + attempts=0, + created_at=now, + payload={"event": "job.completed", "job_id": job_id}, + ) + + message_id = publisher.publish_event(event_id) + + assert message_id == f"msg_{event_id}" + assert len(published_calls) == 1 + assert signed_url_calls == [ + { + "key": result_s3_key, + "expiration": 3600, + "bucket": app_config.S3_RESULTS_BUCKET, + "method": "GET", + "headers": None, + } + ] + + published_payload = json.loads(published_calls[0]["body"]) + assert published_payload["event"] == "job.completed" + assert published_payload["job_id"] == job_id + assert published_payload["result"] == {"checksum": "contract-checksum"} + assert published_payload["result_url"] == ( + f"signed://{app_config.S3_RESULTS_BUCKET}/{result_s3_key}" + "?expires=3600" + ) + + def test_should_reconcile_stale_delivering_webhook_events_from_qstash_logs( worker_contract_environment: None, monkeypatch: MonkeyPatch, diff --git a/apps/worker/tests/contract/test_worker_bootstrap_contract.py b/apps/worker/tests/contract/test_worker_bootstrap_contract.py index fa0ad1c05..c3701b0c6 100644 --- a/apps/worker/tests/contract/test_worker_bootstrap_contract.py +++ b/apps/worker/tests/contract/test_worker_bootstrap_contract.py @@ -2,6 +2,8 @@ import sys +from pytest import MonkeyPatch + def test_should_register_worker_task_modules_for_celery_consumers( worker_contract_environment: None, @@ -10,13 +12,15 @@ def test_should_register_worker_task_modules_for_celery_consumers( from shared.core.celery_app import celery_app expected_task_names: tuple[str, ...] = ( + "app.core.tasks.document_ingestion_tasks.upload_url_file_task", + "app.core.tasks.document_ingestion_tasks.parse_task", "app.core.tasks.kb_tasks.upload_url_file_task", "app.core.tasks.kb_tasks.parse_task", "app.core.tasks.stale_job_sweeper.expire_stale_jobs", "app.core.tasks.webhook_tasks.recover_orphaned_webhooks", ) task_module_names: tuple[str, ...] = ( - "app.core.tasks.kb_tasks", + "app.core.tasks.document_ingestion_tasks", "app.core.tasks.stale_job_sweeper", "app.core.tasks.webhook_tasks", ) @@ -31,3 +35,44 @@ def test_should_register_worker_task_modules_for_celery_consumers( for task_name in expected_task_names: assert task_name in celery_app.tasks + + +def test_should_consume_current_and_legacy_ingestion_queues( + worker_contract_environment: None, + monkeypatch: MonkeyPatch, +) -> None: + from app.core import worker_bootstrap + + worker_main_calls: list[list[str]] = [] + beat_commands: list[list[str]] = [] + + class FakeBeatProcess: + pass + + def record_beat_command(command: list[str]) -> FakeBeatProcess: + beat_commands.append(command) + return FakeBeatProcess() + + def record_worker_main(args: list[str]) -> None: + worker_main_calls.append(args) + + monkeypatch.setattr(worker_bootstrap.subprocess, "Popen", record_beat_command) + monkeypatch.setattr(worker_bootstrap.celery_app, "worker_main", record_worker_main) + + worker_bootstrap.run_worker() + + assert beat_commands != [] + assert len(worker_main_calls) == 1 + + worker_args = worker_main_calls[0] + queue_arg = worker_args[worker_args.index("-Q") + 1] + consumed_queues = set(queue_arg.split(",")) + + assert { + "document_ingestion_high", + "document_ingestion_medium", + "document_ingestion_low", + "kb_high", + "kb_medium", + "kb_low", + }.issubset(consumed_queues) diff --git a/apps/worker/tests/support/contract_database.py b/apps/worker/tests/support/contract_database.py index f586f1339..3cb31c307 100644 --- a/apps/worker/tests/support/contract_database.py +++ b/apps/worker/tests/support/contract_database.py @@ -39,7 +39,7 @@ def insert_contract_job( *, job_id: str, user_id: str, - job_type: str = "kb_management", + job_type: str = "document_ingestion", status: str = "pending", source_type: str = "file", file_path: str | None = None, diff --git a/docs/adr/0001-keep-routes-and-worker-tasks-as-adapters.md b/docs/adr/0001-keep-routes-and-worker-tasks-as-adapters.md new file mode 100644 index 000000000..041ad3271 --- /dev/null +++ b/docs/adr/0001-keep-routes-and-worker-tasks-as-adapters.md @@ -0,0 +1,25 @@ +# 0001 Keep Routes And Worker Tasks As Adapters + +## Status + +Accepted + +## Context + +Knowhere has HTTP routes, internal callback routes, and Celery tasks that start +workflows. Bugs become harder to localize when these adapters also own parser, +publication, billing, retrieval, or webhook implementation details. + +## Decision + +Routes and worker tasks should translate external inputs into application or +shared workflow calls. Domain behavior should live behind workflow modules such +as Document Ingestion, Worker Document Parsing, Publication, Retrieval, Billing +Workflow, Storage Event Intake, and Webhook delivery. + +## Consequences + +Adapters stay thin and contract tests can assert public behavior at HTTP, +worker-task, database, storage, Redis, or event surfaces. New feature work +should prefer deep workflow modules over adding branch-heavy logic to routes or +Celery task functions. diff --git a/docs/adr/0002-use-typed-workflow-outcomes.md b/docs/adr/0002-use-typed-workflow-outcomes.md new file mode 100644 index 000000000..fce8dddd3 --- /dev/null +++ b/docs/adr/0002-use-typed-workflow-outcomes.md @@ -0,0 +1,28 @@ +# 0002 Use Typed Workflow Outcomes + +## Status + +Accepted + +## Context + +Several Knowhere workflows used tuples, dictionaries, or booleans to represent +rich behavior. That made retry, fallback, and transition decisions harder to +reason about because callers lost the reason behind a result. + +## Decision + +New or deepened workflow seams should use typed outcomes when callers need more +than a yes/no answer. Current examples are `JobTransitionOutcome`, +`WorkloadEstimate`, `ParseOutput`, `ParseArtifact`, `GeneratedResultPackage`, +and `PostCommitEffectPlan`. + +Boolean and tuple facades may remain only where there is a real public contract. +The implementation should concentrate reason-bearing behavior behind typed +outcomes. + +## Consequences + +Contract tests can assert stable behavior and failure reasons without reaching +into private helper call order. Avoid internal-only compatibility facades; if +there is no external consumer, update callers to the real module boundary. diff --git a/docs/adr/0003-keep-retrieval-workflow-policy-explicit.md b/docs/adr/0003-keep-retrieval-workflow-policy-explicit.md new file mode 100644 index 000000000..ccf97b846 --- /dev/null +++ b/docs/adr/0003-keep-retrieval-workflow-policy-explicit.md @@ -0,0 +1,29 @@ +# 0003 Keep Retrieval Workflow Policy Explicit + +## Status + +Accepted + +## Context + +Retrieval requests include scope, document/section exclusions, data type, +signal paths, channel selection, channel weights, internal recall, and ranking +policy. When these fields are passed as loose keyword arguments, the legacy and +agentic routes can drift. + +## Decision + +Retrieval should keep request policy in typed request modules. `RetrievalQuery` +owns cache and route policy. `WorkflowRunRequest` and `WorkflowStepRequest` +carry the agentic workflow projection of that policy through planning and step +execution. + +Fields that are accepted but not yet implemented in a route, such as reranking +or threshold semantics for multi-step workflow answers, must remain explicit in +the request type rather than disappearing silently. + +## Consequences + +The request type becomes the test surface for retrieval policy. Product +behavior changes such as enabling workflow reranking or threshold filtering +should be separate decisions with contract tests for the selected phase. diff --git a/docs/adr/README.md b/docs/adr/README.md new file mode 100644 index 000000000..3d0cb4fd0 --- /dev/null +++ b/docs/adr/README.md @@ -0,0 +1,12 @@ +# Architecture Decision Records + +This directory records durable architecture decisions for Knowhere. Keep ADRs +short, repo-facing, and focused on decisions future maintainers or agents might +otherwise re-litigate. + +Use this shape: + +- Status +- Context +- Decision +- Consequences diff --git a/packages/shared-python/shared/core/async_utils.py b/packages/shared-python/shared/core/async_utils.py deleted file mode 100644 index 1975d8b7a..000000000 --- a/packages/shared-python/shared/core/async_utils.py +++ /dev/null @@ -1,32 +0,0 @@ -import asyncio -from typing import Any, Coroutine, TypeVar - -T = TypeVar("T") - - -def run_async_task(coro: Coroutine[Any, Any, T]) -> T: - """ - Run an async task in a synchronous context, reusing the event loop if possible. - - This function attempts to get the current event loop. If it's closed or missing, - it creates a new one but DOES NOT close it after execution (unlike asyncio.run). - This allows long-lived async resources to persist across tasks. - - Args: - coro: The coroutine to run. - - Returns: - The return value of the coroutine. - """ - try: - loop = asyncio.get_event_loop() - if loop.is_closed(): - # Loop exists but closed - create new one - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - except RuntimeError: - # No loop in this thread - create new one - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - - return loop.run_until_complete(coro) diff --git a/packages/shared-python/shared/core/celery_app.py b/packages/shared-python/shared/core/celery_app.py index 8d4f2fbb1..83b6534bb 100644 --- a/packages/shared-python/shared/core/celery_app.py +++ b/packages/shared-python/shared/core/celery_app.py @@ -32,7 +32,11 @@ def get_unique_node_name() -> str: # Priority is handled via Celery's Redis transport priority support # (task_queue_max_priority + task_default_priority). celery_app.conf.task_queues = ( - # Knowledge-base queues (routed by priority in task invocation) + # Document Ingestion queues (routed by priority in task invocation) + Queue("document_ingestion_high", routing_key="document_ingestion.high"), + Queue("document_ingestion_medium", routing_key="document_ingestion.medium"), + Queue("document_ingestion_low", routing_key="document_ingestion.low"), + # Pre-rename ingestion queues kept until old broker messages have drained. Queue("kb_high", routing_key="kb.high"), Queue("kb_medium", routing_key="kb.medium"), Queue("kb_low", routing_key="kb.low"), @@ -87,7 +91,10 @@ def get_unique_node_name() -> str: beat_max_loop_interval=30, # Task routing task_routes={ - # Knowledge base tasks (default medium priority) + # Document Ingestion tasks (default medium priority) + "app.core.tasks.document_ingestion_tasks.*": { + "queue": "document_ingestion_medium" + }, "app.core.tasks.kb_tasks.*": {"queue": "kb_medium"}, # Webhook orphan recovery "app.core.tasks.webhook_tasks.recover_orphaned_webhooks": {"queue": "default"}, diff --git a/packages/shared-python/shared/core/celery_router.py b/packages/shared-python/shared/core/celery_router.py index fe2ca4502..eca089cbc 100644 --- a/packages/shared-python/shared/core/celery_router.py +++ b/packages/shared-python/shared/core/celery_router.py @@ -13,6 +13,20 @@ # Note: get_queue_for_job has been simplified and no longer needs direct DB access. +_DOCUMENT_INGESTION_TASK_TYPE = "document_ingestion" +_LEGACY_DOCUMENT_INGESTION_TASK_TYPES = frozenset( + { + "kb_encoding", + "kb_management", + } +) + + +def _normalize_task_type(task_type: str) -> str: + if task_type in _LEGACY_DOCUMENT_INGESTION_TASK_TYPES: + return _DOCUMENT_INGESTION_TASK_TYPE + return task_type + class TaskType(Enum): """Task-type enum.""" @@ -21,12 +35,11 @@ class TaskType(Enum): USER_AUTH = "user_auth" URGENT_DOCUMENT = "urgent_document" DOCUMENT_PROCESSING = "document_processing" - KB_ENCODING = "kb_encoding" + DOCUMENT_INGESTION = "document_ingestion" BATCH_PROCESSING = "batch_processing" ANALYTICS = "analytics" BACKUP = "backup" LOG_PROCESSING = "log_processing" - KB_MANAGEMENT = "kb_management" class UserLevel(Enum): @@ -75,12 +88,11 @@ def __init__(self): TaskType.USER_AUTH: 10, TaskType.URGENT_DOCUMENT: 10, TaskType.DOCUMENT_PROCESSING: 5, - TaskType.KB_ENCODING: 5, + TaskType.DOCUMENT_INGESTION: 6, TaskType.BATCH_PROCESSING: 5, TaskType.ANALYTICS: 1, TaskType.BACKUP: 1, TaskType.LOG_PROCESSING: 1, - TaskType.KB_MANAGEMENT: 6, } # User-level weights. @@ -171,7 +183,7 @@ def create_task_context( """ try: # Parse the task type. - task_type_enum = TaskType(task_type) + task_type_enum = TaskType(_normalize_task_type(task_type)) # Parse the user level. user_level = UserLevel(kwargs.get("user_level", "standard")) @@ -206,7 +218,7 @@ def get_queue_for_job(self, job_type: str, user_id: str) -> str: Resolve the queue name for a job based on job type and subscription. Args: - job_type: Task type such as kb_management or ai_query. + job_type: Task type such as document_ingestion or ai_query. user_id: User ID. Returns: @@ -214,16 +226,10 @@ def get_queue_for_job(self, job_type: str, user_id: str) -> str: """ try: # TODO: temporary simplified path to avoid async work here. - priority_level = 1 # Default free-subscription level. - - # Choose the queue by task type and priority level. - if job_type in ["kb_management", "kb_encoding"]: - if priority_level >= 9: - return "kb_high" - elif priority_level >= 5: - return "kb_medium" - else: - return "kb_low" + if job_type in _LEGACY_DOCUMENT_INGESTION_TASK_TYPES: + return "kb_low" + if job_type == _DOCUMENT_INGESTION_TASK_TYPE: + return "document_ingestion_low" elif job_type in ["ai_query", "user_auth", "urgent_document"]: return "ai_high_priority" elif job_type in ["document_processing"]: @@ -243,8 +249,10 @@ def get_queue_for_job(self, job_type: str, user_id: str) -> str: except Exception as e: logger.error(f"Failed to get queue for user {user_id}: {e}") # Default to a medium-priority queue on failure. - if job_type in ["kb_management", "kb_encoding"]: + if job_type in _LEGACY_DOCUMENT_INGESTION_TASK_TYPES: return "kb_medium" + if job_type == _DOCUMENT_INGESTION_TASK_TYPE: + return "document_ingestion_medium" elif job_type in ["ai_query", "user_auth", "urgent_document"]: return "ai_high_priority" else: @@ -313,7 +321,7 @@ def route_task(self, name, args, kwargs, options, task=None, **kwds): ) # Celery supports router list: try function router first, fallback to static dict -# 1. task_router.route_task: Dynamic routing (based on user subscription, for kb_tasks) +# 1. task_router.route_task: Dynamic routing for Document Ingestion tasks. # 2. static_routes: Static routing (webhook and other fixed routes, from celery_app.py) celery_app.conf.task_routes = [ task_router.route_task, # Dynamic routing priority diff --git a/packages/shared-python/shared/core/config/base.py b/packages/shared-python/shared/core/config/base.py index 0f098cd4e..d0f865eff 100644 --- a/packages/shared-python/shared/core/config/base.py +++ b/packages/shared-python/shared/core/config/base.py @@ -25,7 +25,7 @@ class BaseConfig(BaseSettings): description="Application version read from the APP_VERSION environment variable", ) APP_DESCRIPTION: str = Field( - default="AI-powered document parsing, retrieval, and knowledge access backend", + default="AI-powered document parsing and retrieval backend", description="Application description", ) diff --git a/packages/shared-python/shared/core/config/celery.py b/packages/shared-python/shared/core/config/celery.py index 99541fc62..e22053780 100644 --- a/packages/shared-python/shared/core/config/celery.py +++ b/packages/shared-python/shared/core/config/celery.py @@ -23,9 +23,13 @@ class CeleryConfig(BaseModel): ) # Task retry configuration - KB_TASK_MAX_RETRIES: int = Field(default=2, description="KB task max retries") - KB_TASK_RETRY_COUNTDOWN: int = Field( - default=120, description="KB task retry countdown (seconds)" + DOCUMENT_INGESTION_TASK_MAX_RETRIES: int = Field( + default=2, + description="Document Ingestion task max retries", + ) + DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN: int = Field( + default=120, + description="Document Ingestion task retry countdown (seconds)", ) PYMUPDF_MAX_CONCURRENT: int = Field( default=2, @@ -40,7 +44,7 @@ class CeleryConfig(BaseModel): "user_auth": 10, "urgent_document": 10, "document_processing": 5, - "kb_encoding": 5, + "document_ingestion": 6, "batch_processing": 5, "analytics": 1, "backup": 1, @@ -56,7 +60,7 @@ class CeleryConfig(BaseModel): "user_auth": "auth_queue", "urgent_document": "document_urgent", "document_processing": "document_processing", - "kb_encoding": "kb_encoding", + "document_ingestion": "document_ingestion_low", "batch_processing": "batch_processing", "analytics": "analytics_queue", "backup": "backup_queue", diff --git a/packages/shared-python/shared/core/logging.py b/packages/shared-python/shared/core/logging.py index c2633b806..bb246f18e 100644 --- a/packages/shared-python/shared/core/logging.py +++ b/packages/shared-python/shared/core/logging.py @@ -3,11 +3,12 @@ from contextlib import contextmanager from contextvars import ContextVar from enum import Enum -from typing import Any, Dict +from typing import TYPE_CHECKING, Any, Dict from loguru import logger -from logfire.types import ExceptionCallbackHelper +if TYPE_CHECKING: + from logfire.types import ExceptionCallbackHelper _log_context: ContextVar[Dict[str, Any]] = ContextVar("log_context", default={}) _DEFAULT_CONSOLE_FORMAT = ( diff --git a/packages/shared-python/shared/core/state_machine/config.py b/packages/shared-python/shared/core/state_machine/config.py deleted file mode 100644 index 2128df763..000000000 --- a/packages/shared-python/shared/core/state_machine/config.py +++ /dev/null @@ -1,92 +0,0 @@ -"""State-machine configuration.""" - -from dataclasses import dataclass, field -from typing import Dict - - -DEFAULT_STATE_TIMEOUTS: Dict[str, int] = { - "pending": 300, - "uploading": 600, - "processing": 1800, - "completed": 0, - "failed": 0, -} - - -@dataclass -class StateMachineConfig: - """State-machine settings.""" - - max_retries: int = 3 # Maximum retry count. - base_retry_delay: float = 0.1 # Base retry delay in seconds. - - # Timeout settings used with Redis Keyspace Notifications. - state_timeouts: Dict[str, int] = field( - default_factory=lambda: DEFAULT_STATE_TIMEOUTS.copy() - ) - - # Synchronization settings. - sync_batch_size: int = 100 # Batch size for sync work. - sync_interval: int = 300 # Sync interval in seconds. - - # Maintenance settings. - maintenance_interval: int = 3600 # Maintenance interval in seconds. - cleanup_interval: int = 1800 # Cleanup interval in seconds. - - # Redis Keyspace Notifications support. - enable_keyspace_notifications: bool = True # Enable Keyspace Notifications. - -# Default state-machine configuration. -DEFAULT_CONFIG = StateMachineConfig() - - -def get_state_machine_config() -> StateMachineConfig: - """Return the active state-machine configuration.""" - return DEFAULT_CONFIG - - -def update_state_machine_config( - *, - max_retries: int | None = None, - base_retry_delay: float | None = None, - state_timeouts: Dict[str, int] | None = None, - sync_batch_size: int | None = None, - sync_interval: int | None = None, - maintenance_interval: int | None = None, - cleanup_interval: int | None = None, - enable_keyspace_notifications: bool | None = None, -) -> StateMachineConfig: - """Update and return the active state-machine configuration.""" - global DEFAULT_CONFIG - current = DEFAULT_CONFIG - DEFAULT_CONFIG = StateMachineConfig( - max_retries=current.max_retries if max_retries is None else max_retries, - base_retry_delay=( - current.base_retry_delay - if base_retry_delay is None - else base_retry_delay - ), - state_timeouts=( - current.state_timeouts.copy() - if state_timeouts is None - else state_timeouts - ), - sync_batch_size=( - current.sync_batch_size if sync_batch_size is None else sync_batch_size - ), - sync_interval=current.sync_interval if sync_interval is None else sync_interval, - maintenance_interval=( - current.maintenance_interval - if maintenance_interval is None - else maintenance_interval - ), - cleanup_interval=( - current.cleanup_interval if cleanup_interval is None else cleanup_interval - ), - enable_keyspace_notifications=( - current.enable_keyspace_notifications - if enable_keyspace_notifications is None - else enable_keyspace_notifications - ), - ) - return DEFAULT_CONFIG diff --git a/packages/shared-python/shared/core/state_machine/service.py b/packages/shared-python/shared/core/state_machine/service.py index ee5f421a4..1206313b9 100644 --- a/packages/shared-python/shared/core/state_machine/service.py +++ b/packages/shared-python/shared/core/state_machine/service.py @@ -7,7 +7,6 @@ import asyncio import time -from datetime import datetime, timezone from typing import Any, Dict, Optional from loguru import logger @@ -17,18 +16,29 @@ from shared.core.state_machine.states import ( JobStatus, - is_valid_transition, +) +from shared.core.state_machine.transition_payloads import ( + build_failure_transition_metadata, + build_progress_cache_payload, + build_retry_transition, + serialize_transition_metadata, + utc_now_naive, +) +from shared.core.state_machine.transition_outcome import JobTransitionOutcome +from shared.core.state_machine.transition_runner import ( + MAX_TRANSITION_ATTEMPTS, + TransitionJobSnapshot, + build_cas_conflict_outcome, + build_rollback_exception_outcome, + build_transition_exception_outcome, + get_cas_retry_delay_seconds, + prepare_transition_attempt, + should_retry_cas_conflict, ) from shared.models.database.job import Job from shared.models.database.job_state_audit_log import JobStateAuditLog from shared.services.redis import RedisServiceFactory -from shared.utils.error_details import normalize_error_details -from shared.utils.json_utils import make_json_safe -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder - - -def _utc_now_naive() -> datetime: - return datetime.now(timezone.utc).replace(tzinfo=None) +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class AsyncStateMachineService: @@ -51,23 +61,69 @@ async def transition( auto_commit: bool = True, ) -> bool: """Execute an optimistic-lock state transition with up to 3 retries.""" - max_retries = 3 + outcome = await self.transition_outcome( + db, + job_id, + to_state, + transition_reason, + operator_id, + operator_type, + metadata, + auto_commit=auto_commit, + ) + return outcome.as_bool() + + async def transition_outcome( + self, + db: AsyncSession, + job_id: str, + to_state: str, + transition_reason: str = "normal_transition", + operator_id: Optional[str] = None, + operator_type: str = "system", + metadata: Optional[Dict[str, Any]] = None, + auto_commit: bool = True, + ) -> JobTransitionOutcome: + """Execute a state transition and preserve the reason if it is rejected.""" + max_retries = MAX_TRANSITION_ATTEMPTS for attempt in range(max_retries): + attempts = attempt + 1 try: job = await self._get_job_with_version(db, job_id) - if not job: + prepared = prepare_transition_attempt( + job_id=job_id, + to_state=to_state, + snapshot=( + TransitionJobSnapshot(status=job.status, version=job.version) + if job + else None + ), + attempts=attempts, + ) + if prepared.outcome is not None: + if prepared.outcome.reason == "job_not_found": + logger.error(f"Job {job_id} does not exist") + elif prepared.outcome.reason == "invalid_transition": + logger.warning( + f"Job {job_id}: illegal transition " + f"{prepared.outcome.from_state} → {to_state}, rejected" + ) + return prepared.outcome + + if not prepared.can_write: logger.error(f"Job {job_id} does not exist") - return False - - if not is_valid_transition(job.status, to_state): - logger.warning( - f"Job {job_id}: illegal transition {job.status} → {to_state}, rejected" + return build_transition_exception_outcome( + job_id=job_id, + to_state=to_state, + attempts=attempts, + error="transition precondition did not provide state/version", ) - return False - old_state = job.status - old_version = job.version + old_state = prepared.from_state + old_version = prepared.version + assert old_state is not None + assert old_version is not None # Record audit log (before CAS so the INSERT is within the same tx) await self._record_audit_log( @@ -99,18 +155,28 @@ async def transition( logger.info( f"Job {job_id} state transition: {old_state} → {to_state}" ) - return True + return JobTransitionOutcome.transitioned( + job_id=job_id, + from_state=old_state, + to_state=to_state, + attempts=attempts, + ) # CAS miss — retry with backoff - if attempt < max_retries - 1: + if should_retry_cas_conflict(attempt, max_attempts=max_retries): logger.warning( f"Job {job_id} CAS conflict, retry {attempt + 1}/{max_retries}" ) - await asyncio.sleep(0.1 * (2**attempt)) + await asyncio.sleep(get_cas_retry_delay_seconds(attempt)) continue else: logger.error(f"Job {job_id} CAS retries exhausted") - return False + return build_cas_conflict_outcome( + job_id=job_id, + from_state=old_state, + to_state=to_state, + attempts=attempts, + ) except Exception as e: logger.error(f"Job {job_id} transition failed: {e}") @@ -119,9 +185,25 @@ async def transition( await db.rollback() except Exception as rollback_err: logger.warning(f"Job {job_id} rollback failed: {rollback_err}") - return False + return build_rollback_exception_outcome( + job_id=job_id, + to_state=to_state, + attempts=attempts, + error=rollback_err, + ) + return build_transition_exception_outcome( + job_id=job_id, + to_state=to_state, + attempts=attempts, + error=e, + ) - return False + return build_cas_conflict_outcome( + job_id=job_id, + to_state=to_state, + attempts=max_retries, + from_state=None, + ) async def mark_failed( self, @@ -135,8 +217,39 @@ async def mark_failed( auto_commit: bool = True, ) -> bool: """Mark a job as failed with error information.""" + outcome = await self.mark_failed_outcome( + db, + job_id, + error_message, + error_code=error_code, + error_details=error_details, + operator_id=operator_id, + metadata=metadata, + auto_commit=auto_commit, + ) + return outcome.as_bool() + + async def mark_failed_outcome( + self, + db: AsyncSession, + job_id: str, + error_message: str, + error_code: str = "UNKNOWN", + error_details: Optional[Dict[str, Any]] = None, + operator_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + auto_commit: bool = True, + ) -> JobTransitionOutcome: + """Mark a job as failed and expose why the transition was rejected.""" try: - normalized_details = normalize_error_details(error_details) + normalized_details, transition_metadata = ( + build_failure_transition_metadata( + error_message=error_message, + error_code=error_code, + error_details=error_details, + metadata=metadata, + ) + ) await self._update_job_error( db, job_id, @@ -145,13 +258,7 @@ async def mark_failed( normalized_details, ) - transition_metadata = (metadata or {}).copy() - transition_metadata["error_message"] = error_message - transition_metadata["error_code"] = error_code - if normalized_details: - transition_metadata["error_details"] = normalized_details - - return await self.transition( + return await self.transition_outcome( db, job_id, JobStatus.FAILED.value, @@ -163,7 +270,13 @@ async def mark_failed( ) except Exception as e: logger.error(f"Failed to mark Job {job_id} as failed: {e}") - return False + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=JobStatus.FAILED.value, + reason="transition_exception", + attempts=1, + error_message=str(e), + ) async def mark_completed( self, @@ -174,8 +287,26 @@ async def mark_completed( auto_commit: bool = True, ) -> bool: """Mark a job as completed.""" + outcome = await self.mark_completed_outcome( + db, + job_id, + result_metadata=result_metadata, + operator_id=operator_id, + auto_commit=auto_commit, + ) + return outcome.as_bool() + + async def mark_completed_outcome( + self, + db: AsyncSession, + job_id: str, + result_metadata: Optional[Dict[str, Any]] = None, + operator_id: Optional[str] = None, + auto_commit: bool = True, + ) -> JobTransitionOutcome: + """Mark a job as completed and expose why the transition was rejected.""" try: - return await self.transition( + return await self.transition_outcome( db, job_id, JobStatus.DONE.value, @@ -187,7 +318,13 @@ async def mark_completed( ) except Exception as e: logger.error(f"Failed to mark Job {job_id} as completed: {e}") - return False + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=JobStatus.DONE.value, + reason="transition_exception", + attempts=1, + error_message=str(e), + ) async def handle_retry( self, @@ -208,17 +345,11 @@ async def handle_retry( logger.error(f"Job {job_id} has no status") return False - retry_target = ( - JobStatus.PENDING.value - if current_state == JobStatus.FAILED.value - else current_state + retry_target, retry_metadata = build_retry_transition( + current_state=current_state, + retry_metadata=retry_metadata, ) - retry_metadata = retry_metadata or {} - retry_metadata["retry_reason"] = "task_retry" - retry_metadata["retry_timestamp"] = str(int(time.time())) - retry_metadata["retry_count"] = retry_metadata.get("retry_count", 0) + 1 - # Always use full transition() for CAS protection — even same-state return await self.transition( db, @@ -291,7 +422,7 @@ async def _cas_update_state( .values( status=to_state, version=old_version + 1, - updated_at=_utc_now_naive(), + updated_at=utc_now_naive(), ) ) return result.rowcount > 0 @@ -307,13 +438,7 @@ async def _record_audit_log( operator_type: str, metadata: Optional[Dict[str, Any]], ) -> None: - serialized = None - if metadata: - try: - serialized = make_json_safe(metadata) - except Exception as e: - logger.warning(f"Metadata serialization failed: {e}") - serialized = {"error": "metadata_serialization_failed"} + serialized = serialize_transition_metadata(metadata) db.add( JobStateAuditLog( @@ -392,13 +517,20 @@ async def _update_redis_cache( ) progress_key = redis_key_builder.task_progress(job_id) - progress_data: Dict[str, Any] = { - "status": status, - "timestamp": str(int(time.time())), - } + progress_data: Dict[str, Any] = build_progress_cache_payload( + status=status, + metadata=None, + timestamp=int(time.time()), + ) if metadata: try: - progress_data.update(make_json_safe(metadata)) + progress_data.update( + build_progress_cache_payload( + status=status, + metadata=metadata, + timestamp=int(time.time()), + ) + ) except Exception as e: logger.warning(f"Metadata serialization skipped: {e}") diff --git a/packages/shared-python/shared/core/state_machine/service_sync.py b/packages/shared-python/shared/core/state_machine/service_sync.py index 88e9aa3c6..6d8ba1cc1 100644 --- a/packages/shared-python/shared/core/state_machine/service_sync.py +++ b/packages/shared-python/shared/core/state_machine/service_sync.py @@ -7,7 +7,6 @@ """ import time -from datetime import datetime, timezone from typing import Any, Dict, Optional from loguru import logger @@ -16,18 +15,29 @@ from shared.core.state_machine.states import ( JobStatus, - is_valid_transition, +) +from shared.core.state_machine.transition_payloads import ( + build_failure_transition_metadata, + build_progress_cache_payload, + build_retry_transition, + serialize_transition_metadata, + utc_now_naive, +) +from shared.core.state_machine.transition_outcome import JobTransitionOutcome +from shared.core.state_machine.transition_runner import ( + MAX_TRANSITION_ATTEMPTS, + TransitionJobSnapshot, + build_cas_conflict_outcome, + build_rollback_exception_outcome, + build_transition_exception_outcome, + get_cas_retry_delay_seconds, + prepare_transition_attempt, + should_retry_cas_conflict, ) from shared.models.database.job import Job from shared.models.database.job_state_audit_log import JobStateAuditLog from shared.services.redis.redis_sync_service import SyncRedisServiceFactory -from shared.utils.error_details import normalize_error_details -from shared.utils.json_utils import make_json_safe -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder - - -def _utc_now_naive() -> datetime: - return datetime.now(timezone.utc).replace(tzinfo=None) +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class SyncStateMachineService: @@ -49,23 +59,66 @@ def transition( metadata: Optional[Dict[str, Any]] = None, ) -> bool: """Execute an optimistic-lock state transition with up to 3 retries.""" - max_retries = 3 + return self.transition_outcome( + db, + job_id, + to_state, + transition_reason, + operator_id, + operator_type, + metadata, + ).as_bool() + + def transition_outcome( + self, + db: Session, + job_id: str, + to_state: str, + transition_reason: str = "normal_transition", + operator_id: Optional[str] = None, + operator_type: str = "system", + metadata: Optional[Dict[str, Any]] = None, + ) -> JobTransitionOutcome: + """Execute a state transition and preserve the reason if it is rejected.""" + max_retries = MAX_TRANSITION_ATTEMPTS for attempt in range(max_retries): + attempts = attempt + 1 try: job = self._get_job_with_version(db, job_id) - if not job: + prepared = prepare_transition_attempt( + job_id=job_id, + to_state=to_state, + snapshot=( + TransitionJobSnapshot(status=job.status, version=job.version) + if job + else None + ), + attempts=attempts, + ) + if prepared.outcome is not None: + if prepared.outcome.reason == "job_not_found": + logger.error(f"Job {job_id} does not exist") + elif prepared.outcome.reason == "invalid_transition": + logger.warning( + f"Job {job_id}: illegal transition " + f"{prepared.outcome.from_state} → {to_state}, rejected" + ) + return prepared.outcome + + if not prepared.can_write: logger.error(f"Job {job_id} does not exist") - return False - - if not is_valid_transition(job.status, to_state): - logger.warning( - f"Job {job_id}: illegal transition {job.status} → {to_state}, rejected" + return build_transition_exception_outcome( + job_id=job_id, + to_state=to_state, + attempts=attempts, + error="transition precondition did not provide state/version", ) - return False - old_state = job.status - old_version = job.version + old_state = prepared.from_state + old_version = prepared.version + assert old_state is not None + assert old_version is not None self._record_audit_log( db, @@ -85,19 +138,29 @@ def transition( logger.info( f"Job {job_id} state transition: {old_state} → {to_state}" ) - return True + return JobTransitionOutcome.transitioned( + job_id=job_id, + from_state=old_state, + to_state=to_state, + attempts=attempts, + ) - if attempt < max_retries - 1: + if should_retry_cas_conflict(attempt, max_attempts=max_retries): logger.warning( f"Job {job_id} CAS conflict, retry {attempt + 1}/{max_retries}" ) import gevent - gevent.sleep(0.1 * (2**attempt)) + gevent.sleep(get_cas_retry_delay_seconds(attempt)) continue else: logger.error(f"Job {job_id} CAS retries exhausted") - return False + return build_cas_conflict_outcome( + job_id=job_id, + from_state=old_state, + to_state=to_state, + attempts=attempts, + ) except Exception as e: logger.error(f"Job {job_id} transition failed: {e}") @@ -106,9 +169,25 @@ def transition( db.rollback() except Exception as rollback_err: logger.warning(f"Job {job_id} rollback failed: {rollback_err}") - return False + return build_rollback_exception_outcome( + job_id=job_id, + to_state=to_state, + attempts=attempts, + error=rollback_err, + ) + return build_transition_exception_outcome( + job_id=job_id, + to_state=to_state, + attempts=attempts, + error=e, + ) - return False + return build_cas_conflict_outcome( + job_id=job_id, + to_state=to_state, + attempts=max_retries, + from_state=None, + ) def mark_failed( self, @@ -121,8 +200,36 @@ def mark_failed( metadata: Optional[Dict[str, Any]] = None, ) -> bool: """Mark a job as failed with error information.""" + return self.mark_failed_outcome( + db, + job_id, + error_message, + error_code=error_code, + error_details=error_details, + operator_id=operator_id, + metadata=metadata, + ).as_bool() + + def mark_failed_outcome( + self, + db: Session, + job_id: str, + error_message: str, + error_code: str = "UNKNOWN", + error_details: Optional[Dict[str, Any]] = None, + operator_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> JobTransitionOutcome: + """Mark a job as failed and expose why the transition was rejected.""" try: - normalized_details = normalize_error_details(error_details) + normalized_details, transition_metadata = ( + build_failure_transition_metadata( + error_message=error_message, + error_code=error_code, + error_details=error_details, + metadata=metadata, + ) + ) self._update_job_error( db, job_id, @@ -131,13 +238,7 @@ def mark_failed( normalized_details, ) - transition_metadata = (metadata or {}).copy() - transition_metadata["error_message"] = error_message - transition_metadata["error_code"] = error_code - if normalized_details: - transition_metadata["error_details"] = normalized_details - - return self.transition( + return self.transition_outcome( db, job_id, JobStatus.FAILED.value, @@ -148,7 +249,13 @@ def mark_failed( ) except Exception as e: logger.error(f"Failed to mark Job {job_id} as failed: {e}") - return False + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=JobStatus.FAILED.value, + reason="transition_exception", + attempts=1, + error_message=str(e), + ) def mark_completed( self, @@ -158,8 +265,23 @@ def mark_completed( operator_id: Optional[str] = None, ) -> bool: """Mark a job as completed.""" + return self.mark_completed_outcome( + db, + job_id, + result_metadata=result_metadata, + operator_id=operator_id, + ).as_bool() + + def mark_completed_outcome( + self, + db: Session, + job_id: str, + result_metadata: Optional[Dict[str, Any]] = None, + operator_id: Optional[str] = None, + ) -> JobTransitionOutcome: + """Mark a job as completed and expose why the transition was rejected.""" try: - return self.transition( + return self.transition_outcome( db, job_id, JobStatus.DONE.value, @@ -170,6 +292,54 @@ def mark_completed( ) except Exception as e: logger.error(f"Failed to mark Job {job_id} as completed: {e}") + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=JobStatus.DONE.value, + reason="transition_exception", + attempts=1, + error_message=str(e), + ) + + def handle_retry( + self, + db: Session, + job_id: str, + retry_metadata: Optional[Dict[str, Any]] = None, + operator_id: Optional[str] = None, + ) -> bool: + """Handle task retry — always goes through CAS-protected transition.""" + try: + job = self._get_job(db, job_id) + if not job: + logger.error(f"Job {job_id} does not exist") + return False + + current_state = job.status + if not current_state: + logger.error(f"Job {job_id} has no status") + return False + + retry_target, retry_metadata = build_retry_transition( + current_state=current_state, + retry_metadata=retry_metadata, + ) + + return self.transition( + db, + job_id, + retry_target, + "retry_transition", + operator_id, + "retry", + retry_metadata, + ) + except Exception as e: + logger.error(f"Job {job_id} retry failed: {e}") + try: + if db.is_active: + db.rollback() + except Exception as rollback_err: + logger.warning(f"Job {job_id} rollback failed: {rollback_err}") return False # ── Private helpers ───────────────────────────────────────────────── @@ -199,7 +369,7 @@ def _cas_update_state( .values( status=to_state, version=old_version + 1, - updated_at=_utc_now_naive(), + updated_at=utc_now_naive(), ) ) return result.rowcount > 0 @@ -215,13 +385,7 @@ def _record_audit_log( operator_type: str, metadata: Optional[Dict[str, Any]], ) -> None: - serialized = None - if metadata: - try: - serialized = make_json_safe(metadata) - except Exception as e: - logger.warning(f"Metadata serialization failed: {e}") - serialized = {"error": "metadata_serialization_failed"} + serialized = serialize_transition_metadata(metadata) db.add( JobStateAuditLog( @@ -281,13 +445,20 @@ def _update_redis_cache( ) progress_key = redis_key_builder.task_progress(job_id) - progress_data: Dict[str, Any] = { - "status": status, - "timestamp": str(int(time.time())), - } + progress_data: Dict[str, Any] = build_progress_cache_payload( + status=status, + metadata=None, + timestamp=int(time.time()), + ) if metadata: try: - progress_data.update(make_json_safe(metadata)) + progress_data.update( + build_progress_cache_payload( + status=status, + metadata=metadata, + timestamp=int(time.time()), + ) + ) except Exception as e: logger.warning(f"Metadata serialization skipped: {e}") diff --git a/packages/shared-python/shared/core/state_machine/transition_outcome.py b/packages/shared-python/shared/core/state_machine/transition_outcome.py new file mode 100644 index 000000000..4caca8b28 --- /dev/null +++ b/packages/shared-python/shared/core/state_machine/transition_outcome.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +TransitionOutcomeReason = Literal[ + "transitioned", + "job_not_found", + "invalid_transition", + "cas_conflict", + "transition_exception", + "rollback_exception", +] + + +@dataclass(frozen=True) +class JobTransitionOutcome: + job_id: str + to_state: str + reason: TransitionOutcomeReason + attempts: int + from_state: str | None = None + error_message: str | None = None + + @property + def succeeded(self) -> bool: + return self.reason == "transitioned" + + def as_bool(self) -> bool: + return self.succeeded + + @classmethod + def transitioned( + cls, + *, + job_id: str, + from_state: str, + to_state: str, + attempts: int, + ) -> JobTransitionOutcome: + return cls( + job_id=job_id, + from_state=from_state, + to_state=to_state, + reason="transitioned", + attempts=attempts, + ) + + @classmethod + def rejected( + cls, + *, + job_id: str, + to_state: str, + reason: TransitionOutcomeReason, + attempts: int, + from_state: str | None = None, + error_message: str | None = None, + ) -> JobTransitionOutcome: + return cls( + job_id=job_id, + from_state=from_state, + to_state=to_state, + reason=reason, + attempts=attempts, + error_message=error_message, + ) diff --git a/packages/shared-python/shared/core/state_machine/transition_payloads.py b/packages/shared-python/shared/core/state_machine/transition_payloads.py new file mode 100644 index 000000000..c2a2cf58b --- /dev/null +++ b/packages/shared-python/shared/core/state_machine/transition_payloads.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any + +from shared.core.state_machine.states import JobStatus +from shared.utils.error_details import normalize_error_details +from shared.utils.json_utils import make_json_safe + + +def utc_now_naive() -> datetime: + return datetime.now(timezone.utc).replace(tzinfo=None) + + +def serialize_transition_metadata( + metadata: dict[str, Any] | None, +) -> dict[str, Any] | None: + if not metadata: + return None + + try: + return make_json_safe(metadata) + except Exception: + return {"error": "metadata_serialization_failed"} + + +def build_failure_transition_metadata( + *, + error_message: str, + error_code: str, + error_details: dict[str, Any] | None, + metadata: dict[str, Any] | None, +) -> tuple[dict[str, Any] | None, dict[str, Any]]: + normalized_details = normalize_error_details(error_details) + transition_metadata = (metadata or {}).copy() + transition_metadata["error_message"] = error_message + transition_metadata["error_code"] = error_code + if normalized_details: + transition_metadata["error_details"] = normalized_details + return normalized_details, transition_metadata + + +def build_retry_transition( + *, + current_state: str, + retry_metadata: dict[str, Any] | None, +) -> tuple[str, dict[str, Any]]: + retry_target = ( + JobStatus.PENDING.value + if current_state == JobStatus.FAILED.value + else current_state + ) + + resolved_metadata = retry_metadata or {} + resolved_metadata["retry_reason"] = "task_retry" + resolved_metadata["retry_timestamp"] = str(int(datetime.now(timezone.utc).timestamp())) + resolved_metadata["retry_count"] = resolved_metadata.get("retry_count", 0) + 1 + return retry_target, resolved_metadata + + +def build_progress_cache_payload( + *, + status: str, + metadata: dict[str, Any] | None, + timestamp: int, +) -> dict[str, Any]: + progress_data: dict[str, Any] = { + "status": status, + "timestamp": str(timestamp), + } + if metadata: + progress_data.update(make_json_safe(metadata)) + return progress_data diff --git a/packages/shared-python/shared/core/state_machine/transition_runner.py b/packages/shared-python/shared/core/state_machine/transition_runner.py new file mode 100644 index 000000000..199ee7e59 --- /dev/null +++ b/packages/shared-python/shared/core/state_machine/transition_runner.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from shared.core.state_machine.states import is_valid_transition +from shared.core.state_machine.transition_outcome import JobTransitionOutcome + +MAX_TRANSITION_ATTEMPTS = 3 + + +@dataclass(frozen=True) +class TransitionJobSnapshot: + status: str + version: int + + +@dataclass(frozen=True) +class PreparedTransitionAttempt: + outcome: JobTransitionOutcome | None + from_state: str | None = None + version: int | None = None + + @property + def can_write(self) -> bool: + return self.outcome is None and self.from_state is not None and self.version is not None + + +def prepare_transition_attempt( + *, + job_id: str, + to_state: str, + snapshot: TransitionJobSnapshot | None, + attempts: int, +) -> PreparedTransitionAttempt: + if snapshot is None: + return PreparedTransitionAttempt( + outcome=JobTransitionOutcome.rejected( + job_id=job_id, + to_state=to_state, + reason="job_not_found", + attempts=attempts, + ) + ) + + if not is_valid_transition(snapshot.status, to_state): + return PreparedTransitionAttempt( + outcome=JobTransitionOutcome.rejected( + job_id=job_id, + from_state=snapshot.status, + to_state=to_state, + reason="invalid_transition", + attempts=attempts, + ) + ) + + return PreparedTransitionAttempt( + outcome=None, + from_state=snapshot.status, + version=snapshot.version, + ) + + +def should_retry_cas_conflict( + attempt_index: int, + *, + max_attempts: int = MAX_TRANSITION_ATTEMPTS, +) -> bool: + return attempt_index < max_attempts - 1 + + +def get_cas_retry_delay_seconds(attempt_index: int) -> float: + return 0.1 * (2**attempt_index) + + +def build_cas_conflict_outcome( + *, + job_id: str, + from_state: str | None, + to_state: str, + attempts: int, +) -> JobTransitionOutcome: + return JobTransitionOutcome.rejected( + job_id=job_id, + from_state=from_state, + to_state=to_state, + reason="cas_conflict", + attempts=attempts, + ) + + +def build_transition_exception_outcome( + *, + job_id: str, + to_state: str, + attempts: int, + error: Exception | str, +) -> JobTransitionOutcome: + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=to_state, + reason="transition_exception", + attempts=attempts, + error_message=str(error), + ) + + +def build_rollback_exception_outcome( + *, + job_id: str, + to_state: str, + attempts: int, + error: Exception | str, +) -> JobTransitionOutcome: + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=to_state, + reason="rollback_exception", + attempts=attempts, + error_message=str(error), + ) diff --git a/packages/shared-python/shared/core/tasks/__init__.py b/packages/shared-python/shared/core/tasks/__init__.py index 849c9f7b6..95e7349b9 100644 --- a/packages/shared-python/shared/core/tasks/__init__.py +++ b/packages/shared-python/shared/core/tasks/__init__.py @@ -1,5 +1,5 @@ """Shared Celery task definitions. -Knowledge-base tasks moved to the Worker service, and state-machine tasks moved -to the API service. Only generic Celery tasks remain here. +Document Ingestion tasks live in the Worker service, and state-machine tasks +live in the API service. Only generic Celery tasks remain here. """ diff --git a/packages/shared-python/shared/models/database/job.py b/packages/shared-python/shared/models/database/job.py index 76c9f1df1..8a4fb6bf3 100644 --- a/packages/shared-python/shared/models/database/job.py +++ b/packages/shared-python/shared/models/database/job.py @@ -6,8 +6,7 @@ from datetime import datetime -# Forward references avoid circular imports. -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import Any, Dict, Optional from uuid import uuid4 from sqlalchemy import ( @@ -28,13 +27,6 @@ from shared.utils.utc_now import utc_now_naive -if TYPE_CHECKING: - from shared.models.database.job_result import JobResult - from shared.models.database.job_state_audit_log import JobStateAuditLog - from shared.models.database.job_state_history import JobStateHistory - from shared.models.database.webhook_log import WebhookLog - - class Job(Base): """Job Model - User API Business Task""" @@ -51,7 +43,7 @@ class Job(Base): ) # Basic Job Info - job_type: Mapped[str] = mapped_column(String(50), nullable=False) # kb_management + job_type: Mapped[str] = mapped_column(String(50), nullable=False) # document_ingestion status: Mapped[str] = mapped_column( String(50), nullable=False, default="pending" ) # PRD Status: pending, waiting-file, running, converting, done, failed @@ -113,22 +105,22 @@ class Job(Base): # Relationships — default to noload to prevent implicit SELECTs. # Use explicit selectinload() in queries that need related data. - state_history: Mapped[list["JobStateHistory"]] = relationship( + state_history: Mapped[list[Any]] = relationship( "JobStateHistory", back_populates="job", cascade="all, delete-orphan", lazy="noload", ) - state_audit_logs: Mapped[list["JobStateAuditLog"]] = relationship( + state_audit_logs: Mapped[list[Any]] = relationship( "JobStateAuditLog", back_populates="job", cascade="all, delete-orphan", lazy="noload", ) - webhook_logs: Mapped[list["WebhookLog"]] = relationship( + webhook_logs: Mapped[list[Any]] = relationship( "WebhookLog", back_populates="job", cascade="all, delete-orphan", lazy="noload" ) - job_result: Mapped[Optional["JobResult"]] = relationship( + job_result: Mapped[Optional[Any]] = relationship( "JobResult", back_populates="job", uselist=False, lazy="noload" ) diff --git a/packages/shared-python/shared/models/database/job_result.py b/packages/shared-python/shared/models/database/job_result.py index c6d4d9e73..f88f3b98b 100644 --- a/packages/shared-python/shared/models/database/job_result.py +++ b/packages/shared-python/shared/models/database/job_result.py @@ -3,7 +3,7 @@ from __future__ import annotations from datetime import datetime -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import Any, Dict, List, Optional from uuid import uuid4 from sqlalchemy import JSON, DateTime, ForeignKey, Index, Integer, String, Text @@ -12,9 +12,6 @@ from shared.core.database import Base from shared.utils.utc_now import utc_now_naive -if TYPE_CHECKING: - from shared.models.database.job import Job - class JobResult(Base): """Primary job-result table.""" @@ -57,7 +54,7 @@ class JobResult(Base): nullable=False, ) - job: Mapped["Job"] = relationship( + job: Mapped[Any] = relationship( "Job", back_populates="job_result", lazy="joined" ) chunks: Mapped[List["JobChunk"]] = relationship( diff --git a/packages/shared-python/shared/models/database/job_state_audit_log.py b/packages/shared-python/shared/models/database/job_state_audit_log.py index 970b2af57..c2edc4ade 100644 --- a/packages/shared-python/shared/models/database/job_state_audit_log.py +++ b/packages/shared-python/shared/models/database/job_state_audit_log.py @@ -3,7 +3,7 @@ from __future__ import annotations from datetime import datetime -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import Any, Dict, Optional from sqlalchemy import JSON, DateTime, ForeignKey, Index, Integer, String from sqlalchemy.orm import Mapped, mapped_column, relationship @@ -11,9 +11,6 @@ from shared.core.database import Base from shared.utils.utc_now import utc_now_naive -if TYPE_CHECKING: - from shared.models.database.job import Job - class JobStateAuditLog(Base): """Job state-transition audit log.""" @@ -58,7 +55,7 @@ class JobStateAuditLog(Base): ) # Relationships. - job: Mapped["Job"] = relationship( + job: Mapped[Any] = relationship( "Job", back_populates="state_audit_logs", lazy="select" ) diff --git a/packages/shared-python/shared/models/database/job_state_history.py b/packages/shared-python/shared/models/database/job_state_history.py index b480ae3f7..a564f160b 100644 --- a/packages/shared-python/shared/models/database/job_state_history.py +++ b/packages/shared-python/shared/models/database/job_state_history.py @@ -3,7 +3,7 @@ from __future__ import annotations from datetime import datetime -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import Any, Dict, Optional from uuid import uuid4 from sqlalchemy import JSON, DateTime, ForeignKey, Index, String @@ -12,9 +12,6 @@ from shared.core.database import Base from shared.utils.utc_now import utc_now_naive -if TYPE_CHECKING: - from shared.models.database.job import Job - class JobStateHistory(Base): """Job state-history model for individual state-machine transitions.""" @@ -44,7 +41,7 @@ class JobStateHistory(Base): ) # Relationships. - job: Mapped["Job"] = relationship("Job", back_populates="state_history") + job: Mapped[Any] = relationship("Job", back_populates="state_history") # Indexes. __table_args__ = ( diff --git a/packages/shared-python/shared/models/database/webhook.py b/packages/shared-python/shared/models/database/webhook.py index 09f7f4dd1..73dc31dc8 100644 --- a/packages/shared-python/shared/models/database/webhook.py +++ b/packages/shared-python/shared/models/database/webhook.py @@ -8,7 +8,7 @@ from __future__ import annotations from datetime import datetime -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import Any, Dict, Optional from uuid import uuid4 from sqlalchemy import DateTime, ForeignKey, Index, Integer, String @@ -18,10 +18,6 @@ from shared.core.database import Base from shared.utils.utc_now import utc_now_naive -if TYPE_CHECKING: - from shared.models.database.job import Job - from shared.models.database.webhook_log import WebhookLog - class WebhookEventStatus: """Webhook event status constants.""" @@ -83,8 +79,8 @@ class WebhookEvent(Base): ) # Relationships - job: Mapped["Job"] = relationship("Job", lazy="select") - deliveries: Mapped[list["WebhookLog"]] = relationship( + job: Mapped[Any] = relationship("Job", lazy="select") + deliveries: Mapped[list[Any]] = relationship( "WebhookLog", back_populates="event", cascade="all, delete-orphan", # how about the manual trigger? diff --git a/packages/shared-python/shared/models/database/webhook_log.py b/packages/shared-python/shared/models/database/webhook_log.py index 2beb6dca2..cd664680a 100644 --- a/packages/shared-python/shared/models/database/webhook_log.py +++ b/packages/shared-python/shared/models/database/webhook_log.py @@ -7,7 +7,7 @@ from __future__ import annotations from datetime import datetime -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import Any, Dict, Optional from uuid import uuid4 from sqlalchemy import JSON, DateTime, ForeignKey, Index, Integer, String, Text @@ -16,9 +16,6 @@ from shared.core.database import Base from shared.utils.utc_now import utc_now_naive -if TYPE_CHECKING: - from shared.models.database.job import Job - from shared.models.database.webhook import WebhookEvent class WebhookLog(Base): """Webhook Log Model - Records webhook delivery history.""" @@ -75,8 +72,8 @@ class WebhookLog(Base): ) # Relationships - job: Mapped["Job"] = relationship("Job", back_populates="webhook_logs") - event: Mapped[Optional["WebhookEvent"]] = relationship( + job: Mapped[Any] = relationship("Job", back_populates="webhook_logs") + event: Mapped[Optional[Any]] = relationship( "WebhookEvent", back_populates="deliveries", foreign_keys=[event_id] ) diff --git a/packages/shared-python/shared/models/schemas/job.py b/packages/shared-python/shared/models/schemas/job.py index 20bad6763..a2b686551 100644 --- a/packages/shared-python/shared/models/schemas/job.py +++ b/packages/shared-python/shared/models/schemas/job.py @@ -17,9 +17,6 @@ class ParsingParams(BaseModel): model: Literal["base", "advanced"] = Field("base", description="Parsing model") ocr_enabled: bool = Field(False, description="Enable OCR") - kb_dir: Optional[str] = Field( - "Default_Root", description="Knowledge-base directory" - ) doc_type: Literal["auto", "pdf", "docx", "txt", "md"] = Field( "auto", description="Document type" ) @@ -45,7 +42,9 @@ class JobCreate(BaseModel): """Request payload for creating a job.""" namespace: Optional[str] = Field( - None, description="Retrieval namespace; defaults to default" + None, + max_length=255, + description="Retrieval namespace; defaults to default", ) document_id: Optional[str] = Field( None, description="Existing document ID for update flows" diff --git a/packages/shared-python/shared/models/schemas/job_metadata.py b/packages/shared-python/shared/models/schemas/job_metadata.py index 79cfbff81..71aa4d2dc 100644 --- a/packages/shared-python/shared/models/schemas/job_metadata.py +++ b/packages/shared-python/shared/models/schemas/job_metadata.py @@ -4,6 +4,8 @@ from pydantic import BaseModel, ConfigDict, Field +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace + class JobMetadataBase(BaseModel): """Base schema for stored job metadata.""" @@ -39,9 +41,10 @@ class JobMetadataHelper: @staticmethod def create_from_request(request, **kwargs) -> Dict[str, Any]: """Build metadata from a JobCreate request without embedding user_config.""" + namespace = normalize_retrieval_namespace(request.namespace) metadata = { "original_request": request.model_dump(), - "namespace": request.namespace or "default", + "namespace": namespace, "document_id": request.document_id, "parsing_params": ( request.parsing_params.model_dump() if request.parsing_params else None @@ -52,6 +55,35 @@ def create_from_request(request, **kwargs) -> Dict[str, Any]: metadata.update(kwargs) return metadata + @staticmethod + def set_document_scope( + metadata: Dict[str, Any], + *, + document_id: str, + namespace: str, + ) -> None: + """Store the effective retrieval document scope.""" + metadata["document_id"] = document_id + metadata["namespace"] = namespace + + @staticmethod + def set_file_source(metadata: Dict[str, Any], *, source_file_name: str) -> None: + """Store source metadata for direct file uploads.""" + metadata["source_file_name"] = source_file_name + metadata["source_type"] = "file" + + @staticmethod + def set_url_source( + metadata: Dict[str, Any], + *, + source_file_name: str, + source_url: str, + ) -> None: + """Store source metadata for URL ingestion.""" + metadata["source_file_name"] = source_file_name + metadata["source_url"] = source_url + metadata["source_type"] = "url" + @staticmethod def get_field( metadata: Optional[Dict[str, Any]], field: str, default: Any = None @@ -61,6 +93,54 @@ def get_field( return default return metadata.get(field, default) + @staticmethod + def get_string_field( + metadata: Optional[Dict[str, Any]], field: str, default: str | None = None + ) -> str | None: + """Read a string field from metadata.""" + value = JobMetadataHelper.get_field(metadata, field, default) + return value if isinstance(value, str) else default + + @staticmethod + def get_original_request(metadata: Optional[Dict[str, Any]]) -> Dict[str, Any]: + """Return the stored creation request payload.""" + original_request = JobMetadataHelper.get_field(metadata, "original_request", {}) + return original_request if isinstance(original_request, dict) else {} + + @staticmethod + def get_parsing_params_dict(metadata: Optional[Dict[str, Any]]) -> Dict[str, Any]: + """Return stored parsing parameters as a dictionary.""" + parsing_params = JobMetadataHelper.get_field(metadata, "parsing_params", {}) + return parsing_params if isinstance(parsing_params, dict) else {} + + @staticmethod + def get_namespace( + metadata: Optional[Dict[str, Any]], default: str | None = None + ) -> str | None: + """Return the retrieval namespace stored in metadata.""" + namespace = JobMetadataHelper.get_string_field(metadata, "namespace", default) + return normalize_retrieval_namespace(namespace) if namespace is not None else None + + @staticmethod + def get_document_id(metadata: Optional[Dict[str, Any]]) -> str | None: + """Return the retrieval document id stored in metadata.""" + return JobMetadataHelper.get_string_field(metadata, "document_id") + + @staticmethod + def get_data_id(metadata: Optional[Dict[str, Any]]) -> str | None: + """Return the user-defined data id stored in metadata.""" + return JobMetadataHelper.get_string_field(metadata, "data_id") + + @staticmethod + def get_source_file_name(metadata: Optional[Dict[str, Any]]) -> str | None: + """Return the source file name stored in metadata.""" + return JobMetadataHelper.get_string_field(metadata, "source_file_name") + + @staticmethod + def get_source_url(metadata: Optional[Dict[str, Any]]) -> str | None: + """Return the source URL stored in metadata.""" + return JobMetadataHelper.get_string_field(metadata, "source_url") + @staticmethod def get_parsing_param( metadata: Optional[Dict[str, Any]], param: str, default: Any = None diff --git a/packages/shared-python/shared/models/schemas/retrieval_namespace.py b/packages/shared-python/shared/models/schemas/retrieval_namespace.py new file mode 100644 index 000000000..ef2be2349 --- /dev/null +++ b/packages/shared-python/shared/models/schemas/retrieval_namespace.py @@ -0,0 +1,9 @@ +from __future__ import annotations + +_DEFAULT_RETRIEVAL_NAMESPACE = "default" + + +def normalize_retrieval_namespace(namespace: str | None) -> str: + """Return the canonical namespace value used by jobs, retrieval, and caches.""" + normalized = str(namespace or "").strip() + return normalized or _DEFAULT_RETRIEVAL_NAMESPACE diff --git a/packages/shared-python/shared/models/schemas/s3_file.py b/packages/shared-python/shared/models/schemas/s3_file.py deleted file mode 100644 index 4b3a0cbbc..000000000 --- a/packages/shared-python/shared/models/schemas/s3_file.py +++ /dev/null @@ -1,8 +0,0 @@ -from pydantic import BaseModel - - -class FliesDownload(BaseModel): - message: str - file_key: str - download_url: str - expires_in_seconds: int diff --git a/packages/shared-python/shared/utils/ali_quota_manager.py b/packages/shared-python/shared/services/ai/ali_quota_manager.py similarity index 96% rename from packages/shared-python/shared/utils/ali_quota_manager.py rename to packages/shared-python/shared/services/ai/ali_quota_manager.py index 535b86c0c..37bcc2e78 100644 --- a/packages/shared-python/shared/utils/ali_quota_manager.py +++ b/packages/shared-python/shared/services/ai/ali_quota_manager.py @@ -14,7 +14,7 @@ SyncRedisService, SyncRedisServiceFactory, ) -from shared.utils.quota_manager import BaseQuotaManager, TokenConfig +from shared.services.quota.token_pool import BaseQuotaManager, TokenConfig class AliQuotaManager(BaseQuotaManager): diff --git a/packages/shared-python/shared/utils/iloveapi_quota_manager.py b/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py similarity index 97% rename from packages/shared-python/shared/utils/iloveapi_quota_manager.py rename to packages/shared-python/shared/services/ai/iloveapi_quota_manager.py index b3b923a41..510719008 100644 --- a/packages/shared-python/shared/utils/iloveapi_quota_manager.py +++ b/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py @@ -26,7 +26,7 @@ SyncRedisService, SyncRedisServiceFactory, ) -from shared.utils.quota_manager import BaseQuotaManager, TokenConfig +from shared.services.quota.token_pool import BaseQuotaManager, TokenConfig class ILoveApiQuotaManager(BaseQuotaManager): @@ -124,8 +124,8 @@ def parse_tokens_from_settings() -> List[TokenConfig]: ), ) ) - except json.JSONDecodeError: - pass + except json.JSONDecodeError as exc: + logger.debug(f"Invalid I Love API token JSON ignored: {exc}") if not specs and legacy_pub and legacy_sec: specs.append( diff --git a/packages/shared-python/shared/utils/llm_mock.py b/packages/shared-python/shared/services/ai/llm_mock.py similarity index 100% rename from packages/shared-python/shared/utils/llm_mock.py rename to packages/shared-python/shared/services/ai/llm_mock.py diff --git a/packages/shared-python/shared/utils/OpenAICompatibleClientSync.py b/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py similarity index 96% rename from packages/shared-python/shared/utils/OpenAICompatibleClientSync.py rename to packages/shared-python/shared/services/ai/openai_compatible_client_sync.py index 66da76a6d..79a00e97a 100644 --- a/packages/shared-python/shared/utils/OpenAICompatibleClientSync.py +++ b/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py @@ -17,8 +17,8 @@ from shared.core.config import settings from shared.core.exceptions.domain_exceptions import LLMServiceException -from shared.utils.http_clients import get_sync_client -from shared.utils.llm_mock import build_mock_chat_completion_response +from shared.services.http.client_pool import get_sync_client +from shared.services.ai.llm_mock import build_mock_chat_completion_response from shared.utils.security_utils import mask_api_key LOCAL_DEBUG = os.getenv("LOCAL_DEBUG", "0") == "1" @@ -185,7 +185,7 @@ def _make_ali_pool_call( api_kwargs: Dict[str, Any], ) -> tuple[str, LLMUsage]: """Acquire a token, make the call, and retry inline on 429.""" - from shared.utils.ali_quota_manager import get_ali_quota_manager + from shared.services.ai.ali_quota_manager import get_ali_quota_manager quota_manager = get_ali_quota_manager() base_url: Optional[str] = self._base_url @@ -390,11 +390,13 @@ def _parse_retry_after(exc: openai.RateLimitError) -> int: """Extract Retry-After seconds from a RateLimitError, with sane bounds.""" try: if hasattr(exc, "response") and exc.response is not None: - header_value = exc.response.headers.get("retry-after") or exc.response.headers.get("Retry-After") + header_value = exc.response.headers.get( + "retry-after" + ) or exc.response.headers.get("Retry-After") if header_value: return max(1, min(int(header_value), 120)) - except (ValueError, TypeError, AttributeError): - pass + except (ValueError, TypeError, AttributeError) as parse_error: + logger.debug(f"Could not parse retry-after header: {parse_error}") return settings.ALI_TOKEN_COOLDOWN_SECONDS diff --git a/packages/shared-python/shared/services/ai/prompt_service.py b/packages/shared-python/shared/services/ai/prompt_service.py index 9ea759675..339e00d7d 100755 --- a/packages/shared-python/shared/services/ai/prompt_service.py +++ b/packages/shared-python/shared/services/ai/prompt_service.py @@ -3,7 +3,7 @@ Only contains prompts required for document parsing workflow. Removed prompts for: -- RAG/Knowledge Base: talk-kb, merge-answers, judge-kb, rerank, connect-kb, detect-contradict +- Retrieval corpus: talk, merge-answers, judge, rerank, connect, detect-contradict - Document Generation: gen-titles-oneoff, gen-root-titles, gen-thoughts, reason-content-layout, rewrite-paras, rewrite-sentence, construct-table, reason-source - Table Filling: filling-tb-kv, filling-tb-ck diff --git a/packages/shared-python/shared/services/chunks/document_path.py b/packages/shared-python/shared/services/chunks/document_path.py new file mode 100644 index 000000000..37baa9943 --- /dev/null +++ b/packages/shared-python/shared/services/chunks/document_path.py @@ -0,0 +1,156 @@ +"""Document path parsing for parser chunk paths.""" + +from __future__ import annotations + +import os + +_DOCUMENT_FILE_EXTENSIONS = { + ".csv", + ".atlas", + ".fragment", + ".gif", + ".doc", + ".docx", + ".htm", + ".html", + ".jpeg", + ".jpg", + ".json", + ".md", + ".markdown", + ".pdf", + ".png", + ".ppt", + ".pptx", + ".rtf", + ".webp", + ".txt", + ".xls", + ".xlsm", + ".xlsx", +} +_MEDIA_ROOT_SEGMENTS = {"images", "tables"} + + +def split_document_path( + path: str | None, + *, + source_file_name: str | None = None, +) -> tuple[list[str], list[str]]: + """Return ``(root_parts, section_parts)`` for new and legacy chunk paths.""" + parts = _split_path(path, source_file_name=source_file_name) + if not parts: + return [], [] + if parts[0] in _MEDIA_ROOT_SEGMENTS and not _is_legacy_namespace_path( + parts, + source_file_name=source_file_name, + ): + return parts[:1], [] + + document_index = _find_document_index(parts, source_file_name=source_file_name) + return parts[: document_index + 1], parts[document_index + 1 :] + + +def _split_path(path: str | None, *, source_file_name: str | None) -> list[str]: + raw = str(path or "").strip() + raw_segments = raw.split("/") + source_segment = _normalize_document_file_name(source_file_name) + parts: list[str] = [] + for index, segment in enumerate(raw_segments): + parts.extend( + _split_arrow_document_segment( + segment, + can_split=_can_split_arrow_document_segment( + index=index, + raw_segments=raw_segments, + segment=segment, + source_segment=source_segment, + ), + ) + ) + return parts + + +def _can_split_arrow_document_segment( + *, + index: int, + raw_segments: list[str], + segment: str, + source_segment: str, +) -> bool: + if index == 0: + return True + if index != 1: + return False + + first_segment = raw_segments[0].strip() if raw_segments else "" + if not _is_document_file_segment(first_segment): + return True + + arrow_document_segment = _normalize_document_file_name( + segment.split("-->", 1)[0] + ) + return bool(source_segment and arrow_document_segment == source_segment) + + +def _split_arrow_document_segment(segment: str, *, can_split: bool) -> list[str]: + normalized_segment = segment.strip() + if not normalized_segment: + return [] + if not can_split or "-->" not in normalized_segment: + return [normalized_segment] + + arrow_parts = [ + part.strip() + for part in normalized_segment.split("-->") + if part.strip() + ] + if arrow_parts and _is_document_file_segment(arrow_parts[0]): + return arrow_parts + return [normalized_segment] + + +def _find_document_index( + parts: list[str], + *, + source_file_name: str | None, +) -> int: + source_segment = _normalize_document_file_name(source_file_name) + if source_segment: + for index in range(min(2, len(parts))): + if _normalize_document_file_name(parts[index]) == source_segment: + return index + + if _is_legacy_namespace_path(parts, source_file_name=source_file_name): + return 1 + if _is_document_file_segment(parts[0]): + return 0 + return 0 + + +def _is_legacy_namespace_path( + parts: list[str], + *, + source_file_name: str | None, +) -> bool: + if len(parts) < 3 or not _is_document_file_segment(parts[1]): + return False + + source_segment = _normalize_document_file_name(source_file_name) + if source_segment: + return _normalize_document_file_name(parts[1]) == source_segment + return not _is_document_file_segment(parts[0]) + + +def _normalize_document_file_name(value: str | None) -> str: + if not value: + return "" + return os.path.basename(str(value).strip().replace("\\", "/")).lower() + + +def _is_document_file_segment(segment: str) -> bool: + normalized_segment = segment.lower().strip() + return any( + normalized_segment.endswith(extension) + for extension in _DOCUMENT_FILE_EXTENSIONS + ) diff --git a/packages/shared-python/shared/services/http/__init__.py b/packages/shared-python/shared/services/http/__init__.py new file mode 100644 index 000000000..e80b56b4c --- /dev/null +++ b/packages/shared-python/shared/services/http/__init__.py @@ -0,0 +1,123 @@ +"""Public URL and outbound HTTP policy service exports.""" + +from __future__ import annotations + +from importlib import import_module +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from shared.services.http.client_pool import ( + close_async_client, + close_sync_client, + get_async_client, + get_sync_client, + ) + from shared.services.http.pinned_outbound import ( + PinnedDownloadResult, + PinnedHTTPConnection, + PinnedHTTPConnectionPool, + PinnedHTTPSConnection, + PinnedHTTPSConnectionPool, + PinnedIPResolver, + PinnedOutboundResponse, + download_pinned_outbound_file, + download_pinned_outbound_file_async, + send_pinned_outbound_request, + ) + from shared.services.http.url_file_type import ( + CONTENT_TYPE_TO_EXTENSION, + MAX_SAFE_REDIRECTS, + REDIRECT_STATUS_CODES, + URL_VALIDATION_DESCRIPTIONS, + resolve_file_extension_async, + resolve_file_extension_sync, + ) + from shared.services.http.url_security import ( + BLOCKED_HOSTNAMES, + DEVELOPMENT_ENVIRONMENTS, + HTTP_SCHEMES, + HTTPURLValidationResult, + SafePublicHTTPURL, + URLValidationFailureReason, + validate_http_url_and_resolve_ip, + validate_http_url_and_resolve_ip_async, + ) + +__all__ = [ + "BLOCKED_HOSTNAMES", + "CONTENT_TYPE_TO_EXTENSION", + "DEVELOPMENT_ENVIRONMENTS", + "HTTP_SCHEMES", + "HTTPURLValidationResult", + "MAX_SAFE_REDIRECTS", + "PinnedDownloadResult", + "PinnedHTTPConnection", + "PinnedHTTPConnectionPool", + "PinnedHTTPSConnection", + "PinnedHTTPSConnectionPool", + "PinnedIPResolver", + "PinnedOutboundResponse", + "REDIRECT_STATUS_CODES", + "SafePublicHTTPURL", + "URLValidationFailureReason", + "URL_VALIDATION_DESCRIPTIONS", + "close_async_client", + "close_sync_client", + "download_pinned_outbound_file", + "download_pinned_outbound_file_async", + "get_async_client", + "get_sync_client", + "resolve_file_extension_async", + "resolve_file_extension_sync", + "send_pinned_outbound_request", + "validate_http_url_and_resolve_ip", + "validate_http_url_and_resolve_ip_async", +] + +_EXPORT_MODULES: dict[str, str] = { + "BLOCKED_HOSTNAMES": "shared.services.http.url_security", + "CONTENT_TYPE_TO_EXTENSION": "shared.services.http.url_file_type", + "DEVELOPMENT_ENVIRONMENTS": "shared.services.http.url_security", + "HTTP_SCHEMES": "shared.services.http.url_security", + "HTTPURLValidationResult": "shared.services.http.url_security", + "MAX_SAFE_REDIRECTS": "shared.services.http.url_file_type", + "PinnedDownloadResult": "shared.services.http.pinned_outbound", + "PinnedHTTPConnection": "shared.services.http.pinned_outbound", + "PinnedHTTPConnectionPool": "shared.services.http.pinned_outbound", + "PinnedHTTPSConnection": "shared.services.http.pinned_outbound", + "PinnedHTTPSConnectionPool": "shared.services.http.pinned_outbound", + "PinnedIPResolver": "shared.services.http.pinned_outbound", + "PinnedOutboundResponse": "shared.services.http.pinned_outbound", + "REDIRECT_STATUS_CODES": "shared.services.http.url_file_type", + "SafePublicHTTPURL": "shared.services.http.url_security", + "URLValidationFailureReason": "shared.services.http.url_security", + "URL_VALIDATION_DESCRIPTIONS": "shared.services.http.url_file_type", + "close_async_client": "shared.services.http.client_pool", + "close_sync_client": "shared.services.http.client_pool", + "download_pinned_outbound_file": "shared.services.http.pinned_outbound", + "download_pinned_outbound_file_async": "shared.services.http.pinned_outbound", + "get_async_client": "shared.services.http.client_pool", + "get_sync_client": "shared.services.http.client_pool", + "resolve_file_extension_async": "shared.services.http.url_file_type", + "resolve_file_extension_sync": "shared.services.http.url_file_type", + "send_pinned_outbound_request": "shared.services.http.pinned_outbound", + "validate_http_url_and_resolve_ip": "shared.services.http.url_security", + "validate_http_url_and_resolve_ip_async": "shared.services.http.url_security", +} + + +def __getattr__(name: str) -> Any: + """Load HTTP service exports on first access.""" + module_name = _EXPORT_MODULES.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + module = import_module(module_name) + value = getattr(module, name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + """Return public HTTP package exports.""" + return sorted([*globals(), *__all__]) diff --git a/packages/shared-python/shared/utils/http_clients.py b/packages/shared-python/shared/services/http/client_pool.py similarity index 91% rename from packages/shared-python/shared/utils/http_clients.py rename to packages/shared-python/shared/services/http/client_pool.py index 928b3f02c..1ffa03893 100644 --- a/packages/shared-python/shared/utils/http_clients.py +++ b/packages/shared-python/shared/services/http/client_pool.py @@ -1,7 +1,8 @@ """ Singleton HTTP clients with connection pooling. -Reduces socket churn under high concurrency by reusing -TCP+TLS connections across requests. + +Reduces socket churn under high concurrency by reusing TCP and TLS connections +across requests. """ import threading @@ -21,7 +22,6 @@ pool=30.0, ) -# --- Sync client (gevent-patched threading.Lock) --- _sync_client: Optional[httpx.Client] = None _sync_lock = threading.Lock() @@ -49,7 +49,6 @@ def close_sync_client() -> None: _sync_client = None -# --- Async client --- _async_client: Optional[httpx.AsyncClient] = None _async_lock = threading.Lock() diff --git a/packages/shared-python/shared/utils/pinned_outbound_http.py b/packages/shared-python/shared/services/http/pinned_outbound.py similarity index 99% rename from packages/shared-python/shared/utils/pinned_outbound_http.py rename to packages/shared-python/shared/services/http/pinned_outbound.py index 808b219c8..c281dceb5 100644 --- a/packages/shared-python/shared/utils/pinned_outbound_http.py +++ b/packages/shared-python/shared/services/http/pinned_outbound.py @@ -19,7 +19,7 @@ from urllib3.util.connection import create_connection from shared.core.exceptions.domain_exceptions import ValidationException -from shared.utils.url_security import SafePublicHTTPURL +from shared.services.http.url_security import SafePublicHTTPURL @dataclass(frozen=True) diff --git a/packages/shared-python/shared/utils/url_file_type.py b/packages/shared-python/shared/services/http/url_file_type.py similarity index 86% rename from packages/shared-python/shared/utils/url_file_type.py rename to packages/shared-python/shared/services/http/url_file_type.py index 7314d7e23..14fd2a16d 100644 --- a/packages/shared-python/shared/utils/url_file_type.py +++ b/packages/shared-python/shared/services/http/url_file_type.py @@ -2,7 +2,7 @@ Resolve file extension from a URL. Tries the URL path first, then falls back to a HEAD request to read Content-Type. -Provides both async (for API) and sync (for worker) variants. +Provides both async and sync variants. """ import os @@ -12,13 +12,12 @@ from shared.core.config import settings from shared.core.exceptions.domain_exceptions import ValidationException -from shared.utils.url_security import ( +from shared.services.http.url_security import ( HTTPURLValidationResult, SafePublicHTTPURL, validate_http_url_and_resolve_ip, ) -# Content-Type to file extension mapping CONTENT_TYPE_TO_EXTENSION: dict[str, str] = { "application/pdf": ".pdf", "application/msword": ".doc", @@ -93,7 +92,6 @@ def _extension_from_content_type(content_type: str | None) -> str | None: """Map a Content-Type header value to a supported file extension.""" if not content_type: return None - # Strip parameters like "; charset=utf-8" mime = content_type.split(";")[0].strip().lower() ext = CONTENT_TYPE_TO_EXTENSION.get(mime) if ext and ext in settings.get_supported_extensions(): @@ -103,7 +101,7 @@ def _extension_from_content_type(content_type: str | None) -> str | None: async def resolve_file_extension_async(url: str) -> str | None: """ - Resolve file extension from a URL (async version for API layer). + Resolve file extension from a URL. 1. Try extracting extension from URL path. 2. If that fails, send a HEAD request and read Content-Type. @@ -116,7 +114,7 @@ async def resolve_file_extension_async(url: str) -> str | None: return ext try: - from shared.utils.http_clients import get_async_client + from shared.services.http.client_pool import get_async_client client = get_async_client() response = None @@ -140,23 +138,19 @@ async def resolve_file_extension_async(url: str) -> str | None: content_type = response.headers.get("content-type") ext = _extension_from_content_type(content_type) if ext: - logger.info( - f"Resolved file extension from Content-Type header: {ext}" - ) + logger.info(f"Resolved file extension from Content-Type header: {ext}") return ext except ValidationException: raise except Exception as exc: - logger.warning( - f"HEAD request failed for URL file type detection: {exc}" - ) + logger.warning(f"HEAD request failed for URL file type detection: {exc}") return None def resolve_file_extension_sync(url: str) -> str | None: """ - Resolve file extension from a URL (sync version for worker layer). + Resolve file extension from a URL. Same logic as async variant but uses the shared sync httpx client. """ @@ -167,7 +161,7 @@ def resolve_file_extension_sync(url: str) -> str | None: return ext try: - from shared.utils.http_clients import get_sync_client + from shared.services.http.client_pool import get_sync_client client = get_sync_client() response = None @@ -191,15 +185,11 @@ def resolve_file_extension_sync(url: str) -> str | None: content_type = response.headers.get("content-type") ext = _extension_from_content_type(content_type) if ext: - logger.info( - f"Resolved file extension from Content-Type header: {ext}" - ) + logger.info(f"Resolved file extension from Content-Type header: {ext}") return ext except ValidationException: raise except Exception as exc: - logger.warning( - f"HEAD request failed for URL file type detection: {exc}" - ) + logger.warning(f"HEAD request failed for URL file type detection: {exc}") return None diff --git a/packages/shared-python/shared/utils/url_security.py b/packages/shared-python/shared/services/http/url_security.py similarity index 100% rename from packages/shared-python/shared/utils/url_security.py rename to packages/shared-python/shared/services/http/url_security.py diff --git a/packages/shared-python/shared/services/job_lifecycle_sync.py b/packages/shared-python/shared/services/job_lifecycle_sync.py deleted file mode 100644 index cd9dcc8ef..000000000 --- a/packages/shared-python/shared/services/job_lifecycle_sync.py +++ /dev/null @@ -1,513 +0,0 @@ -""" -Sync Job Lifecycle Service for Celery worker (gevent pool). - -Encapsulates the complete job success/failure finalization that previously -used an API-side broker consumer. The worker now writes directly to the -database in a single atomic transaction, -using the same transactional outbox pattern for webhook events. -""" - -from __future__ import annotations - -import time -from datetime import datetime, timezone -from typing import Any, Dict, List, Optional -from uuid import uuid4 - -from loguru import logger -from sqlalchemy import delete, select, update -from sqlalchemy.orm import Session - -from shared.core.database_sync import get_sync_db_context -from shared.core.response import build_standard_error_response -from shared.core.state_machine.service_sync import SyncStateMachineService -from shared.models.database.job import Job -from shared.models.database.document import DocumentSection -from shared.models.database.job_result import JobChunk, JobResult -from shared.models.database.webhook import WebhookEvent, WebhookEventStatus -from shared.services.billing.credits_sync_service import SyncCreditsService -from shared.services.redis.redis_sync_service import ( - SyncRedisServiceFactory, -) -from shared.services.retrieval.publication_service import RetrievalPublicationService -from shared.utils.error_details import normalize_error_details -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder - - -def _utc_now_naive() -> datetime: - return datetime.now(timezone.utc).replace(tzinfo=None) - - -class SyncJobLifecycleService: - """Manages job lifecycle transitions in the worker process (sync/gevent). - - Implements the direct worker → DB write path for job completion and failure. - """ - - def __init__(self) -> None: - self._state_machine = SyncStateMachineService() - self._retrieval_publication = RetrievalPublicationService() - - # ── Public API ────────────────────────────────────────────────────── - - def finalize_job_success( - self, - job_id: str, - result_s3_key: str, - checksum: str, - zip_size: int, - chunks: Optional[List[Dict[str, Any]]] = None, - stored_count: int = 0, - delivery_mode: str = "url", - section_summaries: Optional[Dict[str, str]] = None, - ) -> Dict[str, Any]: - """Finalize a successful job in a single atomic transaction. - - Steps (all within one DB transaction): - 1. Upsert JobResult + replace full result chunks - 2. Publish document state from full result chunks - 3. Mark job as DONE via state machine (CAS) - 4. Create WebhookEvent if webhook_enabled - 5. COMMIT - 6. Post-commit: enqueue webhook - """ - logger.info(f"Finalizing job success: job_id={job_id}") - - with get_sync_db_context() as db: - try: - inline_payload = {"checksum": checksum} - job_result = self._upsert_job_result( - db, - job_id, - delivery_mode, - inline_payload=inline_payload, - result_s3_key=result_s3_key, - result_size=zip_size, - ) - - normalized_chunks = chunks or [] - self._replace_chunks(db, job_result.id, normalized_chunks) - previous_document_scope = ( - self._retrieval_publication.get_existing_document_scope( - db, - job_id=job_id, - ) - ) - published_document_state = ( - self._retrieval_publication.publish_document_state( - db, - job_id=job_id, - job_result_id=job_result.id, - chunks=normalized_chunks, - ) - ) - if published_document_state is not None and not published_document_state.get("skipped_all_duplicate"): - # Backfill DocumentSection.summary from enriched doc_nav data - if section_summaries: - self._backfill_section_summaries( - db, - document_id=published_document_state.get("document_id", ""), - job_result_id=job_result.id, - section_summaries=section_summaries, - ) - self._retrieval_publication.publish_document_graph( - db, - job_id=job_id, - job_result_id=job_result.id, - ) - cache_invalidation = self._build_retrieval_cache_invalidation( - db, - job_id=job_id, - published_document_state=published_document_state, - previous_document_scope=previous_document_scope, - ) - - transition_ok = self._state_machine.mark_completed( - db, - job_id, - result_metadata={ - "storage_completed": True, - "stored_count": stored_count, - "delivery_mode": delivery_mode, - }, - ) - if not transition_ok: - logger.error(f"Job {job_id} mark_completed transition failed") - db.rollback() - return { - "status": "failed", - "job_id": job_id, - "reason": "state_transition_failed", - } - - webhook_event = self._maybe_create_webhook_event( - db, - job_id, - event_type="job.completed", - ) - - db.commit() - logger.info(f"Job {job_id} success transaction committed") - - self._post_commit_invalidate_retrieval_cache(cache_invalidation) - self._post_commit_enqueue_webhook(webhook_event) - - return { - "status": "success", - "job_id": job_id, - "stored_count": stored_count, - } - - except Exception as exc: - logger.error(f"Failed to finalize job success {job_id}: {exc}") - db.rollback() - raise - - def finalize_job_failure( - self, - job_id: str, - error_message: str, - error_code: str = "UNKNOWN", - error_details: Optional[Dict[str, Any]] = None, - should_refund: bool = False, - ) -> bool: - """Finalize a failed job in a single atomic transaction. - - Steps (all within one DB transaction): - 1. Mark job as FAILED via state machine (CAS + error fields) - 2. Refund credits if needed - 3. Create WebhookEvent if webhook_enabled - 4. COMMIT - 5. Post-commit: enqueue webhook - """ - logger.info(f"Finalizing job failure: job_id={job_id}") - - with get_sync_db_context() as db: - try: - transition_ok = self._state_machine.mark_failed( - db, - job_id, - error_message, - error_code=error_code, - error_details=error_details, - ) - if not transition_ok: - logger.error(f"Job {job_id} mark_failed transition failed") - db.rollback() - return False - - if should_refund: - self._try_refund_credits(db, job_id) - - normalized_error_details = normalize_error_details(error_details) - webhook_event = self._maybe_create_webhook_event( - db, - job_id, - event_type="job.failed", - extra_payload={ - "error": build_standard_error_response( - code=error_code, - message=error_message, - request_id=job_id, - details=normalized_error_details, - ), - }, - ) - - db.commit() - logger.info(f"Job {job_id} failure transaction committed") - - self._post_commit_enqueue_webhook(webhook_event) - - return True - - except Exception as exc: - logger.error(f"Failed to finalize job failure {job_id}: {exc}") - db.rollback() - raise - - def update_progress( - self, - job_id: str, - progress: int, - message: str = "", - ) -> bool: - """Write job progress directly to Redis (replaces publish_progress_update). - - Best-effort — failures are logged but do not raise. - """ - try: - redis_service = SyncRedisServiceFactory.get_service() - task_ttl = redis_key_builder.get_key_ttl(RedisKeyType.TASK) - progress_key = redis_service._build_key( - redis_key_builder.task_progress(job_id) - ) - - pipe = redis_service.pipeline() - pipe.hset( - progress_key, - mapping={ - "progress": str(progress), - "message": message, - "timestamp": str(int(time.time())), - }, - ) - pipe.expire(progress_key, task_ttl) - pipe.execute() - return True - except Exception as exc: - logger.warning(f"Failed to update progress for job {job_id}: {exc}") - return False - - # ── Private helpers ───────────────────────────────────────────────── - - def _backfill_section_summaries( - self, - db: Session, - *, - document_id: str, - job_result_id: str, - section_summaries: Dict[str, str], - ) -> None: - """Populate DocumentSection.summary from enriched doc_nav data. - - Runs UPDATE statements within the existing transaction so no extra - commit is needed. Overwrites any existing summary value since the - enriched doc_nav data is the authoritative source. - """ - if not document_id or not section_summaries: - return - try: - for path, summary in section_summaries.items(): - if not path or not summary: - continue - db.execute( - update(DocumentSection) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - .where(DocumentSection.section_path == path) - .values(summary=summary) - ) - db.flush() - logger.debug( - f"Backfilled section summaries: document_id={document_id}, " - f"count={len(section_summaries)}" - ) - except Exception as exc: - logger.warning(f"Section summary backfill failed (non-fatal): {exc}") - - def _upsert_job_result( - self, - db: Session, - job_id: str, - delivery_mode: str, - *, - inline_payload: Optional[Dict[str, Any]] = None, - result_s3_key: Optional[str] = None, - result_size: Optional[int] = None, - ) -> JobResult: - """Create or update JobResult row.""" - result = db.execute(select(JobResult).where(JobResult.job_id == job_id)) - existing = result.scalar_one_or_none() - - if existing: - existing.delivery_mode = delivery_mode - existing.inline_payload = inline_payload - existing.result_s3_key = result_s3_key - existing.result_size = result_size - db.flush() - return existing - - job_result = JobResult( - job_id=job_id, - delivery_mode=delivery_mode, - document_metadata={}, - inline_payload=inline_payload, - result_s3_key=result_s3_key, - result_size=result_size, - ) - db.add(job_result) - db.flush() - return job_result - - def _replace_chunks( - self, - db: Session, - job_result_id: str, - chunks: List[Dict[str, Any]], - ) -> Optional[Dict[str, Any]]: - """Delete existing chunks and insert new ones.""" - db.execute(delete(JobChunk).where(JobChunk.job_result_id == job_result_id)) - - if not chunks: - db.flush() - return - - chunk_models = [] - for index, chunk in enumerate(chunks): - chunk_identifier = chunk.get("chunk_id") or str(uuid4()) - chunk_models.append( - JobChunk( - job_result_id=job_result_id, - chunk_id=chunk_identifier, - chunk_type=chunk.get("type", "paragraph"), - text=chunk.get("text"), - path=chunk.get("metadata", {}).get("path"), - chunk_metadata=chunk.get("metadata"), - sort_order=chunk.get("order", index), - ) - ) - db.add_all(chunk_models) - db.flush() - - def _build_retrieval_cache_invalidation( - self, - db: Session, - *, - job_id: str, - published_document_state: Optional[Dict[str, str]], - previous_document_scope: Optional[Dict[str, str]], - ) -> Optional[Dict[str, Any]]: - job = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() - if not job: - return None - - namespaces: list[str] = [] - metadata = job.job_metadata or {} - new_namespace = metadata.get("namespace") or "default" - namespaces.append(new_namespace) - - if previous_document_scope and previous_document_scope.get("namespace"): - namespaces.append(previous_document_scope["namespace"]) - if published_document_state and published_document_state.get("namespace"): - namespaces.append(published_document_state["namespace"]) - - return {"user_id": str(job.user_id), "namespaces": namespaces, "job_id": job_id} - - def _post_commit_invalidate_retrieval_cache( - self, cache_invalidation: Optional[Dict[str, Any]] - ) -> None: - if not cache_invalidation: - return - try: - redis_service = SyncRedisServiceFactory.get_service() - user_id = cache_invalidation["user_id"] - seen: set[str] = set() - for namespace in cache_invalidation["namespaces"]: - if not namespace or namespace in seen: - continue - seen.add(namespace) - redis_service.incr(f"retrieval:version:{user_id}:{namespace}") - except Exception as exc: - logger.warning( - f"Failed to invalidate retrieval cache after publication (ignored): job_id={cache_invalidation.get('job_id')}, error={exc}" - ) - - def _maybe_create_webhook_event( - self, - db: Session, - job_id: str, - event_type: str, - extra_payload: Optional[Dict[str, Any]] = None, - ) -> Optional[WebhookEvent]: - """Create a WebhookEvent if the job has webhooks enabled.""" - result = db.execute(select(Job).where(Job.job_id == job_id)) - job = result.scalar_one_or_none() - - if not job: - logger.warning(f"Job not found for webhook check: {job_id}") - return None - - webhook_url = getattr(job, "webhook_url", None) - if not job.webhook_enabled or not webhook_url: - return None - - status = "completed" if event_type == "job.completed" else "failed" - timestamp_key = f"{status}_at" - payload: Dict[str, Any] = { - "event": event_type, - "job_id": job_id, - "status": status, - timestamp_key: _utc_now_naive().isoformat(), - } - if extra_payload: - payload.update(extra_payload) - - event = WebhookEvent( - job_id=job_id, - target_url=webhook_url, - payload=payload, - status=WebhookEventStatus.PENDING, - attempts=0, - ) - db.add(event) - db.flush() - logger.info(f"WebhookEvent created: event_id={event.id}, job_id={job_id}") - return event - - def _try_refund_credits(self, db: Session, job_id: str) -> None: - """Attempt to refund credits for a failed job.""" - try: - result = db.execute(select(Job).where(Job.job_id == job_id)) - job = result.scalar_one_or_none() - if not job: - return - - amount = getattr(job, "credits_charged", 0) or 0 - billing_status = getattr(job, "billing_status", "") - if amount <= 0 or billing_status != "charged": - return - - credits_service = SyncCreditsService() - credits_service.refund_job_credits( - session=db, - user_id=str(job.user_id), - amount=amount, - job_id=job_id, - ) - job.billing_status = "refunded" - logger.info(f"Refunded {amount} credits for job {job_id}") - except Exception as exc: - logger.error(f"Credit refund failed for job {job_id}: {exc}") - - def _post_commit_enqueue_webhook( - self, - webhook_event: Optional[WebhookEvent], - ) -> None: - """Publish a persisted webhook via QStash after commit (best-effort).""" - if not webhook_event: - return - - try: - from shared.services.webhook.qstash_publisher import ( - get_qstash_webhook_publisher, - ) - - publisher = get_qstash_webhook_publisher() - message_id = publisher.publish_event(webhook_event.id) - if not message_id: - logger.warning( - f"Webhook publish failed after commit: event_id={webhook_event.id}" - ) - return - logger.info( - f"Webhook published after commit: event_id={webhook_event.id}, " - f"message_id={message_id}" - ) - except Exception as exc: - logger.error( - f"Failed to publish webhook after commit (event persisted): " - f"event_id={webhook_event.id}, error={exc}" - ) - - -# Module-level singleton -_lifecycle_service: Optional[SyncJobLifecycleService] = None - - -def get_sync_job_lifecycle_service() -> SyncJobLifecycleService: - """Get the singleton sync job lifecycle service.""" - global _lifecycle_service - if _lifecycle_service is None: - _lifecycle_service = SyncJobLifecycleService() - return _lifecycle_service diff --git a/packages/shared-python/shared/services/jobs/lifecycle/__init__.py b/packages/shared-python/shared/services/jobs/lifecycle/__init__.py new file mode 100644 index 000000000..93da44f12 --- /dev/null +++ b/packages/shared-python/shared/services/jobs/lifecycle/__init__.py @@ -0,0 +1 @@ +"""Synchronous job lifecycle finalization workflow.""" diff --git a/packages/shared-python/shared/services/jobs/lifecycle/failure_finalizer.py b/packages/shared-python/shared/services/jobs/lifecycle/failure_finalizer.py new file mode 100644 index 000000000..b87b270fb --- /dev/null +++ b/packages/shared-python/shared/services/jobs/lifecycle/failure_finalizer.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.orm import Session + +from shared.core.response import build_standard_error_response +from shared.core.state_machine.service_sync import SyncStateMachineService +from shared.models.database.job import Job +from shared.services.billing.credits_sync_service import SyncCreditsService +from shared.services.jobs.lifecycle.post_commit_effects import PostCommitEffectPlan +from shared.services.jobs.lifecycle.webhook_outbox import SyncJobWebhookOutbox +from shared.utils.error_details import normalize_error_details + + +@dataclass(frozen=True) +class JobFailureFinalization: + succeeded: bool + post_commit_effects: PostCommitEffectPlan + + +class SyncJobFailureFinalizer: + """Finalize failed Jobs inside the lifecycle transaction.""" + + def __init__( + self, + *, + state_machine: SyncStateMachineService | None = None, + webhook_outbox: SyncJobWebhookOutbox | None = None, + credits_service: SyncCreditsService | None = None, + ) -> None: + self._state_machine = state_machine or SyncStateMachineService() + self._webhook_outbox = webhook_outbox or SyncJobWebhookOutbox() + self._credits_service = credits_service or SyncCreditsService() + + def finalize( + self, + db: Session, + *, + job_id: str, + error_message: str, + error_code: str, + error_details: dict[str, Any] | None, + should_refund: bool, + ) -> JobFailureFinalization: + transition_outcome = self._state_machine.mark_failed_outcome( + db, + job_id, + error_message, + error_code=error_code, + error_details=error_details, + ) + if not transition_outcome.succeeded: + logger.error( + f"Job {job_id} mark_failed transition failed: " + f"reason={transition_outcome.reason}" + ) + return JobFailureFinalization( + succeeded=False, + post_commit_effects=PostCommitEffectPlan.none(), + ) + + if should_refund: + self._try_refund_credits(db, job_id) + + normalized_error_details = normalize_error_details(error_details) + webhook_event = self._webhook_outbox.create_event( + db, + job_id=job_id, + event_type="job.failed", + extra_payload={ + "error": build_standard_error_response( + code=error_code, + message=error_message, + request_id=job_id, + details=normalized_error_details, + ), + }, + ) + return JobFailureFinalization( + succeeded=True, + post_commit_effects=PostCommitEffectPlan.from_failure( + webhook_event_id=webhook_event.event_id if webhook_event else None, + ), + ) + + def _try_refund_credits(self, db: Session, job_id: str) -> None: + try: + result = db.execute(select(Job).where(Job.job_id == job_id)) + job = result.scalar_one_or_none() + if not job: + return + + amount = getattr(job, "credits_charged", 0) or 0 + billing_status = getattr(job, "billing_status", "") + if amount <= 0 or billing_status != "charged": + return + + self._credits_service.refund_job_credits( + session=db, + user_id=str(job.user_id), + amount=amount, + job_id=job_id, + ) + job.billing_status = "refunded" + logger.info(f"Refunded {amount} credits for job {job_id}") + except Exception as exc: + logger.error(f"Credit refund failed for job {job_id}: {exc}") diff --git a/packages/shared-python/shared/services/jobs/lifecycle/post_commit_effects.py b/packages/shared-python/shared/services/jobs/lifecycle/post_commit_effects.py new file mode 100644 index 000000000..e990ec5a4 --- /dev/null +++ b/packages/shared-python/shared/services/jobs/lifecycle/post_commit_effects.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + +from shared.services.jobs.lifecycle.publication import ( + RetrievalCacheInvalidation, + SyncJobPublicationFinalizer, +) +from shared.services.jobs.lifecycle.webhook_outbox import SyncJobWebhookOutbox + + +@dataclass(frozen=True) +class PostCommitEffectPlan: + retrieval_cache_invalidations: tuple[RetrievalCacheInvalidation, ...] = field( + default_factory=tuple + ) + webhook_event_ids: tuple[str, ...] = field(default_factory=tuple) + + @classmethod + def none(cls) -> PostCommitEffectPlan: + return cls() + + @classmethod + def from_success( + cls, + *, + cache_invalidation: RetrievalCacheInvalidation | None, + webhook_event_id: str | None, + ) -> PostCommitEffectPlan: + return cls( + retrieval_cache_invalidations=( + (cache_invalidation,) if cache_invalidation else () + ), + webhook_event_ids=((webhook_event_id,) if webhook_event_id else ()), + ) + + @classmethod + def from_failure( + cls, + *, + webhook_event_id: str | None, + ) -> PostCommitEffectPlan: + return cls(webhook_event_ids=((webhook_event_id,) if webhook_event_id else ())) + + +class SyncJobPostCommitEffectRunner: + def __init__( + self, + *, + publication_finalizer: SyncJobPublicationFinalizer | None = None, + webhook_outbox: SyncJobWebhookOutbox | None = None, + ) -> None: + self._publication_finalizer = ( + publication_finalizer or SyncJobPublicationFinalizer() + ) + self._webhook_outbox = webhook_outbox or SyncJobWebhookOutbox() + + def run(self, plan: PostCommitEffectPlan) -> None: + for cache_invalidation in plan.retrieval_cache_invalidations: + self._publication_finalizer.invalidate_cache_after_commit( + cache_invalidation + ) + for webhook_event_id in plan.webhook_event_ids: + self._webhook_outbox.enqueue_event_id_after_commit(webhook_event_id) diff --git a/packages/shared-python/shared/services/jobs/lifecycle/publication.py b/packages/shared-python/shared/services/jobs/lifecycle/publication.py new file mode 100644 index 000000000..ee6c78e38 --- /dev/null +++ b/packages/shared-python/shared/services/jobs/lifecycle/publication.py @@ -0,0 +1,179 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from loguru import logger +from sqlalchemy import select, update +from sqlalchemy.orm import Session + +from shared.models.database.document import DocumentSection +from shared.models.database.job import Job +from shared.models.schemas.job_metadata import JobMetadataHelper +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace +from shared.services.redis.redis_sync_service import SyncRedisServiceFactory +from shared.services.retrieval.publication_service import RetrievalPublicationService +from shared.services.retrieval.publication_models import ( + ExistingDocumentScope, + PublishedDocumentState, +) + + +@dataclass(frozen=True) +class RetrievalCacheInvalidation: + user_id: str + namespaces: tuple[str, ...] + job_id: str + + +@dataclass(frozen=True) +class JobPublicationOutcome: + published_document_state: PublishedDocumentState | None + cache_invalidation: RetrievalCacheInvalidation | None + + +class SyncJobPublicationFinalizer: + """Publish terminal parse results and invalidate retrieval cache after commit.""" + + def __init__( + self, + *, + retrieval_publication: RetrievalPublicationService | None = None, + ) -> None: + self._retrieval_publication = ( + retrieval_publication or RetrievalPublicationService() + ) + + def publish_result( + self, + db: Session, + *, + job_id: str, + job_result_id: str, + chunks: list[dict[str, Any]], + section_summaries: dict[str, str] | None, + ) -> JobPublicationOutcome: + previous_document_scope = self._retrieval_publication.get_existing_document_scope( + db, + job_id=job_id, + ) + published_document_state = self._retrieval_publication.publish_document_state( + db, + job_id=job_id, + job_result_id=job_result_id, + chunks=chunks, + ) + if _should_publish_document_graph(published_document_state): + assert published_document_state is not None + if section_summaries: + self._backfill_section_summaries( + db, + document_id=published_document_state.document_id or "", + job_result_id=job_result_id, + section_summaries=section_summaries, + ) + self._retrieval_publication.publish_document_graph( + db, + job_id=job_id, + job_result_id=job_result_id, + ) + + cache_invalidation = self._build_cache_invalidation( + db, + job_id=job_id, + published_document_state=published_document_state, + previous_document_scope=previous_document_scope, + ) + return JobPublicationOutcome( + published_document_state=published_document_state, + cache_invalidation=cache_invalidation, + ) + + def invalidate_cache_after_commit( + self, + cache_invalidation: RetrievalCacheInvalidation | None, + ) -> None: + if not cache_invalidation: + return + + try: + redis_service = SyncRedisServiceFactory.get_service() + user_id = cache_invalidation.user_id + seen: set[str] = set() + for raw_namespace in cache_invalidation.namespaces: + namespace = normalize_retrieval_namespace(str(raw_namespace)) + if not namespace or namespace in seen: + continue + seen.add(namespace) + redis_service.incr(f"retrieval:version:{user_id}:{namespace}") + except Exception as exc: + logger.warning( + "Failed to invalidate retrieval cache after publication " + f"(ignored): job_id={cache_invalidation.job_id}, error={exc}" + ) + + def _backfill_section_summaries( + self, + db: Session, + *, + document_id: str, + job_result_id: str, + section_summaries: dict[str, str], + ) -> None: + if not document_id or not section_summaries: + return + + try: + for path, summary in section_summaries.items(): + if not path or not summary: + continue + db.execute( + update(DocumentSection) + .where(DocumentSection.document_id == document_id) + .where(DocumentSection.job_result_id == job_result_id) + .where(DocumentSection.section_path == path) + .values(summary=summary) + ) + db.flush() + logger.debug( + f"Backfilled section summaries: document_id={document_id}, " + f"count={len(section_summaries)}" + ) + except Exception as exc: + logger.warning(f"Section summary backfill failed (non-fatal): {exc}") + + def _build_cache_invalidation( + self, + db: Session, + *, + job_id: str, + published_document_state: PublishedDocumentState | None, + previous_document_scope: ExistingDocumentScope | None, + ) -> RetrievalCacheInvalidation | None: + job = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() + if not job: + return None + + metadata = job.job_metadata or {} + namespaces: list[str] = [ + JobMetadataHelper.get_namespace(metadata, "default") or "default", + ] + if previous_document_scope: + namespaces.append(previous_document_scope.namespace) + if published_document_state: + namespaces.append(published_document_state.namespace) + + return RetrievalCacheInvalidation( + user_id=str(job.user_id), + namespaces=tuple(namespaces), + job_id=job_id, + ) + + +def _should_publish_document_graph( + published_document_state: PublishedDocumentState | None, +) -> bool: + return ( + published_document_state is not None + and not published_document_state.skipped_all_duplicate + ) diff --git a/packages/shared-python/shared/services/jobs/lifecycle/result_writer.py b/packages/shared-python/shared/services/jobs/lifecycle/result_writer.py new file mode 100644 index 000000000..22983a1ec --- /dev/null +++ b/packages/shared-python/shared/services/jobs/lifecycle/result_writer.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from typing import Any +from uuid import uuid4 + +from sqlalchemy import delete, select +from sqlalchemy.orm import Session + +from shared.models.database.job_result import JobChunk, JobResult + + +class SyncJobResultWriter: + """Persist terminal Job Result artifacts inside an existing transaction.""" + + def upsert_job_result( + self, + db: Session, + job_id: str, + delivery_mode: str, + *, + inline_payload: dict[str, Any] | None = None, + result_s3_key: str | None = None, + result_size: int | None = None, + ) -> JobResult: + result = db.execute(select(JobResult).where(JobResult.job_id == job_id)) + existing = result.scalar_one_or_none() + + if existing: + existing.delivery_mode = delivery_mode + existing.inline_payload = inline_payload + existing.result_s3_key = result_s3_key + existing.result_size = result_size + db.flush() + return existing + + job_result = JobResult( + job_id=job_id, + delivery_mode=delivery_mode, + document_metadata={}, + inline_payload=inline_payload, + result_s3_key=result_s3_key, + result_size=result_size, + ) + db.add(job_result) + db.flush() + return job_result + + def replace_chunks( + self, + db: Session, + job_result_id: str, + chunks: list[dict[str, Any]], + ) -> None: + db.execute(delete(JobChunk).where(JobChunk.job_result_id == job_result_id)) + + if not chunks: + db.flush() + return + + chunk_models = [] + for index, chunk in enumerate(chunks): + chunk_identifier = chunk.get("chunk_id") or str(uuid4()) + metadata = chunk.get("metadata") + chunk_text = chunk.get("text") or chunk.get("content") + chunk_path = ( + metadata.get("path") + if isinstance(metadata, dict) and metadata.get("path") + else chunk.get("path") + ) + chunk_models.append( + JobChunk( + job_result_id=job_result_id, + chunk_id=chunk_identifier, + chunk_type=chunk.get("type", "paragraph"), + text=str(chunk_text) if chunk_text is not None else None, + path=str(chunk_path) if chunk_path is not None else None, + chunk_metadata=metadata, + sort_order=chunk.get("order", index), + ) + ) + db.add_all(chunk_models) + db.flush() diff --git a/packages/shared-python/shared/services/jobs/lifecycle/service.py b/packages/shared-python/shared/services/jobs/lifecycle/service.py new file mode 100644 index 000000000..3a611d1f8 --- /dev/null +++ b/packages/shared-python/shared/services/jobs/lifecycle/service.py @@ -0,0 +1,199 @@ +""" +Sync Job Lifecycle Service for Celery worker (gevent pool). + +Encapsulates the complete job success/failure finalization that previously +used an API-side broker consumer. The worker now writes directly to the +database in a single atomic transaction, +using the same transactional outbox pattern for webhook events. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from typing import Any, Dict, List, Optional, TypeVar + +from loguru import logger +from sqlalchemy.orm import Session + +from shared.core.database_sync import get_sync_db_context +from shared.services.jobs.lifecycle.failure_finalizer import SyncJobFailureFinalizer +from shared.services.jobs.lifecycle.post_commit_effects import ( + PostCommitEffectPlan, + SyncJobPostCommitEffectRunner, +) +from shared.services.jobs.lifecycle.success_finalizer import SyncJobSuccessFinalizer +from shared.services.redis.redis_sync_service import ( + SyncRedisServiceFactory, +) +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder + + +class SyncJobLifecycleService: + """Manages job lifecycle transitions in the worker process (sync/gevent). + + Implements the direct worker → DB write path for job completion and failure. + """ + + def __init__(self) -> None: + self._success_finalizer = SyncJobSuccessFinalizer() + self._failure_finalizer = SyncJobFailureFinalizer() + self._post_commit_effect_runner = SyncJobPostCommitEffectRunner() + + # ── Public API ────────────────────────────────────────────────────── + + def finalize_job_success( + self, + job_id: str, + result_s3_key: str, + checksum: str, + zip_size: int, + chunks: Optional[List[Dict[str, Any]]] = None, + stored_count: int = 0, + delivery_mode: str = "url", + section_summaries: Optional[Dict[str, str]] = None, + ) -> Dict[str, Any]: + """Finalize a successful job in a single atomic transaction. + + Steps (all within one DB transaction): + 1. Upsert JobResult + replace full result chunks + 2. Publish document state from full result chunks + 3. Mark job as DONE via state machine (CAS) + 4. Create WebhookEvent if webhook_enabled + 5. COMMIT + 6. Post-commit: enqueue webhook + """ + logger.info(f"Finalizing job success: job_id={job_id}") + + return _run_lifecycle_transaction( + job_id=job_id, + label="success", + finalize=lambda db: self._success_finalizer.finalize( + db, + job_id=job_id, + result_s3_key=result_s3_key, + checksum=checksum, + zip_size=zip_size, + chunks=chunks or [], + stored_count=stored_count, + delivery_mode=delivery_mode, + section_summaries=section_summaries, + ), + should_commit=lambda finalization: finalization.response.should_commit(), + build_response=lambda finalization: finalization.response.to_dict(), + build_effect_plan=lambda finalization: finalization.post_commit_effects, + run_after_commit_effects=self._post_commit_effect_runner.run, + ) + + def finalize_job_failure( + self, + job_id: str, + error_message: str, + error_code: str = "UNKNOWN", + error_details: Optional[Dict[str, Any]] = None, + should_refund: bool = False, + ) -> bool: + """Finalize a failed job in a single atomic transaction. + + Steps (all within one DB transaction): + 1. Mark job as FAILED via state machine (CAS + error fields) + 2. Refund credits if needed + 3. Create WebhookEvent if webhook_enabled + 4. COMMIT + 5. Post-commit: enqueue webhook + """ + logger.info(f"Finalizing job failure: job_id={job_id}") + + return _run_lifecycle_transaction( + job_id=job_id, + label="failure", + finalize=lambda db: self._failure_finalizer.finalize( + db, + job_id=job_id, + error_message=error_message, + error_code=error_code, + error_details=error_details, + should_refund=should_refund, + ), + should_commit=lambda finalization: finalization.succeeded, + build_response=lambda finalization: finalization.succeeded, + build_effect_plan=lambda finalization: finalization.post_commit_effects, + run_after_commit_effects=self._post_commit_effect_runner.run, + ) + + def update_progress( + self, + job_id: str, + progress: int, + message: str = "", + ) -> bool: + """Write job progress directly to Redis (replaces publish_progress_update). + + Best-effort — failures are logged but do not raise. + """ + try: + redis_service = SyncRedisServiceFactory.get_service() + task_ttl = redis_key_builder.get_key_ttl(RedisKeyType.TASK) + progress_key = redis_service._build_key( + redis_key_builder.task_progress(job_id) + ) + + pipe = redis_service.pipeline() + pipe.hset( + progress_key, + mapping={ + "progress": str(progress), + "message": message, + "timestamp": str(int(time.time())), + }, + ) + pipe.expire(progress_key, task_ttl) + pipe.execute() + return True + except Exception as exc: + logger.warning(f"Failed to update progress for job {job_id}: {exc}") + return False + +# Module-level singleton +_lifecycle_service: Optional[SyncJobLifecycleService] = None + + +def get_sync_job_lifecycle_service() -> SyncJobLifecycleService: + """Get the singleton sync job lifecycle service.""" + global _lifecycle_service + if _lifecycle_service is None: + _lifecycle_service = SyncJobLifecycleService() + return _lifecycle_service + + +_FinalizationT = TypeVar("_FinalizationT") +_ResponseT = TypeVar("_ResponseT") + + +def _run_lifecycle_transaction( + *, + job_id: str, + label: str, + finalize: Callable[[Session], _FinalizationT], + should_commit: Callable[[_FinalizationT], bool], + build_response: Callable[[_FinalizationT], _ResponseT], + build_effect_plan: Callable[[_FinalizationT], PostCommitEffectPlan], + run_after_commit_effects: Callable[[PostCommitEffectPlan], None], +) -> _ResponseT: + with get_sync_db_context() as db: + try: + finalization = finalize(db) + if not should_commit(finalization): + db.rollback() + return build_response(finalization) + + db.commit() + logger.info(f"Job {job_id} {label} transaction committed") + + run_after_commit_effects(build_effect_plan(finalization)) + return build_response(finalization) + + except Exception as exc: + logger.error(f"Failed to finalize job {label} {job_id}: {exc}") + db.rollback() + raise diff --git a/packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py b/packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py new file mode 100644 index 000000000..d05b41436 --- /dev/null +++ b/packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Literal + +from loguru import logger +from sqlalchemy.orm import Session + +from shared.core.state_machine.service_sync import SyncStateMachineService +from shared.services.jobs.lifecycle.post_commit_effects import PostCommitEffectPlan +from shared.services.jobs.lifecycle.publication import SyncJobPublicationFinalizer +from shared.services.jobs.lifecycle.result_writer import SyncJobResultWriter +from shared.services.jobs.lifecycle.webhook_outbox import SyncJobWebhookOutbox + + +@dataclass(frozen=True) +class JobSuccessResponse: + status: Literal["success", "failed"] + job_id: str + stored_count: int | None = None + reason: str | None = None + + @classmethod + def completed(cls, *, job_id: str, stored_count: int) -> JobSuccessResponse: + return cls(status="success", job_id=job_id, stored_count=stored_count) + + @classmethod + def state_transition_failed(cls, *, job_id: str) -> JobSuccessResponse: + return cls(status="failed", job_id=job_id, reason="state_transition_failed") + + def should_commit(self) -> bool: + return self.status == "success" + + def to_dict(self) -> dict[str, Any]: + response: dict[str, Any] = { + "status": self.status, + "job_id": self.job_id, + } + if self.stored_count is not None: + response["stored_count"] = self.stored_count + if self.reason: + response["reason"] = self.reason + return response + + +@dataclass(frozen=True) +class JobSuccessFinalization: + response: JobSuccessResponse + post_commit_effects: PostCommitEffectPlan + + +class SyncJobSuccessFinalizer: + """Finalize successful Jobs inside the lifecycle transaction.""" + + def __init__( + self, + *, + state_machine: SyncStateMachineService | None = None, + result_writer: SyncJobResultWriter | None = None, + publication_finalizer: SyncJobPublicationFinalizer | None = None, + webhook_outbox: SyncJobWebhookOutbox | None = None, + ) -> None: + self._state_machine = state_machine or SyncStateMachineService() + self._result_writer = result_writer or SyncJobResultWriter() + self._publication_finalizer = ( + publication_finalizer or SyncJobPublicationFinalizer() + ) + self._webhook_outbox = webhook_outbox or SyncJobWebhookOutbox() + + def finalize( + self, + db: Session, + *, + job_id: str, + result_s3_key: str, + checksum: str, + zip_size: int, + chunks: list[dict[str, Any]], + stored_count: int, + delivery_mode: str, + section_summaries: dict[str, str] | None, + ) -> JobSuccessFinalization: + job_result = self._result_writer.upsert_job_result( + db, + job_id, + delivery_mode, + inline_payload={"checksum": checksum}, + result_s3_key=result_s3_key, + result_size=zip_size, + ) + self._result_writer.replace_chunks(db, job_result.id, chunks) + publication_outcome = self._publication_finalizer.publish_result( + db, + job_id=job_id, + job_result_id=job_result.id, + chunks=chunks, + section_summaries=section_summaries, + ) + + transition_outcome = self._state_machine.mark_completed_outcome( + db, + job_id, + result_metadata={ + "storage_completed": True, + "stored_count": stored_count, + "delivery_mode": delivery_mode, + }, + ) + if not transition_outcome.succeeded: + logger.error( + f"Job {job_id} mark_completed transition failed: " + f"reason={transition_outcome.reason}" + ) + return JobSuccessFinalization( + response=JobSuccessResponse.state_transition_failed(job_id=job_id), + post_commit_effects=PostCommitEffectPlan.none(), + ) + + webhook_event = self._webhook_outbox.create_event( + db, + job_id=job_id, + event_type="job.completed", + ) + return JobSuccessFinalization( + response=JobSuccessResponse.completed( + job_id=job_id, + stored_count=stored_count, + ), + post_commit_effects=PostCommitEffectPlan.from_success( + cache_invalidation=publication_outcome.cache_invalidation, + webhook_event_id=webhook_event.event_id if webhook_event else None, + ), + ) diff --git a/packages/shared-python/shared/services/jobs/lifecycle/webhook_outbox.py b/packages/shared-python/shared/services/jobs/lifecycle/webhook_outbox.py new file mode 100644 index 000000000..5717d4fa4 --- /dev/null +++ b/packages/shared-python/shared/services/jobs/lifecycle/webhook_outbox.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Any + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.orm import Session + +from shared.models.database.job import Job +from shared.models.database.webhook import WebhookEvent, WebhookEventStatus + + +@dataclass(frozen=True) +class WebhookOutboxEvent: + event_id: str + + +class SyncJobWebhookOutbox: + """Create webhook events in-transaction and publish them after commit.""" + + def create_event( + self, + db: Session, + *, + job_id: str, + event_type: str, + extra_payload: dict[str, Any] | None = None, + ) -> WebhookOutboxEvent | None: + result = db.execute(select(Job).where(Job.job_id == job_id)) + job = result.scalar_one_or_none() + + if not job: + logger.warning(f"Job not found for webhook check: {job_id}") + return None + + webhook_url = getattr(job, "webhook_url", None) + if not job.webhook_enabled or not webhook_url: + return None + + status = "completed" if event_type == "job.completed" else "failed" + timestamp_key = f"{status}_at" + payload: dict[str, Any] = { + "event": event_type, + "job_id": job_id, + "status": status, + timestamp_key: _utc_now_naive().isoformat(), + } + if extra_payload: + payload.update(extra_payload) + + event = WebhookEvent( + job_id=job_id, + target_url=webhook_url, + payload=payload, + status=WebhookEventStatus.PENDING, + attempts=0, + ) + db.add(event) + db.flush() + logger.info(f"WebhookEvent created: event_id={event.id}, job_id={job_id}") + return WebhookOutboxEvent(event_id=event.id) + + def enqueue_after_commit(self, webhook_event: WebhookOutboxEvent | None) -> None: + if not webhook_event: + return + self.enqueue_event_id_after_commit(webhook_event.event_id) + + def enqueue_event_id_after_commit(self, webhook_event_id: str | None) -> None: + if not webhook_event_id: + return + try: + from shared.services.webhook.qstash_publisher import ( + get_qstash_webhook_publisher, + ) + + publisher = get_qstash_webhook_publisher() + message_id = publisher.publish_event(webhook_event_id) + if not message_id: + logger.warning( + f"Webhook publish failed after commit: event_id={webhook_event_id}" + ) + return + logger.info( + f"Webhook published after commit: event_id={webhook_event_id}, " + f"message_id={message_id}" + ) + except Exception as exc: + logger.error( + "Failed to publish webhook after commit (event persisted): " + f"event_id={webhook_event_id}, error={exc}" + ) + + +def _utc_now_naive() -> datetime: + return datetime.now(timezone.utc).replace(tzinfo=None) diff --git a/packages/shared-python/shared/services/jobs/result_delivery.py b/packages/shared-python/shared/services/jobs/result_delivery.py new file mode 100644 index 000000000..0a5c9b0d6 --- /dev/null +++ b/packages/shared-python/shared/services/jobs/result_delivery.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Any + +from shared.services.storage.job_file_storage import JobFileStorage +from shared.utils.utc_now import utc_now_naive + + +@dataclass(frozen=True) +class JobResultDelivery: + result: dict[str, Any] | None + result_url: str | None + result_url_expires_at: datetime | None + + +class JobResultDeliveryResolver: + """Resolve the public delivery fields exposed for a terminal Job Result.""" + + def __init__(self, *, storage: JobFileStorage | None = None) -> None: + self._storage = storage or JobFileStorage() + + def resolve( + self, + job_result: Any | None, + *, + default_expires_at: datetime | None = None, + ) -> JobResultDelivery: + result = None + result_url = None + result_url_expires_at = default_expires_at + + if not job_result: + return JobResultDelivery( + result=result, + result_url=result_url, + result_url_expires_at=result_url_expires_at, + ) + + inline_payload = getattr(job_result, "inline_payload", None) + if inline_payload: + result = inline_payload + + result_s3_key = getattr(job_result, "result_s3_key", None) + if result_s3_key: + url_info = self._storage.generate_download_url( + result_s3_key, + bucket=self._storage.results_bucket, + ) + result_url = url_info["download_url"] + expires_in = int(url_info.get("expires_in", 3600)) + result_url_expires_at = utc_now_naive() + timedelta(seconds=expires_in) + + return JobResultDelivery( + result=result, + result_url=result_url, + result_url_expires_at=result_url_expires_at, + ) + + def enrich_payload( + self, + payload: dict[str, Any], + *, + job_result: Any | None, + ) -> dict[str, Any]: + if payload.get("event") != "job.completed": + return payload + + delivery = self.resolve(job_result) + enriched = dict(payload) + if delivery.result_url: + enriched["result_url"] = delivery.result_url + if delivery.result: + enriched["result"] = delivery.result + return enriched diff --git a/packages/shared-python/shared/services/quota/__init__.py b/packages/shared-python/shared/services/quota/__init__.py new file mode 100644 index 000000000..749a6d0ec --- /dev/null +++ b/packages/shared-python/shared/services/quota/__init__.py @@ -0,0 +1,5 @@ +"""Shared quota and token-pool services.""" + +from shared.services.quota.token_pool import BaseQuotaManager, TokenConfig, TokenLease + +__all__ = ["BaseQuotaManager", "TokenConfig", "TokenLease"] diff --git a/packages/shared-python/shared/utils/quota_manager.py b/packages/shared-python/shared/services/quota/token_pool.py similarity index 100% rename from packages/shared-python/shared/utils/quota_manager.py rename to packages/shared-python/shared/services/quota/token_pool.py diff --git a/packages/shared-python/shared/services/redis/__init__.py b/packages/shared-python/shared/services/redis/__init__.py index b4186f21f..a0cb3f46e 100644 --- a/packages/shared-python/shared/services/redis/__init__.py +++ b/packages/shared-python/shared/services/redis/__init__.py @@ -1,13 +1,21 @@ """Redis service exports.""" -from .job_info_redis_service import JobInfoRedisService -from .job_metadata_service import JobMetadataService -from .redis_alerts import AlertRule, RedisAlertManager, RedisAlertNotifier -from .redis_monitor import RedisMonitor -from .redis_service import RedisService -from .redis_service_factory import RedisServiceFactory -from .task_redis_service import TaskRedisService -from .user_redis_service import UserRedisService +from __future__ import annotations + +from importlib import import_module +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .job_info_redis_service import JobInfoRedisService + from .job_metadata_service import JobMetadataService + from .key_builder import RedisKeyBuilder, RedisKeyType, redis_key_builder + from .redis_alerts import AlertRule, RedisAlertManager, RedisAlertNotifier + from .redis_monitor import RedisMonitor + from .redis_service import RedisService + from .redis_service_factory import RedisServiceFactory + from .retry_policy import RedisHealthChecker, RedisRetry + from .task_redis_service import TaskRedisService + from .user_redis_service import UserRedisService __all__ = [ "RedisService", @@ -20,4 +28,44 @@ "UserRedisService", "JobInfoRedisService", "JobMetadataService", + "RedisKeyBuilder", + "RedisKeyType", + "redis_key_builder", + "RedisHealthChecker", + "RedisRetry", ] + +_EXPORT_MODULES: dict[str, str] = { + "RedisService": "shared.services.redis.redis_service", + "RedisServiceFactory": "shared.services.redis.redis_service_factory", + "RedisMonitor": "shared.services.redis.redis_monitor", + "RedisAlertManager": "shared.services.redis.redis_alerts", + "RedisAlertNotifier": "shared.services.redis.redis_alerts", + "AlertRule": "shared.services.redis.redis_alerts", + "TaskRedisService": "shared.services.redis.task_redis_service", + "UserRedisService": "shared.services.redis.user_redis_service", + "JobInfoRedisService": "shared.services.redis.job_info_redis_service", + "JobMetadataService": "shared.services.redis.job_metadata_service", + "RedisKeyBuilder": "shared.services.redis.key_builder", + "RedisKeyType": "shared.services.redis.key_builder", + "redis_key_builder": "shared.services.redis.key_builder", + "RedisHealthChecker": "shared.services.redis.retry_policy", + "RedisRetry": "shared.services.redis.retry_policy", +} + + +def __getattr__(name: str) -> Any: + """Load Redis package exports on first access.""" + module_name = _EXPORT_MODULES.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + module = import_module(module_name) + value = getattr(module, name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + """Return public Redis package exports.""" + return sorted([*globals(), *__all__]) diff --git a/packages/shared-python/shared/services/redis/distributed_lock.py b/packages/shared-python/shared/services/redis/distributed_lock.py index a36a23bf9..751309031 100644 --- a/packages/shared-python/shared/services/redis/distributed_lock.py +++ b/packages/shared-python/shared/services/redis/distributed_lock.py @@ -18,7 +18,7 @@ from shared.core.config import settings from shared.core.exceptions.domain_exceptions import UnavailableException from shared.services.redis.redis_sync_service import SyncRedisService -from shared.utils.redis_key_builder import redis_key_builder +from shared.services.redis.key_builder import redis_key_builder # Lua script: atomically check owner token then delete. # Prevents a stale owner from accidentally deleting a lock @@ -130,7 +130,7 @@ def __enter__(self) -> "RedisJobLock": internal_message=( f"Could not acquire processing lock for job {self._job_id}" ), - retry_after=settings.KB_TASK_RETRY_COUNTDOWN, + retry_after=settings.DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN, ) return self diff --git a/packages/shared-python/shared/services/redis/job_info_redis_service.py b/packages/shared-python/shared/services/redis/job_info_redis_service.py index 89c1a2ec2..0ee34b0b7 100644 --- a/packages/shared-python/shared/services/redis/job_info_redis_service.py +++ b/packages/shared-python/shared/services/redis/job_info_redis_service.py @@ -8,7 +8,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class JobInfoRedisService: diff --git a/packages/shared-python/shared/services/redis/job_metadata_service.py b/packages/shared-python/shared/services/redis/job_metadata_service.py index 62fb9c135..410a1c052 100644 --- a/packages/shared-python/shared/services/redis/job_metadata_service.py +++ b/packages/shared-python/shared/services/redis/job_metadata_service.py @@ -5,7 +5,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class JobMetadataService: diff --git a/packages/shared-python/shared/utils/redis_key_builder.py b/packages/shared-python/shared/services/redis/key_builder.py similarity index 91% rename from packages/shared-python/shared/utils/redis_key_builder.py rename to packages/shared-python/shared/services/redis/key_builder.py index 62d3d5b28..02d0dce78 100644 --- a/packages/shared-python/shared/utils/redis_key_builder.py +++ b/packages/shared-python/shared/services/redis/key_builder.py @@ -10,7 +10,6 @@ class RedisKeyType(Enum): USER = "user" TASK = "task" CONVERSATION = "conversation" - KNOWLEDGE_BASE = "kb" SESSION = "session" CACHE = "cache" QUEUE = "queue" @@ -106,24 +105,6 @@ def conversation_context(self, conversation_id: str) -> str: """Conversation context key.""" return self.build_key(RedisKeyType.CONVERSATION, conversation_id, "context") - # ==================== Knowledge Base Keys ==================== - - def kb_status(self, user_id: str) -> str: - """Knowledge-base status key.""" - return self.build_key(RedisKeyType.KNOWLEDGE_BASE, user_id, "status") - - def kb_vectors(self, user_id: str) -> str: - """Knowledge-base vectors key.""" - return self.build_key(RedisKeyType.KNOWLEDGE_BASE, user_id, "vectors") - - def kb_metadata(self, user_id: str) -> str: - """Knowledge-base metadata key.""" - return self.build_key(RedisKeyType.KNOWLEDGE_BASE, user_id, "metadata") - - def kb_index(self, user_id: str) -> str: - """Knowledge-base index key.""" - return self.build_key(RedisKeyType.KNOWLEDGE_BASE, user_id, "index") - # ==================== Session Keys ==================== def session_data(self, session_id: str) -> str: @@ -270,7 +251,6 @@ def get_key_ttl(self, key_type: RedisKeyType) -> int: RedisKeyType.USER: 86400, # 1 day (user_config cache). RedisKeyType.TASK: 86400, # 1 day. RedisKeyType.CONVERSATION: 3600 * 2, # 2 hours. - RedisKeyType.KNOWLEDGE_BASE: 86400 * 30, # 30 days. RedisKeyType.SESSION: 3600, # 1 hour. RedisKeyType.CACHE: 3600, # 1 hour. RedisKeyType.QUEUE: 86400, # 1 day. diff --git a/packages/shared-python/shared/services/redis/rate_limit_service.py b/packages/shared-python/shared/services/redis/rate_limit_service.py index de7ef821b..c1cec343b 100644 --- a/packages/shared-python/shared/services/redis/rate_limit_service.py +++ b/packages/shared-python/shared/services/redis/rate_limit_service.py @@ -7,7 +7,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import redis_key_builder +from shared.services.redis.key_builder import redis_key_builder class RateLimitService: diff --git a/packages/shared-python/shared/services/redis/redis_monitor.py b/packages/shared-python/shared/services/redis/redis_monitor.py index 6ae090590..f6887576a 100644 --- a/packages/shared-python/shared/services/redis/redis_monitor.py +++ b/packages/shared-python/shared/services/redis/redis_monitor.py @@ -7,7 +7,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import redis_key_builder +from shared.services.redis.key_builder import redis_key_builder class RedisMonitor: diff --git a/packages/shared-python/shared/services/redis/redis_service.py b/packages/shared-python/shared/services/redis/redis_service.py index 99ed39c1d..fa9d39fa9 100644 --- a/packages/shared-python/shared/services/redis/redis_service.py +++ b/packages/shared-python/shared/services/redis/redis_service.py @@ -13,7 +13,7 @@ RedisConnectionError, RedisOperationError, ) -from shared.utils.redis_retry import RedisHealthChecker, RedisRetry +from shared.services.redis.retry_policy import RedisHealthChecker, RedisRetry ResponseT = TypeVar("ResponseT") diff --git a/packages/shared-python/shared/services/redis/redis_sync_service.py b/packages/shared-python/shared/services/redis/redis_sync_service.py index aa725b314..90e363aa1 100644 --- a/packages/shared-python/shared/services/redis/redis_sync_service.py +++ b/packages/shared-python/shared/services/redis/redis_sync_service.py @@ -12,7 +12,7 @@ from redis.connection import BlockingConnectionPool from shared.core.config.redis import RedisConfigManager -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class SyncRedisService: diff --git a/packages/shared-python/shared/utils/redis_retry.py b/packages/shared-python/shared/services/redis/retry_policy.py similarity index 100% rename from packages/shared-python/shared/utils/redis_retry.py rename to packages/shared-python/shared/services/redis/retry_policy.py diff --git a/packages/shared-python/shared/services/redis/task_redis_service.py b/packages/shared-python/shared/services/redis/task_redis_service.py index 883d8a116..6c9444534 100644 --- a/packages/shared-python/shared/services/redis/task_redis_service.py +++ b/packages/shared-python/shared/services/redis/task_redis_service.py @@ -5,7 +5,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class TaskRedisService: diff --git a/packages/shared-python/shared/services/redis/user_redis_service.py b/packages/shared-python/shared/services/redis/user_redis_service.py index 6db0a2acc..d28ee59d5 100644 --- a/packages/shared-python/shared/services/redis/user_redis_service.py +++ b/packages/shared-python/shared/services/redis/user_redis_service.py @@ -5,7 +5,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class UserRedisService: diff --git a/packages/shared-python/shared/services/retrieval/__init__.py b/packages/shared-python/shared/services/retrieval/__init__.py index 832e81eb6..845269856 100644 --- a/packages/shared-python/shared/services/retrieval/__init__.py +++ b/packages/shared-python/shared/services/retrieval/__init__.py @@ -1,27 +1,7 @@ -from .app_service import merge_channels_rrf, run_retrieval_query -from .cache_service import ( - bump_retrieval_namespace_cache_version, - get_cached_retrieval_query_result, - get_retrieval_namespace_cache_version, - invalidate_retrieval_cache_namespaces, - set_cached_retrieval_query_result, -) -from .graph_service import DocumentGraphService, GraphQueryService, GraphScope -from .hit_stats_service import record_retrieval_hits -from .llm_adapter import create_retrieval_llm_fn, create_retrieval_planner_fn +"""Retrieval service package. -__all__ = [ - "create_retrieval_llm_fn", - "create_retrieval_planner_fn", - "run_retrieval_query", - "merge_channels_rrf", - "DocumentGraphService", - "GraphQueryService", - "GraphScope", - "record_retrieval_hits", - "bump_retrieval_namespace_cache_version", - "get_cached_retrieval_query_result", - "get_retrieval_namespace_cache_version", - "invalidate_retrieval_cache_namespaces", - "set_cached_retrieval_query_result", -] +Import concrete modules directly so importing one retrieval submodule does not +initialize unrelated retrieval dependencies. +""" + +from __future__ import annotations diff --git a/packages/shared-python/shared/services/retrieval/agent_navigate.py b/packages/shared-python/shared/services/retrieval/agent_navigate.py deleted file mode 100644 index d5b1bed78..000000000 --- a/packages/shared-python/shared/services/retrieval/agent_navigate.py +++ /dev/null @@ -1,1137 +0,0 @@ -"""Shared helpers for agentic KG document routing and scope navigation.""" -from __future__ import annotations - -import json -import re -from typing import Any, Sequence, TYPE_CHECKING - -if TYPE_CHECKING: - from shared.services.retrieval.agentic.types import DocTreeNode - -from loguru import logger -from sqlalchemy import func, select, or_ -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.models.database.document import Document, DocumentChunk, DocumentSection, GraphNode, GraphEdge -from shared.services.retrieval.lexical_text import normalize_section_path, split_section_path -from shared.utils.text_utils import tokenize_for_retrieval - -_MAX_OVERVIEW_FILES = 50 - -_FILE_SELECT_PROMPT = """\ -You are a document routing assistant. - -{budget_block} -Below is a knowledge base overview showing all available documents, -their navigation summaries, chunk counts, and media counts. - -=== Knowledge Base Overview === -{overview} -=== End Overview === - -User query: {query} -{revision_context} -Based on the query, select documents that may contain relevant information. -If NO document in the knowledge base is relevant to the query, return an EMPTY array []. -Return ONLY a JSON array of document IDs, e.g.: ["doc_abc123", "doc_def456"] -Do not include any explanation. -""" - - - -_DISCOVERY_SELECT_PROMPT = """\ -You are a document navigation assistant. - -Document: "{doc_name}" - -{budget_block} -After navigating the document's section tree, the following section paths -were additionally discovered via keyword and semantic search. -They may contain relevant evidence not found through hierarchical navigation. - -=== Discovery Candidates === -{items} -=== End Discovery Candidates === - -User query: {query} -{revision_context} -Select section paths whose content is needed to answer the query. -If none are relevant, return an EMPTY list []. - -Return ONLY a JSON object: -{{"selections": [{{"path": "...", "confidence": }}, ...]}} -Do not include any explanation. -""" - - -_ACTION_PROMPT = """\ -You are a document navigation agent. - -Document: "{doc_name}" (id: {doc_id}) - -{budget_block} -{scope_header} -Below is the document's section tree. -Sections tagged [SELECT] are within the current scope and may be selected. -Other sections are shown as structural context only (not selectable). -Nodes marked [Leaf] have no further sub-sections. - -=== Section Tree === -{items_overview} -=== End Section Tree === - -User query: {query} - -=== Available Actions === - -Choose ONE action: - -NAVIGATE — Drill into selected sections for detailed content. - Consider this when the query targets specific topics and you need deeper text evidence. - Select one or more [SELECT] sections. - -STOP — Current scope evidence is sufficient. No further drill-down. - Consider this when: - - The query asks for an outline, overview, or summary - - The query is broad/global, the tree section can fulfill it without drilling into individual sections. - - You have already collected enough evidence at this level. - -{tools_block} - -When action is NAVIGATE, provide selections: -- You may ONLY select sections marked with [SELECT]. - -When action is STOP, selections must be empty. - -Return ONLY a JSON object: -{{"action": "NAVIGATE", "tools": [...], "selections": [{{"path": "...", "confidence": }}, ...]}} -or -{{"action": "STOP", "tools": [...], "selections": []}} -Do not include any explanation. -""" - - -def _parse_action_response(text: str) -> dict: - """Parse the unified action response from LLM. - - Returns dict with keys: - action: 'NAVIGATE' | 'STOP' - tools: list[str] (subset of FIND_IMAGES, FIND_TABLES) - selections: list[dict] (each has 'path' and optional 'confidence') - - When action is STOP, selections are forced to empty. - """ - import json as _json - import re as _re - - text = text.strip() - _ASSET_TOOLS = {'FIND_IMAGES', 'FIND_TABLES'} - default = {'action': 'NAVIGATE', 'tools': [], 'selections': []} - - def _extract(data: dict) -> dict: - action = str(data.get('action', 'NAVIGATE')).strip().upper() - if action not in ('NAVIGATE', 'STOP'): - action = 'NAVIGATE' - - tools_val = data.get('tools') or [] - if isinstance(tools_val, list): - tools = [str(t).strip().upper() for t in tools_val if str(t).strip().upper() in _ASSET_TOOLS] - else: - tools = [] - - # STOP → no selections allowed - if action == 'STOP': - return {'action': action, 'tools': tools, 'selections': []} - - selections_val = data.get('selections') or [] - selections = [] - if isinstance(selections_val, list): - for s in selections_val: - if isinstance(s, dict) and s.get('path'): - conf = _normalize_confidence(s.get('confidence', 0.7)) - selections.append({'path': str(s['path']), 'confidence': conf or 0.7}) - - return {'action': action, 'tools': tools, 'selections': selections} - - # Try JSON parse - try: - data = _json.loads(text) - if isinstance(data, dict): - return _extract(data) - except (ValueError, _json.JSONDecodeError): - pass - - # Try extracting JSON from markdown fences - fence_match = _re.search(r'```(?:json)?\s*\n?(.*?)\n?```', text, _re.DOTALL) - if fence_match: - try: - data = _json.loads(fence_match.group(1).strip()) - if isinstance(data, dict): - return _extract(data) - except (ValueError, _json.JSONDecodeError): - pass - - # Try finding a JSON object anywhere - brace_match = _re.search(r'\{.*\}', text, _re.DOTALL) - if brace_match: - try: - data = _json.loads(brace_match.group()) - if isinstance(data, dict): - return _extract(data) - except (ValueError, _json.JSONDecodeError): - pass - - return default - - -def _format_budget_block(snapshot: dict | None) -> str: - if not snapshot: - return "" - planning = snapshot.get("planning") or {} - context = snapshot.get("context") or {} - return ( - "=== Resource Status ===\n" - f"Planning Budget: {planning.get('status', 'HEALTHY')} " - f"({planning.get('used_pct', 0)}% used)\n" - f"Context Budget: {context.get('status', 'HEALTHY')} " - f"({context.get('used_pct', 0)}% used)\n" - f"KG Coverage: {snapshot.get('explored_chunks', 0)}/" - f"{snapshot.get('total_chunks', 0)} chunks explored\n" - f"Docs Explored: {snapshot.get('explored_docs', 0)}/" - f"{snapshot.get('total_docs', 0)}\n" - "When budget is TIGHT, prefer fewer high-confidence selections over broad exploration. " - "When CRITICAL, be very selective — only pick paths with strong relevance. Return empty if evidence suffices.\n" - "=== End Resource Status ===\n" - ) - - -def _extract_json_array_payload(text: str) -> list[Any]: - """Best-effort extraction of a JSON array payload from LLM response text.""" - text = text.strip() - try: - result = json.loads(text) - if isinstance(result, list): - return result - except (json.JSONDecodeError, ValueError): - pass - match = re.search(r'\[.*?\]', text, re.DOTALL) - if match: - try: - result = json.loads(match.group()) - if isinstance(result, list): - return result - except (json.JSONDecodeError, ValueError): - pass - return [] - - -def _parse_json_array(text: str) -> list[str]: - """Best-effort extraction of a JSON array of strings from LLM response text.""" - result = _extract_json_array_payload(text) - return [str(x) for x in result] - - -def _normalize_confidence(value: Any) -> float | None: - if value is None: - return None - if isinstance(value, str): - value = value.strip().rstrip('%') - try: - parsed = float(value) - except (TypeError, ValueError): - return None - if parsed > 1.0: - parsed = parsed / 100.0 - return max(0.0, min(parsed, 1.0)) - - -async def _build_knowledge_map_overview( - db: AsyncSession, - *, - user_id: str, - namespace: str, -) -> tuple[str, dict[str, str]]: - """Build a file-level knowledge map overview for LLM file selection. - - Returns (overview_text, doc_id_to_name) where doc_id_to_name maps - document_id -> source_file_name for validation after LLM response. - """ - doc_stmt = ( - select(Document) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where(Document.current_job_result_id.is_not(None)) - .order_by(Document.updated_at.desc()) - .limit(_MAX_OVERVIEW_FILES) - ) - doc_result = await db.execute(doc_stmt) - documents = list(doc_result.scalars()) - - if not documents: - return '(empty)', {} - - doc_ids = [d.document_id for d in documents] - doc_id_to_name: dict[str, str] = { - d.document_id: (d.source_file_name or d.document_id) - for d in documents - } - - chunk_stats_stmt = ( - select( - DocumentChunk.document_id, - func.count(DocumentChunk.id).label('chunk_count'), - func.count(func.nullif(DocumentChunk.chunk_type, 'text')).label('media_count'), - ) - .join(Document, (Document.document_id == DocumentChunk.document_id) & (Document.current_job_result_id == DocumentChunk.job_result_id)) - .where(DocumentChunk.document_id.in_(doc_ids)) - .group_by(DocumentChunk.document_id) - ) - chunk_stats_result = await db.execute(chunk_stats_stmt) - chunk_stats: dict[str, dict[str, int]] = {} - for row in chunk_stats_result.all(): - chunk_stats[row[0]] = {'total': row[1], 'media': row[2]} - - graph_summary_stmt = ( - select(GraphNode.owner_document_id, GraphNode.properties) - .where(GraphNode.owner_document_id.in_(doc_ids)) - .where(GraphNode.node_kind == 'document') - ) - graph_summary_result = await db.execute(graph_summary_stmt) - doc_top_summaries: dict[str, str] = {} - for did, properties in graph_summary_result.all(): - if not isinstance(properties, dict): - continue - top_summary = str(properties.get('top_summary') or '').strip() - if top_summary: - doc_top_summaries[did] = top_summary - - lines: list[str] = [] - for doc in documents: - did = doc.document_id - name = doc_id_to_name[did] - stats = chunk_stats.get(did, {'total': 0, 'media': 0}) - top_summary = doc_top_summaries.get(did, '') - - line = f'- [{did}] {name} chunks={stats["total"]}' - if stats['media'] > 0: - line += f' media={stats["media"]}' - if top_summary: - line += f'\n top_summary:\n{_indent_block(top_summary, 4)}' - lines.append(line) - - return '\n'.join(lines), doc_id_to_name - - -def _indent_block(text: str, spaces: int) -> str: - prefix = ' ' * spaces - return '\n'.join(f'{prefix}{line}' for line in str(text or '').splitlines()) - - -def _format_items_for_llm( - items: list[dict], - max_chars: int = 20000, -) -> tuple[str, bool]: - """Format items with ▸ └ [Leaf] hierarchy for scope navigation. - - Supports arbitrary depth levels via absolute ``level`` field. - Items with ``show_summary=False`` render title only (structural context). - ``[LN]`` tags indicate the absolute document depth of each section. - ``[Leaf]`` tags indicate bottom-level sections with no further children. - Summaries are included when within budget, dropped on overflow. - - Returns (text, overflowed). - """ - from shared.utils.text_utils import truncate_content_preview - - if not items: - return '(no items available)', False - - SUMMARY_HEAD_TOKENS = 80 - - def _render_item(item: dict, include_summary: bool) -> str: - level = item.get('level', 1) - show = item.get('show_summary', True) - is_leaf = item.get('is_leaf', False) - leaf_tag = ' [Leaf]' if is_leaf else '' - path = item.get('path', '') - summary = item.get('summary') or '' - - # Build chunk count tags (only for current-scope items) - counts_str = '' - if show: - count_parts: list[str] = [] - chunk_count = item.get('chunk_count', 0) - if chunk_count > 0: - count_parts.append(f'text={chunk_count}') - image_count = item.get('image_count', 0) - if image_count > 0: - count_parts.append(f'image={image_count}') - table_count = item.get('table_count', 0) - if table_count > 0: - count_parts.append(f'table={table_count}') - counts_str = f' [{" ".join(count_parts)}]' if count_parts else '' - - indent = " " * (level - 1) - prefix = '▸' if level == 1 else '└' - level_tag = f'[L{level}]' - select_tag = '[SELECT] ' if item.get('selectable', False) else '' - - lines: list[str] = [] - lines.append(f'{indent}{prefix} {select_tag}{level_tag} path="{path}"{counts_str}{leaf_tag}') - - if include_summary and show and summary: - sub_indent = " " * level - clipped = truncate_content_preview(summary, head=SUMMARY_HEAD_TOKENS, tail=0) - lines.append(f'{sub_indent}{clipped}') - - return '\n'.join(lines) - - # Try full render (with summaries for show_summary=True items) - full_lines = [_render_item(item, include_summary=True) for item in items] - full_text = '\n'.join(full_lines) - if len(full_text) <= max_chars: - return full_text, False - - # Overflow: render without summaries - slim_lines = [_render_item(item, include_summary=False) for item in items] - slim_text = '\n'.join(slim_lines) - return slim_text[:max_chars], True - - -# ------------------------------------------------------------------ -# GREP document discovery (aligned with KB do_discover_files) -# ------------------------------------------------------------------ - -async def _grep_discover_document_ids( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - exclude_document_ids: Sequence[str] = (), - limit: int = 10, -) -> list[str]: - """GREP discovery: search term_search_text for query terms, return parent document_ids. - - Aligned with KB's do_discover_files(): if a chunk's term_search_text - contains query terms, its parent document is included in the KG scope. - """ - units = tokenize_for_retrieval(query, dedupe=True) - logger.info(f' GREP tokenized units (cap 8): {units[:8]} (total={len(units)})') - if not units: - return [] - - # Build OR conditions for ILIKE matching - conditions = [] - params: dict[str, str] = { - 'user_id': user_id, - 'namespace': namespace, - } - for i, unit in enumerate(units[:8]): # cap at 8 terms to avoid huge queries - param_name = f'unit_{i}' - params[param_name] = f'%{unit}%' - conditions.append(DocumentChunk.term_search_text.ilike(f'%{unit}%')) - - if not conditions: - return [] - - stmt = ( - select(Document.document_id) - .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id)) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where(DocumentChunk.term_search_text.is_not(None)) - .where(or_(*conditions)) - .distinct() - .limit(limit) - ) - if exclude_document_ids: - stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) - - result = await db.execute(stmt) - return [row[0] for row in result.all()] - - -# ------------------------------------------------------------------ -# Edge expansion (aligned with KB KGIndex.neighbors) -# ------------------------------------------------------------------ - -async def _expand_by_edges( - db: AsyncSession, - *, - document_ids: list[str], - user_id: str, - namespace: str, - hops: int = 1, -) -> list[str]: - """Expand document set by following GraphEdge relationships. - - Aligned with KB's KGIndex.neighbors(): traverse edges to include - related documents. Only queries document-level nodes (no section nodes). - No weight filtering — edges already passed threshold during publication. - """ - if not document_ids: - return document_ids - - current = set(document_ids) - - for hop_idx in range(hops): - # Find document-level graph nodes for current document set - doc_node_ids = [f"doc:{did}" for did in current] - node_stmt = ( - select(GraphNode.node_id, GraphNode.owner_document_id) - .where(GraphNode.user_id == user_id) - .where(GraphNode.namespace == namespace) - .where(GraphNode.node_kind == 'document') - .where(GraphNode.node_id.in_(doc_node_ids)) - ) - node_result = await db.execute(node_stmt) - node_rows = node_result.all() - logger.info(f' edge_expand hop={hop_idx}: doc_nodes_found={len(node_rows)} (of {len(doc_node_ids)} requested)') - - if not node_rows: - break - - node_ids = {row[0] for row in node_rows} - - # Follow edges from/to these document nodes - edge_stmt = ( - select(GraphEdge.source_node_id, GraphEdge.target_node_id) - .where(GraphEdge.user_id == user_id) - .where(GraphEdge.namespace == namespace) - .where(or_( - GraphEdge.source_node_id.in_(list(node_ids)), - GraphEdge.target_node_id.in_(list(node_ids)), - )) - ) - edge_result = await db.execute(edge_stmt) - edge_rows = edge_result.all() - - neighbor_node_ids: set[str] = set() - for src, tgt in edge_rows: - if src in node_ids: - neighbor_node_ids.add(tgt) - if tgt in node_ids: - neighbor_node_ids.add(src) - logger.info(f' edge_expand hop={hop_idx}: edges_traversed={len(edge_rows)} neighbor_nodes={len(neighbor_node_ids)}') - - if not neighbor_node_ids: - break - - # Resolve neighbor nodes to document_ids - neighbor_doc_stmt = ( - select(GraphNode.owner_document_id) - .where(GraphNode.node_id.in_(list(neighbor_node_ids))) - .where(GraphNode.node_kind == 'document') - ) - neighbor_doc_result = await db.execute(neighbor_doc_stmt) - for (doc_id,) in neighbor_doc_result.all(): - current.add(doc_id) - - # Preserve original order, append new ones at end - ordered = list(document_ids) - for doc_id in current: - if doc_id not in document_ids: - ordered.append(doc_id) - return ordered - - -# --------------------------------------------------------------------------- -# Unified scope navigation: load child sections (2-level) -# --------------------------------------------------------------------------- - -async def _load_child_sections( - db: AsyncSession, - document_id: str, - job_result_id: str, - scope_path: str | list[str] | None = None, - exclude_paths: set[str] | None = None, -) -> list[dict]: - """Load the Continuous Context Tree for *scope_path*. - - Returns a flat list sorted by document order, each item: - {path, title, summary, chunk_count, image_count, table_count, - level, show_summary, is_leaf} - - scope_path can be: - - None: root scope, all items are selectable (2 depth bands). - - str: single scope, descendants are selectable. - - list[str]: multi-scope, descendants of ALL paths are selectable - simultaneously — used when the LLM selected multiple drill-down - paths in the previous step. - - - level: absolute depth in the document (1-based) - - show_summary: controls whether _format_items_for_llm renders summary - - exclude_paths: paths already hydrated; skipped from selectable items - """ - # ── Fetch all sections for this document revision ──────────────────── - stmt = ( - select( - DocumentSection.section_id, - DocumentSection.section_title, - DocumentSection.section_path, - DocumentSection.summary, - DocumentSection.sort_order, - ) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - .order_by(DocumentSection.sort_order) - ) - section_rows = (await db.execute(stmt)).all() - if not section_rows: - return [] - - # ── Normalize scope(s) ─────────────────────────────────────────────── - # Multi-scope: list of paths to expand simultaneously - if isinstance(scope_path, list): - scope_list = [normalize_section_path(p) for p in scope_path] - elif scope_path: - scope_list = [normalize_section_path(scope_path)] - else: - scope_list = [] # root - - # For logging, derive representative scope info - scope_depth = len(split_section_path(scope_list[0])) if scope_list else 0 - - logger.debug( - f' _load_child_sections: scopes={scope_list or ["root"]} ' - f'scope_depth={scope_depth} exclude_paths={_excl if (_excl := exclude_paths or set()) else "none"} ' - f'total_sections={len(section_rows)}' - ) - - # Build full section metadata index - all_sections: dict[str, dict] = {} # path → {title, summary, sort_order, section_id, parts, depth} - for section_id, title, path, summary, sort_order in section_rows: - if not path: - continue - path = normalize_section_path(path) - parts = split_section_path(path) - all_sections[path] = { - 'title': title or parts[-1] if parts else path, - 'summary': summary or '', - 'sort_order': int(sort_order or 0), - 'section_id': section_id, - 'parts': parts, - 'depth': len(parts), - } - - # ── Build the set of ancestor prefixes for pruning ──────────────────── - # For multi-scope, union all ancestor prefixes from all scope paths - ancestor_prefixes: set[str] = set() - for sp in scope_list: - sp_parts = split_section_path(sp) - for i in range(1, len(sp_parts) + 1): - ancestor_prefixes.add(' / '.join(sp_parts[:i])) - - # ── Classify each section ──────────────────────────────────────────── - _excl = exclude_paths or set() - items_by_path: dict[str, dict] = {} - # Per-scope depth bands: track child depths separately per scope - per_scope_child_depths: dict[str, set[int]] = {sp: set() for sp in scope_list} if scope_list else {} - root_child_depths: set[int] = set() # used when scope_list is empty (root) - - def _make_item(path: str, meta: dict, show_summary: bool) -> dict: - return { - 'path': path, - 'title': meta['title'], - 'summary': meta['summary'], - 'level': meta['depth'], - 'sort_order': meta['sort_order'], - 'chunk_count': 0, - 'image_count': 0, - 'table_count': 0, - 'section_id': meta['section_id'], - 'show_summary': show_summary, - } - - def _is_excluded(path: str) -> bool: - return bool(_excl and any( - path == ep or path.startswith(ep + ' / ') for ep in _excl - )) - - for path, meta in all_sections.items(): - parts = meta['parts'] - depth = meta['depth'] - - if not scope_list: - # Root scope: everything is a potential child - if depth < 1 or _is_excluded(path): - continue - root_child_depths.add(depth) - items_by_path[path] = _make_item(path, meta, show_summary=True) - continue - - # --- Non-root scope(s) --- - # Check if this path is a descendant of ANY scope in scope_list - matched_scope: str | None = None - for sp in scope_list: - sp_parts = split_section_path(sp) - sp_depth = len(sp_parts) - if depth > sp_depth and parts[:sp_depth] == sp_parts: - matched_scope = sp - break - - if matched_scope: - # Category 2: descendant of a scope path → selectable - if _is_excluded(path): - continue - per_scope_child_depths[matched_scope].add(depth) - items_by_path[path] = _make_item(path, meta, show_summary=True) - continue - - # Category 1: structural context (ancestors of scope paths only) - # Only show nodes that are on the ancestor chain of a scope path. - # Non-scope siblings (e.g. 法律声明, 前言 when navigating into - # chapters 一~六) are pruned to reduce token waste and prevent - # summary overflow in _format_items_for_llm. - max_scope_depth = max(len(split_section_path(sp)) for sp in scope_list) - if depth <= max_scope_depth: - if depth == 1: - if path in ancestor_prefixes: - items_by_path.setdefault(path, _make_item(path, meta, show_summary=False)) - else: - parent_prefix = ' / '.join(parts[:-1]) - if parent_prefix in ancestor_prefixes: - items_by_path.setdefault(path, _make_item(path, meta, show_summary=False)) - continue - - # Category 3: pruned - - if not items_by_path: - return [] - - # ── Limit children to 2 depth bands (relative to each scope) ──────── - allowed_set: set[int] = set() - if scope_list: - for sp, depths in per_scope_child_depths.items(): - if depths: - allowed_set.update(sorted(depths)[:2]) - else: - if root_child_depths: - allowed_set.update(sorted(root_child_depths)[:2]) - - if allowed_set: - to_remove = [ - path for path, item in items_by_path.items() - if item['show_summary'] and item['level'] not in allowed_set - ] - for path in to_remove: - del items_by_path[path] - - if not items_by_path: - return [] - - # ── Count chunks per section (text / image / table) ────────────────── - # Only count for show_summary=True items (current scope children) - scope_item_sids = {item['section_id'] for item in items_by_path.values() if item['show_summary']} - # Also need all section_ids for upward aggregation - all_section_ids = [meta['section_id'] for meta in all_sections.values()] - if all_section_ids and scope_item_sids: - from sqlalchemy import case, literal_column - chunk_stmt = ( - select( - DocumentChunk.section_id, - func.count( - case( - (DocumentChunk.chunk_type.notin_(['image', 'table']), literal_column('1')), - ) - ).label('text_count'), - func.count( - case( - (DocumentChunk.chunk_type == 'image', literal_column('1')), - ) - ).label('image_count'), - func.count( - case( - (DocumentChunk.chunk_type == 'table', literal_column('1')), - ) - ).label('table_count'), - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(all_section_ids)) - .group_by(DocumentChunk.section_id) - ) - chunk_rows = (await db.execute(chunk_stmt)).all() - section_id_counts: dict[str, tuple[int, int, int]] = { - sid: (int(tc), int(ic), int(tbc)) for sid, tc, ic, tbc in chunk_rows - } - else: - section_id_counts = {} - - # Build section_id → path mapping for aggregation - sid_to_path = {meta['section_id']: path for path, meta in all_sections.items()} - - # Aggregate chunk counts upward: each show_summary item gets counts from itself + descendants - # Phase 1: Direct section assignment — counts from chunks directly under each section - for sid, (text_c, img_c, tbl_c) in section_id_counts.items(): - chunk_path = sid_to_path.get(sid, '') - if not chunk_path: - continue - - for item_path, item in items_by_path.items(): - if not item['show_summary']: - continue - if chunk_path == item_path or chunk_path.startswith(item_path + ' / '): - item['chunk_count'] += text_c - item['image_count'] += img_c - item['table_count'] += tbl_c - - # Phase 2: connect_to reference tracing — Root-level standalone assets - # Images/tables often live in the Root section but are referenced via connect_to - # from text chunks in deeper sections. Trace these references to attribute - # assets to the sections that actually use them. - # - # Algorithm: for each show_summary item, find all text chunks under its subtree, - # collect their connect_to targets, and count how many are image/table chunks. - scope_items_with_zero_assets = [ - item for item in items_by_path.values() - if item['show_summary'] and item['image_count'] == 0 and item['table_count'] == 0 - ] - if scope_items_with_zero_assets: - # Load connect_to metadata for text chunks under all scope sections - scope_section_ids = {item['section_id'] for item in items_by_path.values() if item.get('section_id')} - if scope_section_ids: - from sqlalchemy import literal_column - connect_stmt = ( - select( - DocumentChunk.section_id, - DocumentChunk.chunk_metadata, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(list(scope_section_ids))) - .where(DocumentChunk.chunk_type == 'text') - ) - connect_result = (await db.execute(connect_stmt)).all() - - # Map section_id → set of connected target chunk_ids - section_target_ids: dict[str, set[str]] = {} - for sec_id, metadata in connect_result: - if not isinstance(metadata, dict): - continue - for conn in metadata.get('connect_to') or []: - target_id = conn.get('target', '') - if target_id: - section_target_ids.setdefault(sec_id, set()).add(target_id) - - if section_target_ids: - # Collect all target chunk_ids and look up their types - all_target_ids = set() - for tids in section_target_ids.values(): - all_target_ids.update(tids) - - target_type_stmt = ( - select( - DocumentChunk.chunk_id, - DocumentChunk.chunk_type, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.chunk_id.in_(list(all_target_ids))) - .where(DocumentChunk.chunk_type.in_(['image', 'table'])) - ) - target_type_result = (await db.execute(target_type_stmt)).all() - target_types: dict[str, str] = {cid: ctype for cid, ctype in target_type_result} - - # Aggregate connected asset counts per section path → upward to items - for sec_id, target_ids in section_target_ids.items(): - ref_path = sid_to_path.get(sec_id, '') - if not ref_path: - continue - ref_img = sum(1 for tid in target_ids if target_types.get(tid) == 'image') - ref_tbl = sum(1 for tid in target_ids if target_types.get(tid) == 'table') - if ref_img == 0 and ref_tbl == 0: - continue - for item_path, item in items_by_path.items(): - if not item['show_summary']: - continue - if ref_path == item_path or ref_path.startswith(item_path + ' / '): - item['image_count'] += ref_img - item['table_count'] += ref_tbl - - # ── Sort by native document order ───────────────────────────────────── - sorted_items = sorted(items_by_path.values(), key=lambda x: x['sort_order']) - # Clean up internal fields - for item in sorted_items: - item.pop('sort_order', None) - item.pop('section_id', None) - - # ── Detect leaf status ──────────────────────────────────────────────── - # A section is a leaf if no other section in the database for this - # document has a path that descends from it. - all_section_paths = set(all_sections.keys()) - for item in sorted_items: - item_path = item['path'] - has_descendants = any( - p != item_path and p.startswith(item_path + ' / ') - for p in all_section_paths - ) - item['is_leaf'] = not has_descendants - - # ── Assign selectability ────────────────────────────────────────────── - # Rule: in the 2-band window, only the DEEPER band is selectable. - # Leaf nodes at the shallower band are still selectable (no children - # to drill into). Structural context (show_summary=False) is never - # selectable. - if allowed_set: - shallowest_band = min(allowed_set) - for item in sorted_items: - if not item.get('show_summary', True): - # Structural context → never selectable - item['selectable'] = False - elif item['level'] == shallowest_band and not item.get('is_leaf', False): - # Shallowest band, non-leaf → grouping header, not selectable - item['selectable'] = False - else: - item['selectable'] = True - else: - for item in sorted_items: - item['selectable'] = item.get('show_summary', True) - - return sorted_items - - -# --------------------------------------------------------------------------- - - -# --------------------------------------------------------------------------- -# Unified document tree rendering (DocTreeNode → single coherent hierarchy) -# --------------------------------------------------------------------------- - -def _render_leaf_chunks( - parts: list[str], - chunks: list[dict[str, Any]], - indent: str, - asset_lookup: dict[str, str] | None = None, -) -> None: - """Render hydrated leaf chunks inline with table/image inlining and dedup. - - Uses ``connect_to`` metadata to resolve asset references — the same - pattern as ``assemble_retrieval_results``: - - **Tables**: inline HTML content at the ``ref`` placeholder - - **Images**: inline the ``file_path`` (S3-compatible URL) at the - placeholder for multimodal LLMs - - Connected target chunks (images/tables) are expected to already be - present in ``chunks`` via ``hydrate_connected_target_rows``. - - Phase 2: After rendering all text chunks, standalone image/table - chunks that were NOT inlined via connect_to are rendered separately. - This handles cases where assets exist at root/section level without - a parent text chunk referencing them. - """ - chunk_by_id: dict[str, dict] = { - c.get('chunk_id', ''): c for c in chunks if c.get('chunk_id') - } - rendered_ids: set[str] = set() - - # Phase 1: Render text chunks with inline asset resolution - for chunk in chunks: - cid = chunk.get('chunk_id', '') - if cid and cid in rendered_ids: - continue - - chunk_type = (chunk.get('chunk_type') or chunk.get('type') or 'text').strip().lower() - - # Skip standalone image/table chunks — they'll be rendered in Phase 2 - # if not inlined via connect_to from a parent text chunk. - # NOTE: do NOT add to rendered_ids here — Phase 2 needs to see them. - if chunk_type in ('image', 'table'): - continue - - if cid: - rendered_ids.add(cid) - - content = str(chunk.get('content', '')).strip() - - # Resolve connected assets via connect_to metadata - for conn in (chunk.get('chunk_metadata') or {}).get('connect_to') or []: - target = chunk_by_id.get(conn.get('target', '')) - if not target: - continue - target_cid = target.get('chunk_id', '') - target_type = (target.get('chunk_type') or target.get('type') or '').strip().lower() - ref_str = conn.get('ref', '') - if not ref_str or ref_str not in content: - continue - - if target_cid: - rendered_ids.add(target_cid) - - if target_type == 'table': - table_html = str(target.get('content', '')).strip() - content = content.replace(ref_str, f'\n[表格内容]\n{table_html}\n') - elif target_type == 'image': - file_path = target.get('file_path') or '' - img_desc = str(target.get('content', '')).strip() - # Strip self-reference from image description - if ref_str in img_desc: - img_desc = img_desc.replace(ref_str, '').strip() - # Use pre-generated asset URL if available, fall back to file_path - asset_url = (asset_lookup or {}).get(target_cid, '') if target_cid else '' - display_ref = asset_url or file_path - if display_ref: - content = content.replace(ref_str, f'\n[图片: {display_ref}]\n{img_desc}\n') - elif img_desc: - content = content.replace(ref_str, f'\n[图片描述]\n{img_desc}\n') - - for line in content.split('\n'): - if line.strip(): - parts.append(f'{indent}┈ {line}') - - # Phase 2: Render standalone image/table chunks not inlined via connect_to - for chunk in chunks: - cid = chunk.get('chunk_id', '') - if cid and cid in rendered_ids: - continue - if cid: - rendered_ids.add(cid) - - chunk_type = (chunk.get('chunk_type') or chunk.get('type') or '').strip().lower() - if chunk_type == 'image': - file_path = chunk.get('file_path') or '' - img_desc = str(chunk.get('content', '')).strip() - asset_url = (asset_lookup or {}).get(cid, '') if cid else '' - display_ref = asset_url or file_path - if display_ref: - parts.append(f'{indent}┈ [图片: {display_ref}]') - if img_desc: - for line in img_desc.split('\n'): - if line.strip(): - parts.append(f'{indent}┈ {line}') - elif chunk_type == 'table': - table_html = str(chunk.get('content', '')).strip() - parts.append(f'{indent}┈ [表格内容]') - if table_html: - for line in table_html.split('\n'): - if line.strip(): - parts.append(f'{indent}┈ {line}') - - -def render_unified_doc_tree( - node: DocTreeNode, - doc_name: str, - depth: int = 0, - asset_lookup: dict[str, str] | None = None, -) -> str: - """Render a DocTreeNode as a single coherent hierarchy. - - Summaries are navigation-only aids and NEVER appear in evidence. - The rendered output contains: - 1. Structural titles for ALL sections (positioning context) - 2. Hydrated chunk content (┈ lines) ONLY for selected leaf paths - - Asset references (tables/images) are resolved via ``connect_to`` - metadata in hydrated chunks — no separate lookup needed. - """ - - parts: list[str] = [] - indent = ' ' * depth - - if depth == 0: - parts.append(f'【文档】{doc_name}\n') - - # Collect children keys for path-hierarchy dedup: - child_prefixes = set(node.children.keys()) - - # Helper: min sort_order of a leaf_content entry - def _min_sort(path: str) -> float: - chunks = node.leaf_content.get(path, []) - return min((c.get('sort_order') or float('inf') for c in chunks), default=float('inf')) - - # ── Build a unified render queue ── - # Each entry: (sort_key, render_type, data) - # render_type: 'outline' | 'orphan_leaf' | 'orphan_child' - render_queue: list[tuple[float, str, dict | str]] = [] - - outline_paths: set[str] = set() - # Position counter for outline-only items (no leaf content) to preserve - # their relative ordering among themselves. - outline_position = 0.0 - - for item in node.outline_items: - path = item.get('path', '') - # Skip items belonging to a drilled-into child's subtree - if any(path.startswith(cp + ' / ') for cp in child_prefixes): - continue - outline_paths.add(path) - - # Determine sort_key: use chunk sort_order if content exists, - # else use a synthetic position to maintain outline ordering. - if path in node.leaf_content or path in node.children: - sort_key = _min_sort(path) if path in node.leaf_content else outline_position - else: - sort_key = outline_position - outline_position = max(outline_position, sort_key) + 0.001 - - render_queue.append((sort_key, 'outline', item)) - - # Add orphan leaf_content paths (not covered by outline_items) - for path in node.leaf_content: - if path not in outline_paths: - render_queue.append((_min_sort(path), 'orphan_leaf', path)) - - # Add orphan children (not covered by outline_items) - for path in node.children: - if path not in outline_paths: - render_queue.append((float('inf'), 'orphan_child', path)) - - # Sort by sort_key (stable sort preserves insertion order for ties) - render_queue.sort(key=lambda x: x[0]) - - from typing import cast - - # ── Render the unified queue ── - for _sort_key, rtype, data in render_queue: - if rtype == 'outline': - item = cast(dict, data) - path = item.get('path', '') - title = item.get('title', '') - is_leaf = item.get('is_leaf', False) - level = item.get('level', 1) - leaf_tag = ' [Leaf]' if is_leaf else '' - - level_tag = f'[L{level}] ' if level else '' - if level <= 1: - parts.append(f'{indent}▸ {level_tag}{title}{leaf_tag}') - else: - parts.append(f'{indent}└ {level_tag}{title}{leaf_tag}') - - sub_indent = indent + ' ' - - # Case 1: drilled-into child → render child tree - if path in node.children: - child = node.children[path] - if path in node.leaf_content: - _render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) - child_text = render_unified_doc_tree(child, doc_name, depth + 1, asset_lookup=asset_lookup) - if child_text.strip(): - parts.append(child_text) - - # Case 2: hydrated leaf → show chunk content - elif path in node.leaf_content: - _render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) - - # Case 3: unselected → title only (already rendered above) - - elif rtype == 'orphan_leaf': - path = cast(str, data) - title = path.rsplit(' / ', 1)[-1] if ' / ' in path else path - parts.append(f'{indent}▸ [Leaf] {title}') - sub_indent = indent + ' ' - _render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) - - elif rtype == 'orphan_child': - path = cast(str, data) - title = path.rsplit(' / ', 1)[-1] if ' / ' in path else path - parts.append(f'{indent}▸ {title} [DrillDown]') - child_text = render_unified_doc_tree(node.children[path], doc_name, depth + 1, asset_lookup=asset_lookup) - if child_text.strip(): - parts.append(child_text) - - return '\n'.join(parts) - diff --git a/packages/shared-python/shared/services/retrieval/agentic/core/__init__.py b/packages/shared-python/shared/services/retrieval/agentic/core/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/core/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/agentic/budget.py b/packages/shared-python/shared/services/retrieval/agentic/core/budget.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/agentic/budget.py rename to packages/shared-python/shared/services/retrieval/agentic/core/budget.py diff --git a/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py b/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py new file mode 100644 index 000000000..7f71cf076 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py @@ -0,0 +1,146 @@ +"""Runtime setup helpers for agentic retrieval.""" +from __future__ import annotations + +import json +import os +from typing import Any + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk +from shared.services.retrieval.agentic.core.budget import BudgetExceeded, BudgetPoolName +from shared.services.retrieval.agentic.core.types import AgentRunConfig, AgentState +from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage +from shared.utils.token_estimate import estimate_tokens + + +def build_config_from_env() -> AgentRunConfig: + return AgentRunConfig( + max_revisions=int(os.environ.get("RETRIEVAL_AGENTIC_MAX_REVISIONS", "2")), + max_nav_depth=int(os.environ.get("RETRIEVAL_AGENTIC_MAX_NAV_DEPTH", "3")), + latency_budget_ms=int(os.environ.get("RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS", "12000")), + token_budget_total=int(os.environ.get("RETRIEVAL_AGENTIC_TOKEN_BUDGET_TOTAL", "40000")), + planning_ratio=float(os.environ.get("RETRIEVAL_AGENTIC_PLANNING_RATIO", "0.5")), + bootstrap_budget=int(os.environ.get("RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET", "2000")), + per_doc_min_share=int(os.environ.get("RETRIEVAL_AGENTIC_PER_DOC_MIN_SHARE", "1500")), + inventory_aware=os.environ.get("RETRIEVAL_AGENTIC_INVENTORY_AWARE", "true") == "true", + ) + + +async def load_budget_inventory( + db: AsyncSession, + *, + user_id: str, + namespace: str, + exclude_document_ids: list[str], +) -> tuple[int, int, dict[str, int]]: + stmt = ( + select(Document.document_id, func.count(DocumentChunk.id)) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == "active") + .group_by(Document.document_id) + ) + if exclude_document_ids: + stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) + + result = await db.execute(stmt) + doc_chunks = {str(doc_id): int(count or 0) for doc_id, count in result.all()} + return sum(doc_chunks.values()), len(doc_chunks), doc_chunks + + +class AgentLlmBudget: + def __init__(self, state: AgentState) -> None: + self._state = state + + async def call( + self, + llm_fn: LLMFn, + prompt: Any, + *, + pool: BudgetPoolName, + doc_id: str | None = None, + priority: str = "normal", + ) -> str: + ledger = self._state.ledger + if ledger is None: + return await llm_fn(prompt) + + prompt_text = _stringify_llm_input(prompt) + est = estimate_tokens(prompt_text) + reserved = await ledger.try_reserve( + pool, + est, + doc_id=doc_id, + priority="low" if priority == "low" else "normal", + ) + if not reserved: + raise BudgetExceeded(f"{pool} budget exhausted") + + try: + response = await llm_fn(prompt) + except Exception: + await ledger.refund(pool, est=est, doc_id=doc_id) + raise + + usage = current_llm_usage.get() or {} + actual = int(usage.get("prompt_tokens") or est) + await ledger.commit(pool, actual=actual, est=est, doc_id=doc_id) + return response + + def for_pool(self, llm_fn: LLMFn, *, pool: BudgetPoolName) -> LLMFn: + async def _call(prompt: Any) -> str: + return await self.call(llm_fn, prompt, pool=pool) + + return _call + + def for_document( + self, + llm_fn: LLMFn, + *, + doc_id: str, + depth: int, + ) -> LLMFn: + async def _call(prompt: Any) -> str: + return await self.call( + llm_fn, + prompt, + pool="planning", + doc_id=doc_id, + priority="low" if depth >= 2 else "normal", + ) + + return _call + + def for_discovery( + self, + llm_fn: LLMFn, + *, + doc_id: str, + low_priority: bool, + ) -> LLMFn: + async def _call(prompt: Any) -> str: + return await self.call( + llm_fn, + prompt, + pool="planning", + doc_id=doc_id, + priority="low" if low_priority else "normal", + ) + + return _call + + +def _stringify_llm_input(prompt: Any) -> str: + if isinstance(prompt, str): + return prompt + try: + return json.dumps(prompt, ensure_ascii=False, default=str) + except Exception: + return str(prompt) diff --git a/packages/shared-python/shared/services/retrieval/agentic/trace.py b/packages/shared-python/shared/services/retrieval/agentic/core/trace.py similarity index 93% rename from packages/shared-python/shared/services/retrieval/agentic/trace.py rename to packages/shared-python/shared/services/retrieval/agentic/core/trace.py index 87df0c94c..bc6ccaa02 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/trace.py +++ b/packages/shared-python/shared/services/retrieval/agentic/core/trace.py @@ -15,7 +15,7 @@ from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.agentic.types import AgentRunConfig, ToolResult +from shared.services.retrieval.agentic.core.types import AgentRunConfig, ToolResult def _now_utc() -> datetime: @@ -105,8 +105,11 @@ async def create_run(self) -> None: logger.debug(f'agentic trace: failed to create run {self._run_id}: {e}') try: await self._db.rollback() - except Exception: - pass + except Exception as rollback_error: + logger.debug( + f'agentic trace: failed to roll back create failure ' + f'{self._run_id}: {rollback_error}' + ) def record_step( self, @@ -213,5 +216,8 @@ async def complete( logger.debug(f'agentic trace: failed to complete run {self._run_id}: {e}') try: await self._db.rollback() - except Exception: - pass + except Exception as rollback_error: + logger.debug( + f'agentic trace: failed to roll back completion failure ' + f'{self._run_id}: {rollback_error}' + ) diff --git a/packages/shared-python/shared/services/retrieval/agentic/types.py b/packages/shared-python/shared/services/retrieval/agentic/core/types.py similarity index 99% rename from packages/shared-python/shared/services/retrieval/agentic/types.py rename to packages/shared-python/shared/services/retrieval/agentic/core/types.py index 92c8b7f2f..79981b9f6 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/types.py +++ b/packages/shared-python/shared/services/retrieval/agentic/core/types.py @@ -10,7 +10,7 @@ from dataclasses import dataclass, field from typing import Any -from shared.services.retrieval.agentic.budget import BudgetLedger +from shared.services.retrieval.agentic.core.budget import BudgetLedger @dataclass diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery/__init__.py b/packages/shared-python/shared/services/retrieval/agentic/discovery/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/discovery/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery/phase.py b/packages/shared-python/shared/services/retrieval/agentic/discovery/phase.py new file mode 100644 index 000000000..5a95061ed --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/discovery/phase.py @@ -0,0 +1,254 @@ +"""Discovery and document selection phase for agentic retrieval.""" +from __future__ import annotations + +from typing import Any + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document +from shared.services.retrieval.agentic import tools +from shared.services.retrieval.agentic.core.budget import BudgetExceeded +from shared.services.retrieval.agentic.core.trace import TraceRecorder +from shared.services.retrieval.agentic.core.types import AgentState, CandidateDoc, ToolResult +from shared.services.retrieval.llm_adapter import LLMFn + + +async def run_initial_discovery( + db: AsyncSession, + *, + state: AgentState, + trace: TraceRecorder, + trace_enabled: bool, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + data_type: int, + signal_paths: list[str] | None, + filter_mode: str, + channels: list[str] | None, + channel_weights: dict[str, float] | None, + internal_recall_k: int | None, + bootstrap_llm_fn: LLMFn | None, +) -> list[dict[str, Any]]: + discovery_kwargs: dict[str, Any] = { + "user_id": user_id, + "namespace": namespace, + "query": query, + "top_k": top_k, + "exclude_document_ids": exclude_document_ids, + "exclude_sections": exclude_sections, + "data_type": data_type, + "signal_paths": signal_paths, + "filter_mode": filter_mode, + "channels": channels, + "channel_weights": channel_weights, + "internal_recall_k": internal_recall_k, + } + + logger.info(" agentic: Phase 1 — discovery + document selection") + discovery_result = await tools.bottom_discovery(db, **discovery_kwargs) + state.step_count += 1 + discovery_rows = ( + discovery_result.payload.get("fused_rows", []) + if discovery_result.status != "error" + else [] + ) + state.discovery_top_doc_ids = ( + discovery_result.payload.get("top_doc_ids", []) + if discovery_result.status != "error" + else [] + ) + + if trace_enabled: + trace.record_step( + "bottom_discovery", + discovery_result, + decision_reason="phase_1_mandatory", + ) + + logger.info( + f" agentic step {state.step_count}: bottom_discovery " + f"status={discovery_result.status} latency={discovery_result.latency_ms}ms" + ) + + if bootstrap_llm_fn is not None: + await _select_documents( + db, + state=state, + trace=trace, + trace_enabled=trace_enabled, + user_id=user_id, + namespace=namespace, + query=query, + exclude_document_ids=exclude_document_ids, + bootstrap_llm_fn=bootstrap_llm_fn, + ) + + return discovery_rows + + +async def register_discovery_documents( + db: AsyncSession, + *, + state: AgentState, + discovery_by_doc: dict[str, list[dict[str, Any]]], +) -> None: + selected_doc_ids = {doc.document_id for doc in state.selected_docs} + for doc_id in discovery_by_doc: + if doc_id in selected_doc_ids or doc_id in state.ever_explored_doc_ids: + continue + doc_stmt = ( + select(Document.document_id, Document.source_file_name, Document.current_job_result_id) + .where(Document.document_id == doc_id) + ) + doc_result = await db.execute(doc_stmt) + row_data = doc_result.first() + if row_data is None: + continue + did, fname, job_result_id = row_data + state.selected_docs.append( + CandidateDoc( + document_id=did, + source_file_name=fname or did, + confidence=0.4, + reason="discovery_auto (not in KG selection)", + source="discovery_auto", + ) + ) + state.doc_id_to_name[did] = fname or did + if job_result_id: + state.doc_job_map[did] = job_result_id + + +async def select_revision_documents( + db: AsyncSession, + *, + state: AgentState, + trace: TraceRecorder, + trace_enabled: bool, + user_id: str, + namespace: str, + query: str, + exclude_document_ids: list[str], + bootstrap_llm_fn: LLMFn, + revision_hint: str, +) -> str | None: + try: + kg_result = await tools.kg_document_select( + db, + user_id=user_id, + namespace=namespace, + query=query, + llm_fn=bootstrap_llm_fn, + exclude_document_ids=list(set(exclude_document_ids)), + revision_hint=revision_hint, + budget_snapshot=state.ledger.snapshot() if state.ledger else None, + ) + except BudgetExceeded: + logger.info(" agentic: bootstrap budget exhausted during revision doc selection") + if trace_enabled: + trace.record_budget_stop("bootstrap_exhausted") + return "bootstrap_budget" + state.step_count += 1 + _append_selected_docs(state, kg_result) + return None + + +async def _select_documents( + db: AsyncSession, + *, + state: AgentState, + trace: TraceRecorder, + trace_enabled: bool, + user_id: str, + namespace: str, + query: str, + exclude_document_ids: list[str], + bootstrap_llm_fn: LLMFn, +) -> None: + try: + kg_result = await tools.kg_document_select( + db, + user_id=user_id, + namespace=namespace, + query=query, + llm_fn=bootstrap_llm_fn, + exclude_document_ids=list(state.ever_explored_doc_ids | set(exclude_document_ids)), + budget_snapshot=state.ledger.snapshot() if state.ledger else None, + ) + except BudgetExceeded: + logger.info(" agentic: bootstrap budget exhausted during document selection") + if trace_enabled: + trace.record_budget_stop("bootstrap_exhausted") + kg_result = ToolResult( + status="no_confident_doc", + payload={"reason": "bootstrap budget exhausted"}, + ) + state.step_count += 1 + + if trace_enabled: + trace.record_step( + "kg_document_select", + kg_result, + decision_reason="phase_1_doc_selection", + ) + + _append_selected_docs(state, kg_result) + if not state.selected_docs and state.discovery_top_doc_ids: + await _append_discovery_hints(db, state=state) + + logger.info( + f" agentic step {state.step_count}: kg_document_select " + f"status={kg_result.status} docs={len(state.selected_docs)} " + f"latency={kg_result.latency_ms}ms" + ) + + +async def _append_discovery_hints(db: AsyncSession, *, state: AgentState) -> None: + hint_ids = [ + doc_id + for doc_id in state.discovery_top_doc_ids + if doc_id not in state.ever_explored_doc_ids + ] + if not hint_ids: + return + doc_stmt = ( + select(Document.document_id, Document.source_file_name, Document.current_job_result_id) + .where(Document.document_id.in_(hint_ids)) + ) + doc_result = await db.execute(doc_stmt) + for doc_id, source_file_name, job_result_id in doc_result.all(): + state.selected_docs.append( + CandidateDoc( + document_id=doc_id, + source_file_name=source_file_name or doc_id, + confidence=0.5, + reason="discovery_hint (KG returned 0)", + source="discovery_hint", + ) + ) + state.doc_id_to_name[doc_id] = source_file_name or doc_id + if job_result_id: + state.doc_job_map[doc_id] = job_result_id + + +def _append_selected_docs(state: AgentState, kg_result: ToolResult) -> None: + if kg_result.status != "selected_docs": + return + for doc_data in kg_result.payload.get("candidate_docs", []): + state.selected_docs.append( + CandidateDoc( + document_id=doc_data.get("document_id", ""), + source_file_name=doc_data.get("source_file_name", ""), + confidence=doc_data.get("confidence", 0.0), + reason=doc_data.get("reason", ""), + source=doc_data.get("source", ""), + ) + ) + state.doc_id_to_name.update(kg_result.payload.get("doc_id_to_name", {})) + state.doc_job_map.update(kg_result.payload.get("doc_job_map", {})) diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery/selection.py b/packages/shared-python/shared/services/retrieval/agentic/discovery/selection.py new file mode 100644 index 000000000..55103b94b --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/discovery/selection.py @@ -0,0 +1,206 @@ +"""Post-navigation discovery selection for agentic retrieval.""" +from __future__ import annotations + +import time +from typing import Any + +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.agentic.core.budget import BudgetExceeded +from shared.services.retrieval.agentic.prompts import ( + DISCOVERY_SELECT_PROMPT, + format_budget_block, + parse_action_response, +) +from shared.services.retrieval.agentic.navigation.selection_hydration import ( + hydrate_path_selections_into_node, +) +from shared.services.retrieval.agentic.core.types import DocTreeNode +from shared.services.retrieval.search.lexical_text import normalize_section_path +from shared.services.retrieval.llm_adapter import LLMFn + + +_MAX_DISCOVERY_PER_DOC = 3 + + +async def discovery_select_step( + db: AsyncSession, + *, + document_id: str, + query: str, + llm_fn: LLMFn, + user_id: str, + namespace: str, + doc_name: str = "", + discovery_hints: list[dict[str, Any]], + exclude_paths: set[str] | None = None, + revision_hint: str | None = None, + budget_snapshot: dict | None = None, +) -> DocTreeNode: + """Select and hydrate discovery-found sections after BFS navigation.""" + node = DocTreeNode(scope_path=None) + if not discovery_hints: + return node + + hints = discovery_hints[:_MAX_DISCOVERY_PER_DOC] + + t0 = time.monotonic() + try: + hint_lines, hint_by_path, root_path_selections = _project_discovery_hints( + hints, + exclude_paths=exclude_paths, + ) + if not hint_lines and not root_path_selections: + return node + + selections: list[dict[str, Any]] = [] + if hint_lines: + prompt = _build_discovery_selection_prompt( + document_id=document_id, + doc_name=doc_name, + query=query, + hint_lines=hint_lines, + revision_hint=revision_hint, + budget_snapshot=budget_snapshot, + ) + response = await llm_fn(prompt) + parsed = parse_action_response(response) + selections = parsed.get("selections", []) + + logger.info( + f' discovery_select_step doc="{doc_name}": ' + f"hints={len(hints)} selections={len(selections)} " + f"root_selections={len(root_path_selections)}" + ) + + path_selections = _build_discovery_path_selections( + selections=selections, + hint_by_path=hint_by_path, + root_path_selections=root_path_selections, + node=node, + ) + await hydrate_path_selections_into_node( + db, + node=node, + path_selections=path_selections, + user_id=user_id, + namespace=namespace, + document_id=document_id, + ) + + latency = int((time.monotonic() - t0) * 1000) + logger.info( + f" discovery_select_step done: hydrated={len(node.leaf_content)} " + f"latency={latency}ms" + ) + return node + + except BudgetExceeded: + raise + except Exception as exc: + logger.error(f" discovery_select_step failed for doc={document_id}: {exc}") + return node + + +def _project_discovery_hints( + hints: list[dict[str, Any]], + *, + exclude_paths: set[str] | None, +) -> tuple[list[str], dict[str, dict], list[dict[str, Any]]]: + exclude_set = { + normalize_section_path(path) + for path in (exclude_paths or set()) + if path + } + hint_lines: list[str] = [] + hint_by_path: dict[str, dict] = {} + root_path_selections: list[dict[str, Any]] = [] + for hint in hints: + section_path = normalize_section_path(hint.get("section_path", "")) + if not section_path: + continue + if section_path in exclude_set: + continue + if section_path in hint_by_path: + continue + + hint_by_path[section_path] = hint + if section_path == "Root": + root_path_selections.append({ + "path": section_path, + "confidence": float( + hint.get("discovery_score") or hint.get("score") or 0.7 + ), + "hydrate_mode": "self_only", + }) + continue + + summary = hint.get("summary", "") or "" + hint_lines.append(f'▸ path="{section_path}"') + if summary: + hint_lines.append(f" {summary[:300]}") + + return hint_lines, hint_by_path, root_path_selections + + +def _build_discovery_selection_prompt( + *, + document_id: str, + doc_name: str, + query: str, + hint_lines: list[str], + revision_hint: str | None, + budget_snapshot: dict | None, +) -> str: + revision_context = "" + if revision_hint: + revision_context = ( + "\nIMPORTANT: This is a REVISION round. " + "The previous search attempt failed because:\n" + f'"{revision_hint}"\n' + "Adjust your selection accordingly. " + "If no candidate is relevant, return an EMPTY list [].\n" + ) + + return DISCOVERY_SELECT_PROMPT.format( + doc_name=doc_name or document_id, + budget_block=format_budget_block(budget_snapshot), + items="\n".join(hint_lines), + query=query, + revision_context=revision_context, + ) + + +def _build_discovery_path_selections( + *, + selections: list[dict[str, Any]], + hint_by_path: dict[str, dict], + root_path_selections: list[dict[str, Any]], + node: DocTreeNode, +) -> list[dict[str, Any]]: + valid_selections = [ + selection for selection in selections if selection["path"] in hint_by_path + ] + path_selections = list(root_path_selections) + for selection in valid_selections: + path = selection["path"] + confidence = selection.get("confidence", 0.7) + node.confidence[path] = confidence + path_selections.append({"path": path, "confidence": confidence}) + + if not path_selections and hint_by_path: + fallback_path, fallback_hint = next(iter(hint_by_path.items())) + fallback_confidence = float( + fallback_hint.get("discovery_score") + or fallback_hint.get("score") + or 0.5 + ) + node.confidence[fallback_path] = fallback_confidence + path_selections.append({ + "path": fallback_path, + "confidence": fallback_confidence, + "hydrate_mode": "self_only", + }) + + return path_selections diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery/tools.py b/packages/shared-python/shared/services/retrieval/agentic/discovery/tools.py new file mode 100644 index 000000000..e574d1b94 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/discovery/tools.py @@ -0,0 +1,298 @@ +"""Agentic retrieval discovery tools. + +This Module owns phase-1 retrieval: lexical bottom discovery and document +selection from the document-level knowledge map. The public tool adapter stays +in ``tools.py`` so orchestrator call sites keep a stable interface. +""" +from __future__ import annotations + +import time +from typing import Any + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document +from shared.services.retrieval.agentic.core.budget import BudgetExceeded +from shared.services.retrieval.agentic.navigation.knowledge_map import build_knowledge_map_overview +from shared.services.retrieval.agentic.prompts import ( + FILE_SELECT_PROMPT, + format_budget_block, + parse_json_array, +) +from shared.services.retrieval.agentic.core.types import ToolResult +from shared.services.retrieval.search.channels import content_channel, path_channel, term_channel +from shared.services.retrieval.llm_adapter import LLMFn +from shared.services.retrieval.search.scoring import ( + merge_channels_rrf, + merge_same_section_rows, + normalize_row_scores, +) +from shared.services.retrieval.settings import ( + CHANNEL_WEIGHT_CONTENT, + CHANNEL_WEIGHT_PATH, + CHANNEL_WEIGHT_TERM, + INTERNAL_RECALL_K_MULTIPLIER, + resolve_allowed_chunk_types, +) + + +async def bottom_discovery( + db: AsyncSession, + *, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + data_type: int = 1, + signal_paths: list[str] | None = None, + filter_mode: str = "delete", + channels: list[str] | None = None, + channel_weights: dict[str, float] | None = None, + internal_recall_k: int | None = None, + **_kwargs: Any, +) -> ToolResult: + """Run 3-channel BM25 discovery plus RRF fusion.""" + t0 = time.monotonic() + try: + allowed_chunk_types = resolve_allowed_chunk_types(data_type) + effective_recall_k = ( + internal_recall_k + if internal_recall_k is not None + else top_k * INTERNAL_RECALL_K_MULTIPLIER + ) + active_channels = set(channels) if channels else {"path", "content", "term"} + + path_rows: list[dict[str, Any]] = [] + content_rows: list[dict[str, Any]] = [] + term_rows: list[dict[str, Any]] = [] + + if "path" in active_channels: + path_rows = await path_channel( + db, + user_id=user_id, + namespace=namespace, + query=query, + top_k=effective_recall_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + allowed_chunk_types=allowed_chunk_types, + signal_paths=signal_paths, + filter_mode=filter_mode, + ) + + if "content" in active_channels: + content_rows = await content_channel( + db, + user_id=user_id, + namespace=namespace, + query=query, + top_k=effective_recall_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + allowed_chunk_types=allowed_chunk_types, + signal_paths=signal_paths, + filter_mode=filter_mode, + ) + + if "term" in active_channels: + term_rows = await term_channel( + db, + user_id=user_id, + namespace=namespace, + query=query, + top_k=effective_recall_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + allowed_chunk_types=allowed_chunk_types, + signal_paths=signal_paths, + filter_mode=filter_mode, + ) + + default_weights = { + "path": CHANNEL_WEIGHT_PATH, + "content": CHANNEL_WEIGHT_CONTENT, + "term": CHANNEL_WEIGHT_TERM, + } + effective_weights = {**default_weights, **(channel_weights or {})} + + channel_lists: list[list[dict[str, Any]]] = [] + weight_list: list[float] = [] + if path_rows: + channel_lists.append(path_rows) + weight_list.append(effective_weights.get("path", CHANNEL_WEIGHT_PATH)) + if content_rows: + channel_lists.append(content_rows) + weight_list.append(effective_weights.get("content", CHANNEL_WEIGHT_CONTENT)) + if term_rows: + channel_lists.append(term_rows) + weight_list.append(effective_weights.get("term", CHANNEL_WEIGHT_TERM)) + + fused_rows = ( + merge_channels_rrf(channel_lists, weight_list, effective_recall_k) + if channel_lists + else [] + ) + fused_rows = merge_same_section_rows(fused_rows) + + if fused_rows: + normalize_row_scores( + fused_rows, + source_field="score", + target_field="discovery_score", + default=0.5, + ) + + doc_id_counts: dict[str, int] = {} + for row in fused_rows: + did = row.get("document_id", "") + if did: + doc_id_counts[did] = doc_id_counts.get(did, 0) + 1 + top_doc_ids = sorted( + doc_id_counts, + key=lambda document_id: doc_id_counts[document_id], + reverse=True, + )[:5] + + latency = int((time.monotonic() - t0) * 1000) + logger.info( + f" agentic.bottom_discovery: {len(fused_rows)} fused rows, " + f"top_doc_ids={top_doc_ids}, {latency}ms" + ) + return ToolResult( + status="discovery_done", + payload={ + "fused_rows": fused_rows, + "top_doc_ids": top_doc_ids, + "channel_counts": { + "path": len(path_rows), + "content": len(content_rows), + "term": len(term_rows), + }, + }, + latency_ms=latency, + ) + except Exception as exc: + latency = int((time.monotonic() - t0) * 1000) + logger.error(f" agentic.bottom_discovery failed: {exc}") + return ToolResult(status="error", error=str(exc), latency_ms=latency) + + +async def kg_document_select( + db: AsyncSession, + *, + user_id: str, + namespace: str, + query: str, + llm_fn: LLMFn | None, + exclude_document_ids: list[str], + revision_hint: str | None = None, + **_kwargs: Any, +) -> ToolResult: + """Select candidate documents from document-level KG.""" + t0 = time.monotonic() + try: + overview_text, doc_id_to_name = await build_knowledge_map_overview( + db, + user_id=user_id, + namespace=namespace, + ) + if overview_text == "(empty)": + latency = int((time.monotonic() - t0) * 1000) + return ToolResult( + status="no_confident_doc", + payload={"reason": "no active documents in namespace"}, + latency_ms=latency, + ) + + if llm_fn is None: + latency = int((time.monotonic() - t0) * 1000) + return ToolResult( + status="no_confident_doc", + payload={"reason": "LLM not available"}, + latency_ms=latency, + ) + + revision_context = "" + if revision_hint: + revision_context = ( + "\nIMPORTANT: This is a REVISION round. " + "The previous search attempt failed because:\n" + f'"{revision_hint}"\n' + "Adjust your document selection accordingly. " + "If no document can address this, return an EMPTY array [].\n" + ) + + file_prompt = FILE_SELECT_PROMPT.format( + overview=overview_text, + query=query, + revision_context=revision_context, + budget_block=format_budget_block(_kwargs.get("budget_snapshot")), + ) + file_response = await llm_fn(file_prompt) + selected_ids = parse_json_array(file_response) + + exclude_set = set(exclude_document_ids) + valid_ids = [ + document_id + for document_id in selected_ids + if document_id in doc_id_to_name and document_id not in exclude_set + ] + + if not valid_ids: + latency = int((time.monotonic() - t0) * 1000) + logger.info( + f" agentic.kg_document_select: LLM returned no valid docs, {latency}ms" + ) + return ToolResult( + status="no_confident_doc", + payload={ + "reason": "LLM returned no valid document IDs", + "raw_ids": selected_ids, + }, + latency_ms=latency, + ) + + doc_job_map: dict[str, str] = {} + doc_stmt = ( + select(Document.document_id, Document.current_job_result_id) + .where(Document.document_id.in_(valid_ids)) + ) + doc_result = await db.execute(doc_stmt) + for document_id, job_result_id in doc_result.all(): + if job_result_id: + doc_job_map[document_id] = job_result_id + + candidate_docs = [ + { + "document_id": document_id, + "source_file_name": doc_id_to_name.get(document_id, ""), + "confidence": 1.0, + "reason": "LLM selected from KG overview", + "source": "kg_llm_select", + } + for document_id in valid_ids + ] + + latency = int((time.monotonic() - t0) * 1000) + logger.info( + f" agentic.kg_document_select: {len(candidate_docs)} docs selected, {latency}ms" + ) + return ToolResult( + status="selected_docs", + payload={ + "candidate_docs": candidate_docs, + "doc_id_to_name": doc_id_to_name, + "doc_job_map": doc_job_map, + }, + latency_ms=latency, + ) + except BudgetExceeded: + raise + except Exception as exc: + latency = int((time.monotonic() - t0) * 1000) + logger.error(f" agentic.kg_document_select failed: {exc}") + return ToolResult(status="error", error=str(exc), latency_ms=latency) diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence/__init__.py b/packages/shared-python/shared/services/retrieval/agentic/evidence/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py b/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py new file mode 100644 index 000000000..6da147f42 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py @@ -0,0 +1,325 @@ +from __future__ import annotations + +from typing import Any + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import RetrievalHitStat +from shared.services.retrieval.agentic.core.budget import BudgetLedger +from shared.services.retrieval.agentic.core.types import DocTreeNode +from shared.services.retrieval.hydration.assets import build_retrieval_asset_url_map +from shared.services.retrieval.stats.service import compute_importance_score +from shared.utils.token_estimate import estimate_tokens + + +def with_context_prompt_projection( + snapshot: dict[str, object], + *, + prompt_tokens: int, +) -> dict[str, object]: + projected: dict[str, object] = dict(snapshot) + context_raw = projected.get("context") or {} + if not isinstance(context_raw, dict): + return projected + + context = dict(context_raw) + used = int(context.get("used", 0) or 0) + reserved = int(context.get("reserved", 0) or 0) + capacity = int(context.get("capacity", 0) or 0) + projected_used = min(capacity, used + max(int(prompt_tokens), 0)) + projected_remaining = max(capacity - projected_used - reserved, 0) + context.update( + { + "used_projected_before_answer": projected_used, + "answer_prompt_estimate": max(int(prompt_tokens), 0), + "remaining": projected_remaining, + "used_pct": 100 + if capacity <= 0 + else min(100, int(round((projected_used + reserved) * 100 / capacity))), + } + ) + if projected_remaining <= 0: + context["status"] = "EXHAUSTED" + elif context["used_pct"] >= 80: + context["status"] = "CRITICAL" + elif context["used_pct"] >= 50: + context["status"] = "TIGHT" + else: + context["status"] = "HEALTHY" + projected["context"] = context + return projected + + +def collect_media_chunks(node: DocTreeNode) -> list[dict[str, Any]]: + media: list[dict[str, Any]] = [] + for chunks in node.leaf_content.values(): + for chunk in chunks: + chunk_type = ( + chunk.get("chunk_type") or chunk.get("type") or "" + ).strip().lower() + if chunk_type in ("image", "table"): + media.append(chunk) + for child in node.children.values(): + media.extend(collect_media_chunks(child)) + return media + + +def collect_media_chunks_all( + doc_trees: dict[str, DocTreeNode], +) -> list[dict[str, Any]]: + media: list[dict[str, Any]] = [] + for tree in doc_trees.values(): + media.extend(collect_media_chunks(tree)) + return media + + +async def build_asset_url_map( + media_chunks: list[dict[str, Any]], +) -> dict[str, str]: + return await build_retrieval_asset_url_map( + media_chunks, + log_context="agentic evidence", + ) + + +def _collect_all_leaf_paths(node: DocTreeNode) -> set[str]: + paths = set(node.leaf_content.keys()) + for child in node.children.values(): + paths.update(_collect_all_leaf_paths(child)) + return paths + + +def _collect_visible_paths(node: DocTreeNode) -> set[str]: + paths = {item["path"] for item in node.outline_items if item.get("path")} + for child in node.children.values(): + paths.update(_collect_visible_paths(child)) + return paths + + +def _find_closest_ancestor(path: str, target_paths: set[str]) -> str | None: + parts = path.split(" / ") + for i in range(len(parts) - 1, 0, -1): + ancestor = " / ".join(parts[:i]) + if ancestor in target_paths: + return ancestor + return None + + +def reconcile_deferred_assets( + tree: DocTreeNode, + pending_assets: list[dict], +) -> None: + final_paths = _collect_all_leaf_paths(tree) + visible_paths = _collect_visible_paths(tree) + all_target_paths = final_paths | visible_paths + + if not all_target_paths: + return + + existing_ids = { + str(row.get("chunk_id") or "") + for row in tree.flatten_chunk_rows() + if row.get("chunk_id") + } + + placed = 0 + ancestor_placed = 0 + for asset in pending_assets: + chunk_id = str(asset.get("chunk_id") or "") + if chunk_id and chunk_id in existing_ids: + continue + + owner_path = asset.get("owner_section_path") or asset.get("section_path") + if not owner_path: + continue + + target_path = owner_path if owner_path in all_target_paths else None + if target_path is None: + target_path = _find_closest_ancestor(owner_path, all_target_paths) + if target_path: + ancestor_placed += 1 + + if target_path is None: + continue + + tree.add_leaf_chunks(target_path, [asset]) + if chunk_id: + existing_ids.add(chunk_id) + placed += 1 + + if placed: + tree.reparent_leaf_content() + logger.info( + f" deferred asset reconcile: {placed}/{len(pending_assets)} " + f"assets placed into {len(final_paths)} leaf + {len(visible_paths)} visible paths " + f"(ancestor_fallback={ancestor_placed})" + ) + + +async def render_evidence( + db: AsyncSession, + doc_trees: dict[str, DocTreeNode], + doc_id_to_name: dict[str, str], +) -> str: + del db + + from shared.services.retrieval.agentic.evidence.renderer import render_unified_doc_tree + + asset_url_map = await build_asset_url_map(collect_media_chunks_all(doc_trees)) + + evidence_parts: list[str] = [] + for doc_id, doc_tree in doc_trees.items(): + if doc_tree.has_content(): + doc_name = doc_id_to_name.get(doc_id, doc_id) + rendered = render_unified_doc_tree( + doc_tree, + doc_name, + asset_lookup=asset_url_map, + ) + if rendered.strip(): + evidence_parts.append(rendered) + + return "\n\n".join(evidence_parts) if evidence_parts else "(no evidence collected)" + + +def _iter_leaf_content(node: DocTreeNode): + for path, chunks in node.leaf_content.items(): + yield path, chunks + for child in node.children.values(): + yield from _iter_leaf_content(child) + + +def _collect_confidences(node: DocTreeNode) -> dict[str, float]: + values = dict(node.confidence) + for child in node.children.values(): + for path, score in _collect_confidences(child).items(): + values[path] = max(values.get(path, 0.0), score) + return values + + +def _pop_leaf_path(node: DocTreeNode, path: str) -> bool: + if path in node.leaf_content: + node.leaf_content.pop(path) + return True + for child in node.children.values(): + if _pop_leaf_path(child, path): + return True + return False + + +def _estimate_chunks_tokens(chunks: list[dict[str, Any]]) -> int: + text = "\n".join(str(chunk.get("content") or "") for chunk in chunks) + return estimate_tokens(text) + + +async def _fetch_importance_norm_scores( + db: AsyncSession, + *, + user_id: str, + namespace: str, + chunk_ids: list[str], +) -> dict[str, float]: + if not chunk_ids: + return {} + stmt = ( + select( + RetrievalHitStat.chunk_id, + RetrievalHitStat.hit_count, + RetrievalHitStat.last_hit_at, + RetrievalHitStat.created_at, + ) + .where(RetrievalHitStat.user_id == user_id) + .where(RetrievalHitStat.namespace == namespace) + .where(RetrievalHitStat.hit_kind == "chunk") + .where(RetrievalHitStat.chunk_id.in_(chunk_ids)) + ) + result = await db.execute(stmt) + scores: dict[str, float] = {} + for chunk_id, hit_count, last_hit_at, created_at in result.all(): + if chunk_id and last_hit_at and created_at: + scores[str(chunk_id)] = compute_importance_score( + hit_count, + last_hit_at, + created_at, + ) + return scores + + +async def trim_evidence_to_budget( + db: AsyncSession, + *, + doc_trees: dict[str, DocTreeNode], + doc_id_to_name: dict[str, str], + context_remaining: int, + user_id: str, + namespace: str, + ledger: BudgetLedger | None, + safety_margin: float = 0.9, +) -> str: + full_text = await render_evidence(db, doc_trees, doc_id_to_name) + target = int(max(context_remaining, 0) * safety_margin) + if estimate_tokens(full_text) <= target: + return full_text + + candidates: list[tuple[str, str, tuple[float, float, float], int]] = [] + for doc_id, tree in doc_trees.items(): + confidence = _collect_confidences(tree) + for path, chunks in _iter_leaf_content(tree): + chunk_ids = [ + str(chunk.get("chunk_id")) + for chunk in chunks + if chunk.get("chunk_id") + ] + importance = 0.0 + importance_scores = await _fetch_importance_norm_scores( + db, + user_id=user_id, + namespace=namespace, + chunk_ids=chunk_ids, + ) + if importance_scores: + importance = max(importance_scores.values()) + discovery_score = ( + float(chunks[0].get("discovery_score", 0.0) or 0.0) + if chunks + else 0.0 + ) + score = ( + float(confidence.get(path, 0.0) or 0.0), + discovery_score, + importance, + ) + candidates.append((doc_id, path, score, _estimate_chunks_tokens(chunks))) + + current_estimate = estimate_tokens(full_text) + removed: list[dict[str, Any]] = [] + for doc_id, path, score, token_estimate in sorted( + candidates, + key=lambda item: (item[2], -item[3]), + ): + if current_estimate <= target: + break + if _pop_leaf_path(doc_trees[doc_id], path): + confidence_score, discovery_score, importance_score = score + removed.append( + { + "document_id": doc_id, + "document_name": doc_id_to_name.get(doc_id, doc_id), + "path": path, + "confidence_score": round(confidence_score, 4), + "discovery_score": round(discovery_score, 4), + "importance_score": round(importance_score, 4), + "token_estimate": token_estimate, + } + ) + current_estimate = max(current_estimate - token_estimate, 0) + + if ledger is not None: + ledger.trimmed_paths.extend(removed) + logger.info( + f" agentic.trim_evidence: removed={len(removed)} " + f"est_tokens={current_estimate} target={target}" + ) + return await render_evidence(db, doc_trees, doc_id_to_name) diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py b/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py new file mode 100644 index 000000000..016e245d4 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py @@ -0,0 +1,182 @@ +"""Render agentic document trees into evidence text.""" +from __future__ import annotations + +from typing import Any, cast + +from shared.services.retrieval.agentic.core.types import DocTreeNode + + +def render_unified_doc_tree( + node: DocTreeNode, + doc_name: str, + depth: int = 0, + asset_lookup: dict[str, str] | None = None, +) -> str: + """Render a DocTreeNode as one coherent hierarchy.""" + parts: list[str] = [] + indent = " " * depth + + if depth == 0: + parts.append(f"【文档】{doc_name}\n") + + child_prefixes = set(node.children.keys()) + + def min_sort(path: str) -> float: + chunks = node.leaf_content.get(path, []) + return min((chunk.get("sort_order") or float("inf") for chunk in chunks), default=float("inf")) + + render_queue: list[tuple[float, str, dict | str]] = [] + outline_paths: set[str] = set() + outline_position = 0.0 + + for item in node.outline_items: + path = item.get("path", "") + if any(path.startswith(child_prefix + " / ") for child_prefix in child_prefixes): + continue + outline_paths.add(path) + + if path in node.leaf_content or path in node.children: + sort_key = min_sort(path) if path in node.leaf_content else outline_position + else: + sort_key = outline_position + outline_position = max(outline_position, sort_key) + 0.001 + + render_queue.append((sort_key, "outline", item)) + + for path in node.leaf_content: + if path not in outline_paths: + render_queue.append((min_sort(path), "orphan_leaf", path)) + + for path in node.children: + if path not in outline_paths: + render_queue.append((float("inf"), "orphan_child", path)) + + render_queue.sort(key=lambda item: item[0]) + + for _sort_key, render_type, data in render_queue: + if render_type == "outline": + item = cast(dict, data) + path = item.get("path", "") + title = item.get("title", "") + is_leaf = item.get("is_leaf", False) + level = item.get("level", 1) + leaf_tag = " [Leaf]" if is_leaf else "" + + level_tag = f"[L{level}] " if level else "" + if level <= 1: + parts.append(f"{indent}▸ {level_tag}{title}{leaf_tag}") + else: + parts.append(f"{indent}└ {level_tag}{title}{leaf_tag}") + + sub_indent = indent + " " + if path in node.children: + child = node.children[path] + if path in node.leaf_content: + render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) + child_text = render_unified_doc_tree(child, doc_name, depth + 1, asset_lookup=asset_lookup) + if child_text.strip(): + parts.append(child_text) + elif path in node.leaf_content: + render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) + + elif render_type == "orphan_leaf": + path = cast(str, data) + title = path.rsplit(" / ", 1)[-1] if " / " in path else path + parts.append(f"{indent}▸ [Leaf] {title}") + render_leaf_chunks(parts, node.leaf_content[path], indent + " ", asset_lookup=asset_lookup) + + elif render_type == "orphan_child": + path = cast(str, data) + title = path.rsplit(" / ", 1)[-1] if " / " in path else path + parts.append(f"{indent}▸ {title} [DrillDown]") + child_text = render_unified_doc_tree(node.children[path], doc_name, depth + 1, asset_lookup=asset_lookup) + if child_text.strip(): + parts.append(child_text) + + return "\n".join(parts) + + +def render_leaf_chunks( + parts: list[str], + chunks: list[dict[str, Any]], + indent: str, + asset_lookup: dict[str, str] | None = None, +) -> None: + chunk_by_id = { + chunk.get("chunk_id", ""): chunk + for chunk in chunks + if chunk.get("chunk_id") + } + rendered_ids: set[str] = set() + + for chunk in chunks: + chunk_id = chunk.get("chunk_id", "") + if chunk_id and chunk_id in rendered_ids: + continue + + chunk_type = (chunk.get("chunk_type") or chunk.get("type") or "text").strip().lower() + if chunk_type in ("image", "table"): + continue + + if chunk_id: + rendered_ids.add(chunk_id) + + content = str(chunk.get("content", "")).strip() + for connection in (chunk.get("chunk_metadata") or {}).get("connect_to") or []: + target = chunk_by_id.get(connection.get("target", "")) + if not target: + continue + target_id = target.get("chunk_id", "") + target_type = (target.get("chunk_type") or target.get("type") or "").strip().lower() + ref_str = connection.get("ref", "") + if not ref_str or ref_str not in content: + continue + + if target_id: + rendered_ids.add(target_id) + + if target_type == "table": + table_html = str(target.get("content", "")).strip() + content = content.replace(ref_str, f"\n[表格内容]\n{table_html}\n") + elif target_type == "image": + file_path = target.get("file_path") or "" + image_description = str(target.get("content", "")).strip() + if ref_str in image_description: + image_description = image_description.replace(ref_str, "").strip() + asset_url = (asset_lookup or {}).get(target_id, "") if target_id else "" + display_ref = asset_url or file_path + if display_ref: + content = content.replace(ref_str, f"\n[图片: {display_ref}]\n{image_description}\n") + elif image_description: + content = content.replace(ref_str, f"\n[图片描述]\n{image_description}\n") + + for line in content.split("\n"): + if line.strip(): + parts.append(f"{indent}┈ {line}") + + for chunk in chunks: + chunk_id = chunk.get("chunk_id", "") + if chunk_id and chunk_id in rendered_ids: + continue + if chunk_id: + rendered_ids.add(chunk_id) + + chunk_type = (chunk.get("chunk_type") or chunk.get("type") or "").strip().lower() + if chunk_type == "image": + file_path = chunk.get("file_path") or "" + image_description = str(chunk.get("content", "")).strip() + asset_url = (asset_lookup or {}).get(chunk_id, "") if chunk_id else "" + display_ref = asset_url or file_path + if display_ref: + parts.append(f"{indent}┈ [图片: {display_ref}]") + if image_description: + for line in image_description.split("\n"): + if line.strip(): + parts.append(f"{indent}┈ {line}") + elif chunk_type == "table": + table_html = str(chunk.get("content", "")).strip() + parts.append(f"{indent}┈ [表格内容]") + if table_html: + for line in table_html.split("\n"): + if line.strip(): + parts.append(f"{indent}┈ {line}") diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/__init__.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/assets.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/assets.py new file mode 100644 index 000000000..48403e3ca --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/assets.py @@ -0,0 +1,342 @@ +from __future__ import annotations + +import time +from typing import Any + +from loguru import logger +from sqlalchemy import func as sa_func +from sqlalchemy import or_, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.job_result import JobResult + + +def build_connected_owner_map(text_chunks: list[dict[str, Any]]) -> dict[str, str]: + owner_map: dict[str, str] = {} + for chunk in text_chunks: + if (chunk.get("chunk_type") or "text") != "text": + continue + section_path = chunk.get("section_path") or "" + if not section_path: + continue + metadata = chunk.get("chunk_metadata") or {} + if not isinstance(metadata, dict): + continue + for conn in metadata.get("connect_to") or []: + if not isinstance(conn, dict): + continue + target_id = str(conn.get("target") or "").strip() + if target_id and target_id not in owner_map: + owner_map[target_id] = section_path + return owner_map + + +async def _load_scope_sections( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + scope_paths: list[str], +) -> list[tuple[str, str]]: + section_stmt = ( + select(DocumentSection.section_id, DocumentSection.section_path) + .where(DocumentSection.document_id == document_id) + .where(DocumentSection.job_result_id == job_result_id) + ) + if scope_paths: + scope_filters = [] + for scope in scope_paths: + scope_filters.append(DocumentSection.section_path == scope) + scope_filters.append(DocumentSection.section_path.like(f"{scope} / %")) + section_stmt = section_stmt.where(or_(*scope_filters)) + rows = (await db.execute(section_stmt)).all() + return [(section_id, section_path or "") for section_id, section_path in rows] + + +async def count_assets_under_scope( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + scope_paths: list[str], +) -> tuple[int, int]: + section_rows = await _load_scope_sections( + db, + document_id=document_id, + job_result_id=job_result_id, + scope_paths=scope_paths, + ) + all_section_ids = [section_id for section_id, _section_path in section_rows] + + if not all_section_ids: + return 0, 0 + + count_stmt = ( + select( + DocumentChunk.chunk_type, + sa_func.count(DocumentChunk.id), + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(all_section_ids)) + .where(DocumentChunk.chunk_type.in_(["image", "table"])) + .group_by(DocumentChunk.chunk_type) + ) + count_result = await db.execute(count_stmt) + + total_images = 0 + total_tables = 0 + for chunk_type, count in count_result.all(): + if chunk_type == "image": + total_images = count + elif chunk_type == "table": + total_tables = count + return total_images, total_tables + + +def build_asset_tools_block(total_images: int, total_tables: int) -> str: + if total_images <= 0 and total_tables <= 0: + return "" + + tools_lines = ["\nOptional asset tools (usable with NAVIGATE or STOP):\n"] + if total_images > 0: + tools_lines.append( + f" FIND_IMAGES — Extract image/chart assets under the current scope ({total_images} available).\n" + ) + if total_tables > 0: + tools_lines.append( + f" FIND_TABLES — Extract table/data assets under the current scope ({total_tables} available).\n" + ) + tools_lines.append( + " Note: with NAVIGATE selections, asset tools are limited to the selected sections; " + "with STOP or no selections, they use the current scope.\n" + ) + return "".join(tools_lines) + + +async def resolve_root_asset_owners( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + chunks: list[dict[str, Any]], +) -> dict[str, str]: + root_asset_ids = [ + str(chunk.get("chunk_id") or "") + for chunk in chunks + if not chunk.get("owner_section_path") + and (chunk.get("section_path") or "") == "Root" + and (chunk.get("chunk_type") or "").lower() in ("image", "table") + and chunk.get("chunk_id") + ] + if not root_asset_ids: + return {} + + root_asset_set = set(root_asset_ids) + text_stmt = ( + select( + DocumentChunk.chunk_metadata, + DocumentSection.section_path, + ) + .outerjoin( + DocumentSection, + DocumentSection.section_id == DocumentChunk.section_id, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.chunk_type == "text") + ) + result = await db.execute(text_stmt) + + owner_map: dict[str, str] = {} + for metadata, section_path in result.all(): + if not isinstance(metadata, dict) or not section_path: + continue + for conn in metadata.get("connect_to") or []: + if not isinstance(conn, dict): + continue + target_id = str(conn.get("target") or "").strip() + if target_id in root_asset_set and target_id not in owner_map: + owner_map[target_id] = section_path + + if owner_map: + logger.info( + f" resolve_root_asset_owners: resolved {len(owner_map)}/{len(root_asset_ids)} " + f"Root assets to their owner sections" + ) + return owner_map + + +async def asset_filter_step( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + scope_path: str | list[str] | None, + asset_type: str, +) -> list[dict[str, Any]]: + t0 = time.monotonic() + try: + scope_list = ( + scope_path + if isinstance(scope_path, list) + else [scope_path] + if scope_path + else [] + ) + + section_rows = await _load_scope_sections( + db, + document_id=document_id, + job_result_id=job_result_id, + scope_paths=scope_list, + ) + section_ids = {row[0] for row in section_rows} + + if not section_ids: + logger.info(f" asset_filter_step: no sections found under scope={scope_path}") + return [] + + section_path_by_id = { + section_id: section_path for section_id, section_path in section_rows + } + asset_rows = ( + await db.execute( + select( + DocumentChunk.chunk_id, + DocumentChunk.chunk_type, + DocumentChunk.content, + DocumentChunk.file_path, + DocumentChunk.section_id, + DocumentChunk.source_chunk_path, + DocumentChunk.chunk_metadata, + DocumentChunk.sort_order, + DocumentChunk.job_result_id, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(list(section_ids))) + .where(DocumentChunk.chunk_type == asset_type) + .order_by(DocumentChunk.sort_order) + ) + ).all() + + text_rows = ( + await db.execute( + select( + DocumentChunk.section_id, + DocumentChunk.chunk_type, + DocumentChunk.chunk_metadata, + DocumentChunk.source_chunk_path, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(list(section_ids))) + .where(DocumentChunk.chunk_type == "text") + ) + ).all() + text_row_dicts = [ + { + "chunk_type": chunk_type, + "chunk_metadata": metadata or {}, + "section_id": section_id, + "section_path": section_path_by_id.get(section_id, ""), + "source_chunk_path": source_chunk_path, + } + for section_id, chunk_type, metadata, source_chunk_path in text_rows + ] + owner_by_target_id = build_connected_owner_map(text_row_dicts) + + if any(value == "Root" for value in owner_by_target_id.values()): + doc_stmt = select(Document.source_file_name).where( + Document.document_id == document_id + ) + doc_file_name = (await db.execute(doc_stmt)).scalar() or "" + if doc_file_name: + for target_id in list(owner_by_target_id): + if owner_by_target_id[target_id] == "Root": + owner_by_target_id[target_id] = doc_file_name + + connected_target_ids: set[str] = set(owner_by_target_id.keys()) + if connected_target_ids: + connected_rows = ( + await db.execute( + select( + DocumentChunk.chunk_id, + DocumentChunk.chunk_type, + DocumentChunk.content, + DocumentChunk.file_path, + DocumentChunk.section_id, + DocumentChunk.source_chunk_path, + DocumentChunk.chunk_metadata, + DocumentChunk.sort_order, + DocumentChunk.job_result_id, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.chunk_id.in_(list(connected_target_ids))) + .where(DocumentChunk.chunk_type == asset_type) + .order_by(DocumentChunk.sort_order) + ) + ).all() + else: + connected_rows = [] + + job_id = ( + await db.execute(select(JobResult.job_id).where(JobResult.id == job_result_id)) + ).scalar() or "" + seen_ids: set[str] = set() + chunks: list[dict[str, Any]] = [] + for row in list(asset_rows) + list(connected_rows): + chunk_id = row[0] + if chunk_id in seen_ids: + continue + seen_ids.add(chunk_id) + + owner_section_path = owner_by_target_id.get(chunk_id) + if not owner_section_path: + own_section_path = section_path_by_id.get(row[4]) + if own_section_path and own_section_path == "Root": + logger.warning( + " asset_filter_step: rejecting root-level owner fallback " + f"chunk_id={chunk_id} section_path={own_section_path}" + ) + own_section_path = None + owner_section_path = own_section_path + + if not owner_section_path: + logger.warning( + f" asset_filter_step unresolved owner: chunk_id={chunk_id} " + f"file_path={row[3]} scope={scope_path or 'root'}" + ) + continue + + chunks.append( + { + "document_id": document_id, + "chunk_id": chunk_id, + "chunk_type": row[1], + "content": row[2], + "file_path": row[3], + "section_id": row[4], + "section_path": owner_section_path, + "owner_section_path": owner_section_path, + "source_chunk_path": row[5], + "chunk_metadata": row[6] or {}, + "sort_order": row[7], + "job_result_id": job_result_id, + "job_id": job_id, + } + ) + + latency = int((time.monotonic() - t0) * 1000) + logger.info( + f" asset_filter_step scope={scope_path or 'root'} " + f"type={asset_type}: {len(chunks)} chunks found, {latency}ms" + ) + return chunks + + except Exception as exc: + logger.error(f" asset_filter_step failed: {exc}") + return [] diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py new file mode 100644 index 000000000..6720bd28c --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py @@ -0,0 +1,414 @@ +"""Per-document navigation for agentic retrieval.""" +from __future__ import annotations + +from typing import Any, cast + +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.agentic import tools +from shared.services.retrieval.agentic.core.budget import BudgetExceeded +from shared.services.retrieval.agentic.evidence.builder import reconcile_deferred_assets +from shared.services.retrieval.agentic.core.runtime import AgentLlmBudget +from shared.services.retrieval.agentic.core.trace import TraceRecorder +from shared.services.retrieval.agentic.core.types import ( + AgentRunConfig, + AgentState, + CandidateDoc, + DocTreeNode, + ToolResult, +) +from shared.services.retrieval.llm_adapter import LLMFn + + +class DocumentNavigationRunner: + def __init__( + self, + *, + db: AsyncSession, + state: AgentState, + trace: TraceRecorder, + trace_enabled: bool, + user_id: str, + namespace: str, + query: str, + config: AgentRunConfig, + discovery_by_doc: dict[str, list[dict[str, Any]]], + llm_fn: LLMFn | None, + llm_budget: AgentLlmBudget, + ) -> None: + self._db = db + self._state = state + self._trace = trace + self._trace_enabled = trace_enabled + self._user_id = user_id + self._namespace = namespace + self._query = query + self._config = config + self._discovery_by_doc = discovery_by_doc + self._llm_fn = llm_fn + self._llm_budget = llm_budget + + async def navigate_selected_documents(self, *, revision_hint: str | None) -> None: + logger.info( + f" agentic: Phase 2 — navigating {len(self._state.selected_docs)} documents" + ) + for doc in self._state.selected_docs: + if self._state.elapsed_ms >= self._config.latency_budget_ms: + logger.info(" agentic: latency budget hit during Phase 2, stopping") + break + await self._navigate_document(doc, revision_hint=revision_hint) + + async def _navigate_document( + self, + doc: CandidateDoc, + *, + revision_hint: str | None, + ) -> None: + job_result_id = self._state.doc_job_map.get(doc.document_id, "") + if not job_result_id: + logger.info(f" agentic: skipping doc {doc.document_id} — no job_result_id") + self._state.ever_explored_doc_ids.add(doc.document_id) + return + + doc_name = doc.source_file_name or self._state.doc_id_to_name.get(doc.document_id, "") + is_discovery_only_doc = doc.source == "discovery_auto" + root = DocTreeNode(scope_path=None) + doc_pending_assets: list[dict[str, Any]] = [] + + if not is_discovery_only_doc: + doc_pending_assets = await self._navigate_bfs( + doc=doc, + root=root, + doc_name=doc_name, + job_result_id=job_result_id, + revision_hint=revision_hint, + ) + + await self._hydrate_discovery_hints( + doc=doc, + root=root, + doc_name=doc_name, + revision_hint=revision_hint, + ) + + if not is_discovery_only_doc and doc_pending_assets: + self._reconcile_pending_assets( + doc=doc, + root=root, + doc_name=doc_name, + doc_pending_assets=doc_pending_assets, + ) + + if doc.document_id in self._state.doc_trees: + self._state.doc_trees[doc.document_id].merge(root) + else: + self._state.doc_trees[doc.document_id] = root + self._state.ever_explored_doc_ids.add(doc.document_id) + if self._state.ledger is not None: + self._state.ledger.mark_explored(docs=1) + + async def _navigate_bfs( + self, + *, + doc: CandidateDoc, + root: DocTreeNode, + doc_name: str, + job_result_id: str, + revision_hint: str | None, + ) -> list[dict[str, Any]]: + doc_exclude: set[str] = { + key.split("::", 1)[1] + for key in self._state.seen_section_keys + if key.startswith(f"{doc.document_id}::") + } if self._state.seen_section_keys else set() + pending: list[tuple[str | list[str] | None, DocTreeNode, int]] = [(None, root, 0)] + doc_pending_assets: list[dict[str, Any]] = [] + + while pending: + if self._state.elapsed_ms >= self._config.latency_budget_ms: + break + + scope, parent_node, depth = pending.pop(0) + if depth >= self._config.max_nav_depth: + continue + if self._llm_fn is None: + break + if self._state.ledger and self._state.ledger.status("planning") in ("CRITICAL", "EXHAUSTED"): + logger.info(" agentic: planning budget critical, ending BFS for current doc") + break + + doc_llm_fn = self._llm_budget.for_document( + cast(LLMFn, self._llm_fn), + doc_id=doc.document_id, + depth=depth, + ) + try: + action, asset_tools, step_node, drill_paths = await tools.navigate_step( + self._db, + document_id=doc.document_id, + job_result_id=job_result_id, + query=self._query, + llm_fn=doc_llm_fn, + user_id=self._user_id, + namespace=self._namespace, + doc_name=doc_name, + scope_path=scope, + exclude_paths=doc_exclude, + revision_hint=revision_hint if depth == 0 else None, + budget_snapshot=self._state.ledger.snapshot() if self._state.ledger else None, + ) + except BudgetExceeded: + logger.info(" agentic: planning budget exhausted during navigation") + if self._trace_enabled: + self._trace.record_budget_stop("planning_exhausted") + break + self._state.step_count += 1 + + await self._collect_assets( + doc=doc, + scope=scope, + step_node=step_node, + asset_tools=asset_tools, + pending_assets=doc_pending_assets, + round_scope="nav", + ) + _merge_step_node(parent_node, step_node) + _update_excluded_leaf_paths(doc_exclude, step_node, drill_paths) + _queue_drill_paths(pending, parent_node, drill_paths, depth) + parent_node.reparent_leaf_content() + self._record_navigation_step( + doc=doc, + scope=scope, + depth=depth, + action=action, + asset_tools=asset_tools, + step_node=step_node, + drill_paths=drill_paths, + ) + if self._state.ledger is not None: + self._state.ledger.mark_explored( + chunks=sum(len(chunks) for chunks in step_node.leaf_content.values()), + ) + + return doc_pending_assets + + async def _collect_assets( + self, + *, + doc: CandidateDoc, + scope: str | list[str] | None, + step_node: DocTreeNode, + asset_tools: list[str], + pending_assets: list[dict[str, Any]], + round_scope: str, + ) -> None: + selected_asset_scopes = list(step_node.confidence.keys()) + asset_scope = selected_asset_scopes or scope + for asset_tool in asset_tools: + if asset_tool not in ("FIND_IMAGES", "FIND_TABLES"): + continue + asset_type = "image" if asset_tool == "FIND_IMAGES" else "table" + asset_chunks = await tools.asset_filter_step( + self._db, + document_id=doc.document_id, + job_result_id=self._state.doc_job_map.get(doc.document_id, ""), + scope_path=asset_scope, + asset_type=asset_type, + ) + if asset_chunks: + pending_assets.extend(asset_chunks) + + scope_display = ( + asset_scope + if isinstance(asset_scope, list) + else (asset_scope or "root") + ) + if self._trace_enabled: + self._trace.record_step( + "asset_filter_step", + ToolResult( + status="filtered" if asset_chunks else "empty", + payload={ + "document_id": doc.document_id, + "scope": scope_display, + "navigation_scope": scope if isinstance(scope, str) else (scope or "root"), + "asset_type": asset_type, + "chunks_found": len(asset_chunks) if asset_chunks else 0, + }, + ), + decision_reason=f"asset_{round_scope}_{doc.source_file_name}", + ) + logger.info( + f" agentic step {self._state.step_count}: asset_filter_step " + f'doc="{doc.source_file_name}" scope={scope_display} ' + f"type={asset_type} chunks={len(asset_chunks) if asset_chunks else 0}" + ) + + async def _hydrate_discovery_hints( + self, + *, + doc: CandidateDoc, + root: DocTreeNode, + doc_name: str, + revision_hint: str | None, + ) -> None: + doc_hints = self._discovery_by_doc.get(doc.document_id, []) + if not doc_hints or self._llm_fn is None: + return + if self._state.elapsed_ms >= self._config.latency_budget_ms: + return + + discovery_exclude_paths = { + key.split("::", 1)[1] + for key in root.collect_all_paths(doc.document_id) + } + doc_discovery_llm_fn = self._llm_budget.for_discovery( + cast(LLMFn, self._llm_fn), + doc_id=doc.document_id, + low_priority=root.has_content(), + ) + try: + discovery_node = await tools.discovery_select_step( + self._db, + document_id=doc.document_id, + query=self._query, + llm_fn=doc_discovery_llm_fn, + user_id=self._user_id, + namespace=self._namespace, + doc_name=doc_name, + discovery_hints=doc_hints, + exclude_paths=discovery_exclude_paths, + revision_hint=revision_hint, + budget_snapshot=self._state.ledger.snapshot() if self._state.ledger else None, + ) + except BudgetExceeded: + logger.info(" agentic: planning budget exhausted during discovery selection") + if self._trace_enabled: + self._trace.record_budget_stop("planning_exhausted") + discovery_node = DocTreeNode(scope_path=None) + self._state.step_count += 1 + + if self._trace_enabled: + self._trace.record_step( + "discovery_select_step", + ToolResult( + status="selected" if discovery_node.has_content() else "empty", + payload={ + "document_id": doc.document_id, + "hints_count": len(doc_hints), + "hydrated_count": len(discovery_node.leaf_content), + }, + ), + decision_reason=f"discovery_{doc.source_file_name}", + ) + root.merge(discovery_node) + if self._state.ledger is not None: + self._state.ledger.mark_explored( + chunks=sum(len(chunks) for chunks in discovery_node.leaf_content.values()), + ) + + def _reconcile_pending_assets( + self, + *, + doc: CandidateDoc, + root: DocTreeNode, + doc_name: str, + doc_pending_assets: list[dict[str, Any]], + ) -> None: + if doc_name and not root.children and not any( + item.get("path") == doc_name for item in root.outline_items + ): + root.outline_items.insert(0, {"path": doc_name, "level": 0}) + reconcile_deferred_assets(root, doc_pending_assets) + if self._trace_enabled: + self._trace.record_step( + "deferred_asset_reconcile", + ToolResult( + status="reconciled", + payload={ + "document_id": doc.document_id, + "pending_count": len(doc_pending_assets), + "placed_count": sum( + 1 for asset in doc_pending_assets + if str(asset.get("chunk_id") or "") in { + str(row.get("chunk_id") or "") + for row in root.flatten_chunk_rows() + } + ), + }, + ), + decision_reason=f"deferred_reconcile_{doc.source_file_name}", + ) + + def _record_navigation_step( + self, + *, + doc: CandidateDoc, + scope: str | list[str] | None, + depth: int, + action: str, + asset_tools: list[str], + step_node: DocTreeNode, + drill_paths: list[dict[str, Any]], + ) -> None: + if self._trace_enabled: + self._trace.record_step( + "navigate_step", + ToolResult( + status=f"{action.lower()}" + (" (content)" if step_node.has_content() else ""), + payload={ + "document_id": doc.document_id, + "scope": scope if isinstance(scope, str) else (scope or "root"), + "depth": depth, + "action": action, + "asset_tools": asset_tools, + "outline_count": len(step_node.outline_items), + "leaf_count": len(step_node.leaf_content), + "pending_drills": len(drill_paths), + }, + ), + decision_reason=f"nav_d{depth}_{doc.source_file_name}", + ) + scope_log = scope if isinstance(scope, str) else (", ".join(scope) if scope else "root") + logger.info( + f" agentic step {self._state.step_count}: navigate_step " + f'doc="{doc.source_file_name}" scope={scope_log} ' + f"depth={depth} action={action} tools={asset_tools} " + f"outline={len(step_node.outline_items)} " + f"leaves={len(step_node.leaf_content)} " + f"drills={len(drill_paths)}" + ) + + +def _merge_step_node(parent_node: DocTreeNode, step_node: DocTreeNode) -> None: + parent_node.outline_items = step_node.outline_items + for leaf_path, chunks in step_node.leaf_content.items(): + parent_node.add_leaf_chunks(leaf_path, chunks) + parent_node.confidence = step_node.confidence + + +def _update_excluded_leaf_paths( + doc_exclude: set[str], + step_node: DocTreeNode, + drill_paths: list[dict[str, Any]], +) -> None: + drill_path_set = {str(selection["path"]) for selection in drill_paths} + for leaf_path in step_node.leaf_content: + if leaf_path not in drill_path_set: + doc_exclude.add(leaf_path) + + +def _queue_drill_paths( + pending: list[tuple[str | list[str] | None, DocTreeNode, int]], + parent_node: DocTreeNode, + drill_paths: list[dict[str, Any]], + depth: int, +) -> None: + if not drill_paths: + return + for selection in drill_paths: + child = DocTreeNode(scope_path=selection["path"]) + parent_node.children[selection["path"]] = child + batch_scope = [selection["path"] for selection in drill_paths] + pending.append((batch_scope, parent_node, depth + 1)) diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/knowledge_map.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/knowledge_map.py new file mode 100644 index 000000000..c0196be66 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/knowledge_map.py @@ -0,0 +1,93 @@ +"""Knowledge-map overview for agentic document selection.""" +from __future__ import annotations + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk, GraphNode + + +_MAX_OVERVIEW_FILES = 50 + + +async def build_knowledge_map_overview( + db: AsyncSession, + *, + user_id: str, + namespace: str, +) -> tuple[str, dict[str, str]]: + """Build a file-level knowledge map overview for LLM file selection.""" + doc_stmt = ( + select(Document) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == "active") + .where(Document.current_job_result_id.is_not(None)) + .order_by(Document.updated_at.desc()) + .limit(_MAX_OVERVIEW_FILES) + ) + doc_result = await db.execute(doc_stmt) + documents = list(doc_result.scalars()) + + if not documents: + return "(empty)", {} + + doc_ids = [document.document_id for document in documents] + doc_id_to_name = { + document.document_id: (document.source_file_name or document.document_id) + for document in documents + } + + chunk_stats_stmt = ( + select( + DocumentChunk.document_id, + func.count(DocumentChunk.id).label("chunk_count"), + func.count(func.nullif(DocumentChunk.chunk_type, "text")).label("media_count"), + ) + .join( + Document, + (Document.document_id == DocumentChunk.document_id) + & (Document.current_job_result_id == DocumentChunk.job_result_id), + ) + .where(DocumentChunk.document_id.in_(doc_ids)) + .group_by(DocumentChunk.document_id) + ) + chunk_stats_result = await db.execute(chunk_stats_stmt) + chunk_stats: dict[str, dict[str, int]] = {} + for document_id, chunk_count, media_count in chunk_stats_result.all(): + chunk_stats[document_id] = {"total": chunk_count, "media": media_count} + + graph_summary_stmt = ( + select(GraphNode.owner_document_id, GraphNode.properties) + .where(GraphNode.owner_document_id.in_(doc_ids)) + .where(GraphNode.node_kind == "document") + ) + graph_summary_result = await db.execute(graph_summary_stmt) + doc_top_summaries: dict[str, str] = {} + for document_id, properties in graph_summary_result.all(): + if not isinstance(properties, dict): + continue + top_summary = str(properties.get("top_summary") or "").strip() + if top_summary: + doc_top_summaries[document_id] = top_summary + + lines: list[str] = [] + for document in documents: + document_id = document.document_id + name = doc_id_to_name[document_id] + stats = chunk_stats.get(document_id, {"total": 0, "media": 0}) + top_summary = doc_top_summaries.get(document_id, "") + + line = f'- [{document_id}] {name} chunks={stats["total"]}' + if stats["media"] > 0: + line += f' media={stats["media"]}' + if top_summary: + line += f"\n top_summary:\n{indent_block(top_summary, 4)}" + lines.append(line) + + return "\n".join(lines), doc_id_to_name + + +def indent_block(text: str, spaces: int) -> str: + prefix = " " * spaces + return "\n".join(f"{prefix}{line}" for line in str(text or "").splitlines()) diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/section_counts.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_counts.py new file mode 100644 index 000000000..74ba465d3 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_counts.py @@ -0,0 +1,174 @@ +"""Section count aggregation for agentic navigation.""" +from __future__ import annotations + +from sqlalchemy import case, func, literal_column, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import DocumentChunk + + +async def attach_section_counts( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + all_sections: dict[str, dict], + items_by_path: dict[str, dict], +) -> None: + """Attach direct chunk and connected asset counts to visible section items.""" + scope_item_sids = { + item["section_id"] + for item in items_by_path.values() + if item["show_summary"] + } + all_section_ids = [meta["section_id"] for meta in all_sections.values()] + if not all_section_ids or not scope_item_sids: + return + + section_id_counts = await _load_direct_chunk_counts( + db, + document_id=document_id, + job_result_id=job_result_id, + all_section_ids=all_section_ids, + ) + + sid_to_path = {meta["section_id"]: path for path, meta in all_sections.items()} + for section_id, (text_count, image_count, table_count) in section_id_counts.items(): + chunk_path = sid_to_path.get(section_id, "") + if not chunk_path: + continue + + for item_path, item in items_by_path.items(): + if not item["show_summary"]: + continue + if chunk_path == item_path or chunk_path.startswith(item_path + " / "): + item["chunk_count"] += text_count + item["image_count"] += image_count + item["table_count"] += table_count + + await _attach_connected_asset_counts( + db, + document_id=document_id, + job_result_id=job_result_id, + items_by_path=items_by_path, + sid_to_path=sid_to_path, + ) + + +async def _load_direct_chunk_counts( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + all_section_ids: list[str], +) -> dict[str, tuple[int, int, int]]: + chunk_stmt = ( + select( + DocumentChunk.section_id, + func.count( + case( + (DocumentChunk.chunk_type.notin_(["image", "table"]), literal_column("1")), + ) + ).label("text_count"), + func.count( + case( + (DocumentChunk.chunk_type == "image", literal_column("1")), + ) + ).label("image_count"), + func.count( + case( + (DocumentChunk.chunk_type == "table", literal_column("1")), + ) + ).label("table_count"), + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(all_section_ids)) + .group_by(DocumentChunk.section_id) + ) + chunk_rows = (await db.execute(chunk_stmt)).all() + return { + section_id: (int(text_count), int(image_count), int(table_count)) + for section_id, text_count, image_count, table_count in chunk_rows + } + + +async def _attach_connected_asset_counts( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + items_by_path: dict[str, dict], + sid_to_path: dict[str, str], +) -> None: + scope_items_with_zero_assets = [ + item + for item in items_by_path.values() + if item["show_summary"] and item["image_count"] == 0 and item["table_count"] == 0 + ] + if not scope_items_with_zero_assets: + return + + scope_section_ids = { + item["section_id"] + for item in items_by_path.values() + if item.get("section_id") + } + if not scope_section_ids: + return + + connect_stmt = ( + select( + DocumentChunk.section_id, + DocumentChunk.chunk_metadata, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(list(scope_section_ids))) + .where(DocumentChunk.chunk_type == "text") + ) + connect_result = (await db.execute(connect_stmt)).all() + + section_target_ids: dict[str, set[str]] = {} + for section_id, metadata in connect_result: + if not isinstance(metadata, dict): + continue + for connection in metadata.get("connect_to") or []: + target_id = connection.get("target", "") + if target_id: + section_target_ids.setdefault(section_id, set()).add(target_id) + + if not section_target_ids: + return + + all_target_ids: set[str] = set() + for target_ids in section_target_ids.values(): + all_target_ids.update(target_ids) + + target_type_stmt = ( + select( + DocumentChunk.chunk_id, + DocumentChunk.chunk_type, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.chunk_id.in_(list(all_target_ids))) + .where(DocumentChunk.chunk_type.in_(["image", "table"])) + ) + target_type_result = (await db.execute(target_type_stmt)).all() + target_types = {chunk_id: chunk_type for chunk_id, chunk_type in target_type_result} + + for section_id, target_ids in section_target_ids.items(): + ref_path = sid_to_path.get(section_id, "") + if not ref_path: + continue + referenced_images = sum(1 for target_id in target_ids if target_types.get(target_id) == "image") + referenced_tables = sum(1 for target_id in target_ids if target_types.get(target_id) == "table") + if referenced_images == 0 and referenced_tables == 0: + continue + for item_path, item in items_by_path.items(): + if not item["show_summary"]: + continue + if ref_path == item_path or ref_path.startswith(item_path + " / "): + item["image_count"] += referenced_images + item["table_count"] += referenced_tables diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/section_prompt_projection.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_prompt_projection.py new file mode 100644 index 000000000..b7a251805 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_prompt_projection.py @@ -0,0 +1,59 @@ +"""Prompt projection for agentic section navigation.""" +from __future__ import annotations + +from shared.utils.text_utils import truncate_content_preview + + +def format_items_for_llm( + items: list[dict], + max_chars: int = 20000, +) -> tuple[str, bool]: + """Format section items with hierarchy, selectability, counts, and summaries.""" + if not items: + return "(no items available)", False + + full_text = "\n".join(_render_item(item, include_summary=True) for item in items) + if len(full_text) <= max_chars: + return full_text, False + + slim_text = "\n".join(_render_item(item, include_summary=False) for item in items) + return slim_text[:max_chars], True + + +def _render_item(item: dict, include_summary: bool) -> str: + level = item.get("level", 1) + show_summary = item.get("show_summary", True) + is_leaf = item.get("is_leaf", False) + leaf_tag = " [Leaf]" if is_leaf else "" + path = item.get("path", "") + summary = item.get("summary") or "" + + counts_str = "" + if show_summary: + count_parts: list[str] = [] + chunk_count = item.get("chunk_count", 0) + if chunk_count > 0: + count_parts.append(f"text={chunk_count}") + image_count = item.get("image_count", 0) + if image_count > 0: + count_parts.append(f"image={image_count}") + table_count = item.get("table_count", 0) + if table_count > 0: + count_parts.append(f"table={table_count}") + counts_str = f' [{" ".join(count_parts)}]' if count_parts else "" + + indent = " " * (level - 1) + prefix = "▸" if level == 1 else "└" + level_tag = f"[L{level}]" + select_tag = "[SELECT] " if item.get("selectable", False) else "" + + lines = [ + f'{indent}{prefix} {select_tag}{level_tag} path="{path}"{counts_str}{leaf_tag}' + ] + + if include_summary and show_summary and summary: + sub_indent = " " * level + clipped = truncate_content_preview(summary, head=80, tail=0) + lines.append(f"{sub_indent}{clipped}") + + return "\n".join(lines) diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/section_tree.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_tree.py new file mode 100644 index 000000000..7125fdf4f --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_tree.py @@ -0,0 +1,231 @@ +"""Section-tree loading and prompt projection for agentic navigation.""" +from __future__ import annotations + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import DocumentSection +from shared.services.retrieval.agentic.navigation.section_counts import attach_section_counts +from shared.services.retrieval.search.lexical_text import normalize_section_path, split_section_path + + +async def load_child_sections( + db: AsyncSession, + document_id: str, + job_result_id: str, + scope_path: str | list[str] | None = None, + exclude_paths: set[str] | None = None, +) -> list[dict]: + """Load the continuous context tree for a navigation scope.""" + stmt = ( + select( + DocumentSection.section_id, + DocumentSection.section_title, + DocumentSection.section_path, + DocumentSection.summary, + DocumentSection.sort_order, + ) + .where(DocumentSection.document_id == document_id) + .where(DocumentSection.job_result_id == job_result_id) + .order_by(DocumentSection.sort_order) + ) + section_rows = (await db.execute(stmt)).all() + if not section_rows: + return [] + + if isinstance(scope_path, list): + scope_list = [normalize_section_path(path) for path in scope_path] + elif scope_path: + scope_list = [normalize_section_path(scope_path)] + else: + scope_list = [] + + scope_depth = len(split_section_path(scope_list[0])) if scope_list else 0 + excluded_paths = exclude_paths or set() + + logger.debug( + f" load_child_sections: scopes={scope_list or ['root']} " + f"scope_depth={scope_depth} exclude_paths={excluded_paths if excluded_paths else 'none'} " + f"total_sections={len(section_rows)}" + ) + + all_sections: dict[str, dict] = {} + for section_id, title, path, summary, sort_order in section_rows: + if not path: + continue + normalized_path = normalize_section_path(path) + parts = split_section_path(normalized_path) + all_sections[normalized_path] = { + "title": title or parts[-1] if parts else normalized_path, + "summary": summary or "", + "sort_order": int(sort_order or 0), + "section_id": section_id, + "parts": parts, + "depth": len(parts), + } + + ancestor_prefixes: set[str] = set() + for scope in scope_list: + scope_parts = split_section_path(scope) + for index in range(1, len(scope_parts) + 1): + ancestor_prefixes.add(" / ".join(scope_parts[:index])) + + items_by_path = _select_scope_items( + all_sections, + scope_list=scope_list, + ancestor_prefixes=ancestor_prefixes, + exclude_paths=excluded_paths, + ) + if not items_by_path: + return [] + + allowed_set = _resolve_allowed_depths(items_by_path, scope_list) + if allowed_set: + to_remove = [ + path + for path, item in items_by_path.items() + if item["show_summary"] and item["level"] not in allowed_set + ] + for path in to_remove: + del items_by_path[path] + + if not items_by_path: + return [] + + await attach_section_counts( + db, + document_id=document_id, + job_result_id=job_result_id, + all_sections=all_sections, + items_by_path=items_by_path, + ) + + sorted_items = sorted(items_by_path.values(), key=lambda item: item["sort_order"]) + for item in sorted_items: + item.pop("sort_order", None) + item.pop("section_id", None) + + _mark_leaf_and_selectable(sorted_items, all_section_paths=set(all_sections.keys()), allowed_set=allowed_set) + return sorted_items + + +def _select_scope_items( + all_sections: dict[str, dict], + *, + scope_list: list[str], + ancestor_prefixes: set[str], + exclude_paths: set[str], +) -> dict[str, dict]: + items_by_path: dict[str, dict] = {} + + def is_excluded(path: str) -> bool: + return bool( + exclude_paths + and any(path == excluded or path.startswith(excluded + " / ") for excluded in exclude_paths) + ) + + for path, meta in all_sections.items(): + parts = meta["parts"] + depth = meta["depth"] + + if not scope_list: + if depth < 1 or is_excluded(path): + continue + items_by_path[path] = _make_item(path, meta, show_summary=True) + continue + + matched_scope = _find_matched_scope(parts, depth=depth, scope_list=scope_list) + if matched_scope: + if is_excluded(path): + continue + items_by_path[path] = _make_item(path, meta, show_summary=True) + continue + + max_scope_depth = max(len(split_section_path(scope)) for scope in scope_list) + if depth <= max_scope_depth: + if depth == 1 and path in ancestor_prefixes: + items_by_path.setdefault(path, _make_item(path, meta, show_summary=False)) + elif depth > 1: + parent_prefix = " / ".join(parts[:-1]) + if parent_prefix in ancestor_prefixes: + items_by_path.setdefault(path, _make_item(path, meta, show_summary=False)) + + return items_by_path + + +def _make_item(path: str, meta: dict, show_summary: bool) -> dict: + return { + "path": path, + "title": meta["title"], + "summary": meta["summary"], + "level": meta["depth"], + "sort_order": meta["sort_order"], + "chunk_count": 0, + "image_count": 0, + "table_count": 0, + "section_id": meta["section_id"], + "show_summary": show_summary, + } + + +def _find_matched_scope(parts: list[str], *, depth: int, scope_list: list[str]) -> str | None: + for scope in scope_list: + scope_parts = split_section_path(scope) + scope_depth = len(scope_parts) + if depth > scope_depth and parts[:scope_depth] == scope_parts: + return scope + return None + + +def _resolve_allowed_depths(items_by_path: dict[str, dict], scope_list: list[str]) -> set[int]: + if not scope_list: + depths = { + item["level"] + for item in items_by_path.values() + if item.get("show_summary", True) + } + return set(sorted(depths)[:2]) + + allowed_set: set[int] = set() + for scope in scope_list: + scope_parts = split_section_path(scope) + scope_depth = len(scope_parts) + child_depths = { + item["level"] + for item in items_by_path.values() + if item.get("show_summary", True) + and item["level"] > scope_depth + and split_section_path(item["path"])[:scope_depth] == scope_parts + } + if child_depths: + allowed_set.update(sorted(child_depths)[:2]) + return allowed_set + + +def _mark_leaf_and_selectable( + sorted_items: list[dict], + *, + all_section_paths: set[str], + allowed_set: set[int], +) -> None: + for item in sorted_items: + item_path = item["path"] + has_descendants = any( + path != item_path and path.startswith(item_path + " / ") + for path in all_section_paths + ) + item["is_leaf"] = not has_descendants + + if allowed_set: + shallowest_band = min(allowed_set) + for item in sorted_items: + if not item.get("show_summary", True): + item["selectable"] = False + elif item["level"] == shallowest_band and not item.get("is_leaf", False): + item["selectable"] = False + else: + item["selectable"] = True + else: + for item in sorted_items: + item["selectable"] = item.get("show_summary", True) diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/selection_hydration.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/selection_hydration.py new file mode 100644 index 000000000..e75e32759 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/selection_hydration.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.agentic.core.types import DocTreeNode +from shared.services.retrieval.agentic.navigation import assets as asset_tools +from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows +from shared.services.retrieval.hydration.path import hydrate_paths_to_rows + + +async def hydrate_path_selections_into_node( + db: AsyncSession, + *, + node: DocTreeNode, + path_selections: list[dict[str, Any]], + user_id: str, + namespace: str, + document_id: str, + job_result_id: str | None = None, +) -> None: + chunks = await hydrate_paths_to_rows( + db, + path_selections=path_selections, + user_id=user_id, + namespace=namespace, + document_id=document_id, + ) + if not chunks: + return + + chunks = await _append_connected_asset_targets(db, chunks) + resolved_job_result_id = job_result_id or _find_job_result_id(chunks) + if resolved_job_result_id: + await _attach_root_asset_owners( + db, + document_id=document_id, + job_result_id=resolved_job_result_id, + chunks=chunks, + ) + + add_chunks_to_node(node, chunks) + + +async def _append_connected_asset_targets( + db: AsyncSession, chunks: list[dict[str, Any]] +) -> list[dict[str, Any]]: + connected = await hydrate_connected_target_rows( + db=db, + rows=chunks, + exclude_document_ids=[], + exclude_sections=[], + ) + if not connected: + return chunks + + owner_map = asset_tools.build_connected_owner_map(chunks) + for chunk in connected: + if not chunk.get("owner_section_path"): + chunk["owner_section_path"] = owner_map.get(str(chunk.get("chunk_id") or "")) + return [*chunks, *connected] + + +async def _attach_root_asset_owners( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + chunks: list[dict[str, Any]], +) -> None: + root_map = await asset_tools.resolve_root_asset_owners( + db, + document_id=document_id, + job_result_id=job_result_id, + chunks=chunks, + ) + if not root_map: + return + + for chunk in chunks: + if chunk.get("owner_section_path"): + continue + chunk_id = str(chunk.get("chunk_id") or "") + if chunk_id in root_map: + chunk["owner_section_path"] = root_map[chunk_id] + + +def _find_job_result_id(chunks: list[dict[str, Any]]) -> str | None: + return next( + (str(chunk["job_result_id"]) for chunk in chunks if chunk.get("job_result_id")), + None, + ) + + +def add_chunks_to_node(node: DocTreeNode, chunks: list[dict[str, Any]]) -> None: + for chunk in chunks: + real_path = ( + chunk.get("owner_section_path") + or chunk.get("section_path") + or chunk.get("source_chunk_path") + ) + if real_path: + node.add_leaf_chunks(str(real_path), [chunk]) diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py new file mode 100644 index 000000000..79753a2f3 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py @@ -0,0 +1,186 @@ +"""Agentic retrieval navigation tools. + +This Module owns document-scope navigation and post-navigation discovery +selection. It keeps the LLM prompt, section traversal, hydration, and asset +owner reconciliation local to the navigation seam. +""" +from __future__ import annotations + +from typing import Any + +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.agentic.navigation.assets import ( + build_asset_tools_block, + count_assets_under_scope, +) +from shared.services.retrieval.agentic.core.budget import BudgetExceeded +from shared.services.retrieval.agentic.prompts import ( + ACTION_PROMPT, + format_budget_block, + parse_action_response, +) +from shared.services.retrieval.agentic.navigation.section_prompt_projection import format_items_for_llm +from shared.services.retrieval.agentic.navigation.section_tree import load_child_sections +from shared.services.retrieval.agentic.navigation.selection_hydration import ( + hydrate_path_selections_into_node, +) +from shared.services.retrieval.agentic.core.types import DocTreeNode +from shared.services.retrieval.llm_adapter import LLMFn + + +async def navigate_step( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + query: str, + llm_fn: LLMFn, + user_id: str, + namespace: str, + doc_name: str = "", + scope_path: str | list[str] | None = None, + exclude_paths: set[str] | None = None, + revision_hint: str | None = None, + budget_snapshot: dict | None = None, +) -> tuple[str, list[str], DocTreeNode, list[dict]]: + """Navigate one document scope and hydrate selected sections.""" + scope_paths = ( + scope_path if isinstance(scope_path, list) + else [scope_path] if scope_path + else [] + ) + scope_path_set = set(scope_paths) + + empty = DocTreeNode.empty(scope_paths[0] if scope_paths else None) + + try: + items = await load_child_sections( + db, + document_id, + job_result_id, + scope_path, + exclude_paths=exclude_paths, + ) + if not items: + return "STOP", [], empty, [] + + selectable = { + item["path"]: item for item in items if item.get("selectable", False) + } + total_images, total_tables = await count_assets_under_scope( + db, + document_id=document_id, + job_result_id=job_result_id, + scope_paths=scope_paths, + ) + tools_block = build_asset_tools_block(total_images, total_tables) + + items_text, overflowed = format_items_for_llm(items) + prompt = _build_navigation_prompt( + document_id=document_id, + doc_name=doc_name, + query=query, + scope_paths=scope_paths, + budget_snapshot=budget_snapshot, + items_text=items_text, + tools_block=tools_block, + revision_hint=revision_hint, + ) + + response = await llm_fn(prompt) + parsed = parse_action_response(response) + action = parsed["action"] + selected_tools = parsed["tools"] + selections = parsed["selections"] + + scope_label = ", ".join(scope_paths) if scope_paths else "root" + logger.info( + f" navigate_step scope={scope_label}: " + f"action={action} tools={selected_tools} " + f"selections={len(selections)} selectable={len(selectable)} " + f"overflowed={overflowed}" + ) + + node = DocTreeNode(scope_path=scope_paths[0] if scope_paths else None) + node.outline_items = [item for item in items if item.get("show_summary", True)] + + valid_selections = [ + selection + for selection in selections + if selection["path"] in selectable and selection["path"] not in scope_path_set + ] + + pending: list[dict] = [] + path_selections: list[dict[str, Any]] = [] + for selection in valid_selections: + path = selection["path"] + confidence = selection.get("confidence", 0.7) + item = selectable[path] + node.confidence[path] = confidence + + if item.get("is_leaf"): + path_selections.append({ + "path": path, + "confidence": confidence, + "hydrate_mode": "chunks", + }) + else: + pending.append({"path": path, "confidence": confidence}) + path_selections.append({ + "path": path, + "confidence": confidence, + "hydrate_mode": "self_only", + }) + + await hydrate_path_selections_into_node( + db, + node=node, + path_selections=path_selections, + user_id=user_id, + namespace=namespace, + document_id=document_id, + job_result_id=job_result_id, + ) + + return action, selected_tools, node, pending + + except BudgetExceeded: + raise + except Exception as exc: + logger.error(f" navigate_step failed for doc={document_id}: {exc}") + return "STOP", [], empty, [] +def _build_navigation_prompt( + *, + document_id: str, + doc_name: str, + query: str, + scope_paths: list[str], + budget_snapshot: dict | None, + items_text: str, + tools_block: str, + revision_hint: str | None, +) -> str: + if not scope_paths: + scope_header = "Current scope: root (document top level)" + elif len(scope_paths) == 1: + scope_header = f'Current scope: navigating into "{scope_paths[0]}"' + else: + scope_header = f"Current scope: navigating into {len(scope_paths)} sections" + + prompt = ACTION_PROMPT.format( + doc_name=doc_name or document_id, + doc_id=document_id, + scope_header=scope_header, + budget_block=format_budget_block(budget_snapshot), + items_overview=items_text, + query=query, + tools_block=tools_block, + ) + if revision_hint: + prompt += ( + "\n\nIMPORTANT: Previous round feedback: " + f'"{revision_hint}". Adjust your selections accordingly.' + ) + return prompt diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py index 8aa58964d..4ffdda5d8 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py @@ -18,446 +18,38 @@ from __future__ import annotations import os -import json from typing import Any, cast from loguru import logger -from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import Document, DocumentChunk, RetrievalHitStat - -from shared.services.retrieval.agentic.budget import BudgetExceeded, BudgetLedger, BudgetPoolName -from shared.services.retrieval.agentic.trace import TraceRecorder -from shared.services.retrieval.agentic.types import ( +from shared.services.retrieval.agentic.core.budget import BudgetExceeded, BudgetLedger +from shared.services.retrieval.agentic.discovery.phase import ( + register_discovery_documents, + run_initial_discovery, + select_revision_documents, +) +from shared.services.retrieval.agentic.navigation.document import DocumentNavigationRunner +from shared.services.retrieval.agentic.evidence.builder import ( + build_asset_url_map as _build_asset_url_map, + collect_media_chunks_all as _collect_media_chunks_all, + render_evidence as _render_evidence, + trim_evidence_to_budget as _trim_evidence_to_budget, + with_context_prompt_projection as _with_context_prompt_projection, +) +from shared.services.retrieval.agentic.core.runtime import ( + AgentLlmBudget, + build_config_from_env as _build_config_from_env, + load_budget_inventory as _load_budget_inventory, +) +from shared.services.retrieval.agentic.core.trace import TraceRecorder +from shared.services.retrieval.agentic.core.types import ( AgentRunConfig, AgentState, AgenticResult, - CandidateDoc, - DocTreeNode, ToolResult, ) -from shared.services.retrieval.app_service import ( - generate_retrieval_asset_url, - _is_client_result_artifact_ref, -) from shared.services.retrieval.llm_adapter import LLMFn -from shared.services.retrieval.llm_adapter import current_llm_usage -from shared.services.retrieval.hit_stats_service import compute_importance_score -from shared.utils.token_estimate import estimate_tokens - - - -def _with_context_prompt_projection( - snapshot: dict[str, object], - *, - prompt_tokens: int, -) -> dict[str, object]: - """Return a display snapshot that includes the upcoming answer prompt.""" - projected: dict[str, object] = dict(snapshot) - context_raw = projected.get('context') or {} - if not isinstance(context_raw, dict): - return projected - context = dict(context_raw) - used = int(context.get('used', 0) or 0) - reserved = int(context.get('reserved', 0) or 0) - capacity = int(context.get('capacity', 0) or 0) - projected_used = min(capacity, used + max(int(prompt_tokens), 0)) - projected_remaining = max(capacity - projected_used - reserved, 0) - context.update({ - 'used_projected_before_answer': projected_used, - 'answer_prompt_estimate': max(int(prompt_tokens), 0), - 'remaining': projected_remaining, - 'used_pct': 100 if capacity <= 0 else min( - 100, - int(round((projected_used + reserved) * 100 / capacity)), - ), - }) - if projected_remaining <= 0: - context['status'] = 'EXHAUSTED' - elif context['used_pct'] >= 80: - context['status'] = 'CRITICAL' - elif context['used_pct'] >= 50: - context['status'] = 'TIGHT' - else: - context['status'] = 'HEALTHY' - projected['context'] = context - return projected - - - - - -def _collect_media_chunks(node: DocTreeNode) -> list[dict[str, Any]]: - """Recursively collect image/table chunks from a doc tree's leaf_content.""" - media: list[dict[str, Any]] = [] - for chunks in node.leaf_content.values(): - for chunk in chunks: - ct = (chunk.get('chunk_type') or chunk.get('type') or '').strip().lower() - if ct in ('image', 'table'): - media.append(chunk) - for child in node.children.values(): - media.extend(_collect_media_chunks(child)) - return media - - -def _collect_media_chunks_all(doc_trees: dict[str, DocTreeNode]) -> list[dict[str, Any]]: - """Collect media chunks from all doc trees.""" - media: list[dict[str, Any]] = [] - for tree in doc_trees.values(): - media.extend(_collect_media_chunks(tree)) - return media - - -async def _build_asset_url_map( - media_chunks: list[dict[str, Any]], -) -> dict[str, str]: - """Generate presigned asset URLs for media chunks. - - Uses the same ``generate_retrieval_asset_url`` as ``_to_public_response`` - in ``app_service.py`` — no separate logic. - """ - url_map: dict[str, str] = {} - for chunk in media_chunks: - chunk_id = str(chunk.get('chunk_id') or '').strip() - file_path = chunk.get('file_path') or '' - job_id = chunk.get('job_id') or '' - if not chunk_id or not file_path or not job_id: - continue - if not _is_client_result_artifact_ref(file_path): - continue - try: - url = await generate_retrieval_asset_url( - job_id=str(job_id), - artifact_ref=str(file_path), - ) - if url: - url_map[chunk_id] = url - except Exception as e: - logger.warning(f'Failed to generate asset URL for {chunk_id} (ignored): {e}') - return url_map - - -def _collect_all_leaf_paths(node: DocTreeNode) -> set[str]: - """Recursively collect all leaf_content keys across the entire tree.""" - paths = set(node.leaf_content.keys()) - for child in node.children.values(): - paths.update(_collect_all_leaf_paths(child)) - return paths - - -def _collect_visible_paths(node: DocTreeNode) -> set[str]: - """Collect all outline_items paths across the entire tree. - - These are sections that are "visible" in the rendered tree (shown to the - LLM during navigation) even if no chunks have been hydrated into them yet. - Used as fallback targets for asset reconciliation. - """ - paths = {item['path'] for item in node.outline_items if item.get('path')} - for child in node.children.values(): - paths.update(_collect_visible_paths(child)) - return paths - - -def _find_closest_ancestor(path: str, target_paths: set[str]) -> str | None: - """Walk up a section path to find the closest ancestor in target_paths. - - Example: path="kb/file/Ch1/S1.1/S1.1.1", target_paths={"kb/file/Ch1/S1.1"} - → returns "kb/file/Ch1/S1.1" - - Uses the ' / ' separator convention from the section path format. - """ - parts = path.split(' / ') - # Walk from most specific to least specific (skip the full path itself) - for i in range(len(parts) - 1, 0, -1): - ancestor = ' / '.join(parts[:i]) - if ancestor in target_paths: - return ancestor - return None - - -def _reconcile_deferred_assets( - tree: DocTreeNode, - pending_assets: list[dict], -) -> None: - """Place collected assets into the tree based on final navigated paths. - - Called ONCE after the entire BFS + discovery merge completes for a - document. Asset placement uses a two-tier strategy: - - 1. **Exact match**: If the asset's ``owner_section_path`` matches a - leaf_content key, place directly (existing behavior). - 2. **Closest visible ancestor**: If exact match fails, walk up the - owner_section_path hierarchy to find the nearest ancestor that - appears in either leaf_content or outline_items. This handles - the case where the LLM stopped navigation early (e.g. at root) - but still requested images/tables — assets at L3 get attributed - to the visible L2 section on their path. - """ - final_paths = _collect_all_leaf_paths(tree) - visible_paths = _collect_visible_paths(tree) - all_target_paths = final_paths | visible_paths - - if not all_target_paths: - return - - # Collect existing chunk_ids to avoid duplicates - existing_ids = { - str(row.get('chunk_id') or '') - for row in tree.flatten_chunk_rows() - if row.get('chunk_id') - } - - placed = 0 - ancestor_placed = 0 - for asset in pending_assets: - chunk_id = str(asset.get('chunk_id') or '') - if chunk_id and chunk_id in existing_ids: - continue # already in tree via hydrate_connected_target_rows - - owner_path = ( - asset.get('owner_section_path') - or asset.get('section_path') - ) - if not owner_path: - continue - - # Tier 1: exact match in leaf_content or visible outline - target_path = owner_path if owner_path in all_target_paths else None - - # Tier 2: closest visible ancestor fallback - if target_path is None: - target_path = _find_closest_ancestor(owner_path, all_target_paths) - if target_path: - ancestor_placed += 1 - - if target_path is None: - continue # no visible ancestor → discard - - # Place into root; reparent_leaf_content will move to correct child - tree.add_leaf_chunks(target_path, [asset]) - if chunk_id: - existing_ids.add(chunk_id) - placed += 1 - - if placed: - tree.reparent_leaf_content() - logger.info( - f' deferred asset reconcile: {placed}/{len(pending_assets)} ' - f'assets placed into {len(final_paths)} leaf + {len(visible_paths)} visible paths ' - f'(ancestor_fallback={ancestor_placed})' - ) - - -def _build_config_from_env() -> AgentRunConfig: - """Read agent config from environment, with sensible defaults.""" - return AgentRunConfig( - max_revisions=int(os.environ.get('RETRIEVAL_AGENTIC_MAX_REVISIONS', '2')), - max_nav_depth=int(os.environ.get('RETRIEVAL_AGENTIC_MAX_NAV_DEPTH', '3')), - latency_budget_ms=int(os.environ.get('RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS', '12000')), - token_budget_total=int(os.environ.get('RETRIEVAL_AGENTIC_TOKEN_BUDGET_TOTAL', '40000')), - planning_ratio=float(os.environ.get('RETRIEVAL_AGENTIC_PLANNING_RATIO', '0.5')), - bootstrap_budget=int(os.environ.get('RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET', '2000')), - per_doc_min_share=int(os.environ.get('RETRIEVAL_AGENTIC_PER_DOC_MIN_SHARE', '1500')), - inventory_aware=os.environ.get('RETRIEVAL_AGENTIC_INVENTORY_AWARE', 'true') == 'true', - ) - - -def _stringify_llm_input(prompt: Any) -> str: - if isinstance(prompt, str): - return prompt - try: - return json.dumps(prompt, ensure_ascii=False, default=str) - except Exception: - return str(prompt) - - -async def _render_evidence( - db: AsyncSession, - doc_trees: dict[str, DocTreeNode], - doc_id_to_name: dict[str, str], -) -> str: - """Render unified evidence text from doc trees. - - Discovery paths are now handled by ``discovery_select_step`` in Phase 2 - and merged into doc_trees — no separate fallback needed. - """ - from shared.services.retrieval.agent_navigate import render_unified_doc_tree - - # Build asset URL map for all media chunks (images/tables) - # — same pattern as _to_public_response in app_service.py - all_media_chunks: list[dict[str, Any]] = [] - for doc_tree in doc_trees.values(): - all_media_chunks.extend(_collect_media_chunks(doc_tree)) - asset_url_map = await _build_asset_url_map(all_media_chunks) - - # Render unified evidence from doc trees - evidence_parts: list[str] = [] - for doc_id, doc_tree in doc_trees.items(): - if doc_tree.has_content(): - doc_name = doc_id_to_name.get(doc_id, doc_id) - rendered = render_unified_doc_tree(doc_tree, doc_name, asset_lookup=asset_url_map) - if rendered.strip(): - evidence_parts.append(rendered) - - return '\n\n'.join(evidence_parts) if evidence_parts else '(no evidence collected)' - - -async def _load_budget_inventory( - db: AsyncSession, - *, - user_id: str, - namespace: str, - exclude_document_ids: list[str], -) -> tuple[int, int, dict[str, int]]: - stmt = ( - select(Document.document_id, func.count(DocumentChunk.id)) - .join( - DocumentChunk, - (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id), - ) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .group_by(Document.document_id) - ) - if exclude_document_ids: - stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) - - result = await db.execute(stmt) - doc_chunks = {str(doc_id): int(count or 0) for doc_id, count in result.all()} - return sum(doc_chunks.values()), len(doc_chunks), doc_chunks - - -def _iter_leaf_content(node: DocTreeNode): - for path, chunks in node.leaf_content.items(): - yield path, chunks - for child in node.children.values(): - yield from _iter_leaf_content(child) - - -def _collect_confidences(node: DocTreeNode) -> dict[str, float]: - values = dict(node.confidence) - for child in node.children.values(): - for path, score in _collect_confidences(child).items(): - values[path] = max(values.get(path, 0.0), score) - return values - - -def _pop_leaf_path(node: DocTreeNode, path: str) -> bool: - if path in node.leaf_content: - node.leaf_content.pop(path) - return True - for child in node.children.values(): - if _pop_leaf_path(child, path): - return True - return False - - -def _estimate_chunks_tokens(chunks: list[dict[str, Any]]) -> int: - text = '\n'.join(str(chunk.get('content') or '') for chunk in chunks) - return estimate_tokens(text) - - -async def _fetch_importance_norm_scores( - db: AsyncSession, - *, - user_id: str, - namespace: str, - chunk_ids: list[str], -) -> dict[str, float]: - if not chunk_ids: - return {} - stmt = ( - select( - RetrievalHitStat.chunk_id, - RetrievalHitStat.hit_count, - RetrievalHitStat.last_hit_at, - RetrievalHitStat.created_at, - ) - .where(RetrievalHitStat.user_id == user_id) - .where(RetrievalHitStat.namespace == namespace) - .where(RetrievalHitStat.hit_kind == 'chunk') - .where(RetrievalHitStat.chunk_id.in_(chunk_ids)) - ) - result = await db.execute(stmt) - scores: dict[str, float] = {} - for chunk_id, hit_count, last_hit_at, created_at in result.all(): - if chunk_id and last_hit_at and created_at: - scores[str(chunk_id)] = compute_importance_score(hit_count, last_hit_at, created_at) - return scores - - -async def _trim_evidence_to_budget( - db: AsyncSession, - *, - doc_trees: dict[str, DocTreeNode], - doc_id_to_name: dict[str, str], - context_remaining: int, - user_id: str, - namespace: str, - ledger: BudgetLedger | None, - safety_margin: float = 0.9, -) -> str: - full_text = await _render_evidence(db, doc_trees, doc_id_to_name) - target = int(max(context_remaining, 0) * safety_margin) - if estimate_tokens(full_text) <= target: - return full_text - - candidates: list[tuple[str, str, tuple[float, float, float], int]] = [] - for doc_id, tree in doc_trees.items(): - confidence = _collect_confidences(tree) - for path, chunks in _iter_leaf_content(tree): - chunk_ids = [ - str(chunk.get('chunk_id')) - for chunk in chunks - if chunk.get('chunk_id') - ] - importance = 0.0 - importance_scores = await _fetch_importance_norm_scores( - db, - user_id=user_id, - namespace=namespace, - chunk_ids=chunk_ids, - ) - if importance_scores: - importance = max(importance_scores.values()) - discovery_score = ( - float(chunks[0].get('discovery_score', 0.0) or 0.0) - if chunks else 0.0 - ) - score = (float(confidence.get(path, 0.0) or 0.0), discovery_score, importance) - candidates.append((doc_id, path, score, _estimate_chunks_tokens(chunks))) - - current_estimate = estimate_tokens(full_text) - removed: list[dict[str, Any]] = [] - for doc_id, path, _score, token_estimate in sorted( - candidates, - key=lambda item: (item[2], -item[3]), - ): - if current_estimate <= target: - break - if _pop_leaf_path(doc_trees[doc_id], path): - confidence_score, discovery_score, importance_score = _score - removed.append({ - 'document_id': doc_id, - 'document_name': doc_id_to_name.get(doc_id, doc_id), - 'path': path, - 'confidence_score': round(confidence_score, 4), - 'discovery_score': round(discovery_score, 4), - 'importance_score': round(importance_score, 4), - 'token_estimate': token_estimate, - }) - current_estimate = max(current_estimate - token_estimate, 0) - - if ledger is not None: - ledger.trimmed_paths.extend(removed) - logger.info( - f' agentic.trim_evidence: removed={len(removed)} ' - f'est_tokens={current_estimate} target={target}' - ) - return await _render_evidence(db, doc_trees, doc_id_to_name) class RetrievalAgent: @@ -477,82 +69,6 @@ class RetrievalAgent: If ``llm_fn`` is None, the run returns discovery-only results. """ - async def _call_llm_with_budget( - self, - state: AgentState, - llm_fn: LLMFn, - prompt: Any, - *, - pool: BudgetPoolName, - doc_id: str | None = None, - priority: str = 'normal', - ) -> str: - ledger = state.ledger - if ledger is None: - return await llm_fn(prompt) - - prompt_text = _stringify_llm_input(prompt) - est = estimate_tokens(prompt_text) - reserved = await ledger.try_reserve( - pool, - est, - doc_id=doc_id, - priority='low' if priority == 'low' else 'normal', - ) - if not reserved: - raise BudgetExceeded(f'{pool} budget exhausted') - - try: - response = await llm_fn(prompt) - except Exception: - await ledger.refund(pool, est=est, doc_id=doc_id) - raise - - usage = current_llm_usage.get() or {} - actual = int(usage.get('prompt_tokens') or est) - await ledger.commit(pool, actual=actual, est=est, doc_id=doc_id) - return response - - def _budgeted_doc_llm_fn( - self, - state: AgentState, - llm_fn: LLMFn, - *, - doc_id: str, - depth: int, - ) -> LLMFn: - async def _call(prompt): - return await self._call_llm_with_budget( - state, - llm_fn, - prompt, - pool='planning', - doc_id=doc_id, - priority='low' if depth >= 2 else 'normal', - ) - - return _call - - def _budgeted_discovery_llm_fn( - self, - state: AgentState, - llm_fn: LLMFn, - *, - doc_id: str, - low_priority: bool, - ) -> LLMFn: - async def _call(prompt): - return await self._call_llm_with_budget( - state, - llm_fn, - prompt, - pool='planning', - doc_id=doc_id, - priority='low' if low_priority else 'normal', - ) - - return _call - async def run( self, db: AsyncSession, @@ -569,6 +85,7 @@ async def run( filter_mode: str = 'delete', channels: list[str] | None = None, channel_weights: dict[str, float] | None = None, + internal_recall_k: int | None = None, config: AgentRunConfig | None = None, ledger: BudgetLedger | None = None, parent_run_id: str | None = None, @@ -581,7 +98,6 @@ async def run( errors are captured in trace and the best available result is returned. """ - from shared.services.retrieval.agentic import tools from shared.services.retrieval.agentic.policy import ( attempt_answer, estimate_attempt_answer_prompt_tokens, @@ -618,6 +134,7 @@ async def run( 'exclude_document_ids': exclude_document_ids, 'exclude_sections': exclude_sections, 'signal_paths': signal_paths, + 'internal_recall_k': internal_recall_k, }, parent_run_id=parent_run_id, workflow_step_id=workflow_step_id, @@ -636,135 +153,33 @@ async def run( if llm_fn is None: logger.warning('agentic: no llm_fn provided — running discovery-only mode') - planning_llm_fn: LLMFn | None = None bootstrap_llm_fn: LLMFn | None = None context_llm_fn: LLMFn | None = None + llm_budget = AgentLlmBudget(state) if llm_fn is not None: - base_llm_fn = llm_fn - - async def _planning_llm_call(prompt): - return await self._call_llm_with_budget( - state, base_llm_fn, prompt, pool='planning' - ) - - async def _bootstrap_llm_call(prompt): - return await self._call_llm_with_budget( - state, base_llm_fn, prompt, pool='bootstrap' - ) - - async def _context_llm_call(prompt): - return await self._call_llm_with_budget( - state, base_llm_fn, prompt, pool='context' - ) - - planning_llm_fn = _planning_llm_call - bootstrap_llm_fn = _bootstrap_llm_call - context_llm_fn = _context_llm_call - - # Shared kwargs for bottom_discovery - discovery_kwargs: dict[str, Any] = { - 'user_id': user_id, - 'namespace': namespace, - 'query': query, - 'top_k': top_k, - 'exclude_document_ids': exclude_document_ids, - 'exclude_sections': exclude_sections, - 'data_type': data_type, - 'signal_paths': signal_paths, - 'filter_mode': filter_mode, - 'channels': channels, - 'channel_weights': channel_weights, - } - - # ══════════════════════════════════════════════════════════════════ - # Phase 1: Discovery + Document Selection - # ══════════════════════════════════════════════════════════════════ - logger.info(' agentic: Phase 1 — discovery + document selection') + bootstrap_llm_fn = llm_budget.for_pool(llm_fn, pool='bootstrap') + context_llm_fn = llm_budget.for_pool(llm_fn, pool='context') - # 1a. Bottom discovery (always runs) - discovery_result = await tools.bottom_discovery(db, **discovery_kwargs) - state.step_count += 1 - discovery_rows = discovery_result.payload.get('fused_rows', []) if discovery_result.status != 'error' else [] - state.discovery_top_doc_ids = discovery_result.payload.get('top_doc_ids', []) if discovery_result.status != 'error' else [] - - if trace_enabled: - trace.record_step( - 'bottom_discovery', discovery_result, - decision_reason='phase_1_mandatory', - ) - - logger.info( - f' agentic step {state.step_count}: bottom_discovery ' - f'status={discovery_result.status} latency={discovery_result.latency_ms}ms' + discovery_rows = await run_initial_discovery( + db, + state=state, + trace=trace, + trace_enabled=trace_enabled, + user_id=user_id, + namespace=namespace, + query=query, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + data_type=data_type, + signal_paths=signal_paths, + filter_mode=filter_mode, + channels=channels, + channel_weights=channel_weights, + internal_recall_k=internal_recall_k, + bootstrap_llm_fn=bootstrap_llm_fn, ) - # 1b. KG document selection (requires LLM) - if bootstrap_llm_fn is not None: - try: - kg_result = await tools.kg_document_select( - db, - user_id=user_id, - namespace=namespace, - query=query, - llm_fn=bootstrap_llm_fn, - exclude_document_ids=list(state.ever_explored_doc_ids | set(exclude_document_ids)), - budget_snapshot=state.ledger.snapshot() if state.ledger else None, - ) - except BudgetExceeded: - logger.info(' agentic: bootstrap budget exhausted during document selection') - if trace_enabled: - trace.record_budget_stop('bootstrap_exhausted') - kg_result = ToolResult( - status='no_confident_doc', - payload={'reason': 'bootstrap budget exhausted'}, - ) - state.step_count += 1 - - if trace_enabled: - trace.record_step( - 'kg_document_select', kg_result, - decision_reason='phase_1_doc_selection', - ) - - if kg_result.status == 'selected_docs': - for doc_data in kg_result.payload.get('candidate_docs', []): - state.selected_docs.append(CandidateDoc( - document_id=doc_data.get('document_id', ''), - source_file_name=doc_data.get('source_file_name', ''), - confidence=doc_data.get('confidence', 0.0), - reason=doc_data.get('reason', ''), - source=doc_data.get('source', ''), - )) - state.doc_id_to_name.update(kg_result.payload.get('doc_id_to_name', {})) - state.doc_job_map.update(kg_result.payload.get('doc_job_map', {})) - - # If KG returned nothing, use discovery hints - if not state.selected_docs and state.discovery_top_doc_ids: - hint_ids = [d for d in state.discovery_top_doc_ids if d not in state.ever_explored_doc_ids] - if hint_ids: - doc_stmt = ( - select(Document.document_id, Document.source_file_name, Document.current_job_result_id) - .where(Document.document_id.in_(hint_ids)) - ) - doc_result = await db.execute(doc_stmt) - for did, fname, jrid in doc_result.all(): - state.selected_docs.append(CandidateDoc( - document_id=did, - source_file_name=fname or did, - confidence=0.5, - reason='discovery_hint (KG returned 0)', - source='discovery_hint', - )) - state.doc_id_to_name[did] = fname or did - if jrid: - state.doc_job_map[did] = jrid - - logger.info( - f' agentic step {state.step_count}: kg_document_select ' - f'status={kg_result.status} docs={len(state.selected_docs)} ' - f'latency={kg_result.latency_ms}ms' - ) - # If no LLM or no docs selected, return discovery rows directly if not state.selected_docs: logger.info('agentic: no documents selected — returning discovery results') @@ -796,38 +211,17 @@ async def _context_llm_call(prompt): router_used='agentic_discovery_only', ) - # ══════════════════════════════════════════════════════════════════ - # Discovery → Navigation integration - # Group discovery_rows by document for post-BFS discovery selection - # ══════════════════════════════════════════════════════════════════ discovery_by_doc: dict[str, list[dict[str, Any]]] = {} for row in discovery_rows: doc_id = row.get('document_id', '') if doc_id: discovery_by_doc.setdefault(doc_id, []).append(row) - # Auto-register B-class docs (discovery-only, not selected by KG) - selected_doc_ids = {d.document_id for d in state.selected_docs} - for doc_id in discovery_by_doc: - if doc_id not in selected_doc_ids and doc_id not in state.ever_explored_doc_ids: - doc_stmt = ( - select(Document.document_id, Document.source_file_name, Document.current_job_result_id) - .where(Document.document_id == doc_id) - ) - doc_result = await db.execute(doc_stmt) - row_data = doc_result.first() - if row_data: - did, fname, jrid = row_data - state.selected_docs.append(CandidateDoc( - document_id=did, - source_file_name=fname or did, - confidence=0.4, - reason='discovery_auto (not in KG selection)', - source='discovery_auto', - )) - state.doc_id_to_name[did] = fname or did - if jrid: - state.doc_job_map[did] = jrid + await register_discovery_documents( + db, + state=state, + discovery_by_doc=discovery_by_doc, + ) if state.ledger is not None: await state.ledger.allocate_doc_caps({ @@ -849,286 +243,20 @@ async def _context_llm_call(prompt): stop_reason = 'latency_budget' break - # ── Phase 2: Per-Document Navigation ──────────────────────── - logger.info( - f' agentic: Phase 2 (round {round_idx}) — ' - f'navigating {len(state.selected_docs)} documents' + navigation_runner = DocumentNavigationRunner( + db=db, + state=state, + trace=trace, + trace_enabled=trace_enabled, + user_id=user_id, + namespace=namespace, + query=query, + config=config, + discovery_by_doc=discovery_by_doc, + llm_fn=llm_fn, + llm_budget=llm_budget, ) - - for doc in state.selected_docs: - if state.elapsed_ms >= config.latency_budget_ms: - logger.info(' agentic: latency budget hit during Phase 2, stopping') - break - - job_result_id = state.doc_job_map.get(doc.document_id, '') - if not job_result_id: - logger.info(f' agentic: skipping doc {doc.document_id} — no job_result_id') - state.ever_explored_doc_ids.add(doc.document_id) - continue - - doc_name = doc.source_file_name or state.doc_id_to_name.get(doc.document_id, '') - - # B-class docs (discovery_auto) skip BFS, go to discovery_select - is_b_class = doc.source == 'discovery_auto' - - if not is_b_class: - # Build exclude_paths for this doc from seen_section_keys - # Starts with revision-carried paths, then accumulates - # leaf paths hydrated during THIS BFS round to prevent - # re-selection in deeper drill-downs. - doc_exclude: set[str] = { - key.split('::', 1)[1] - for key in state.seen_section_keys - if key.startswith(f'{doc.document_id}::') - } if state.seen_section_keys else set() - - # BFS queue: (scope_path(s), parent_node, depth) - # scope can be: None (root), str, or list[str] (multi-scope) - root = DocTreeNode(scope_path=None) - pending: list[tuple[str | list[str] | None, DocTreeNode, int]] = [(None, root, 0)] - doc_pending_assets: list[dict] = [] # deferred asset reconcile - - while pending: - if state.elapsed_ms >= config.latency_budget_ms: - break - - scope, parent_node, depth = pending.pop(0) - if depth >= config.max_nav_depth: - continue - - if planning_llm_fn is None: - break - if state.ledger and state.ledger.status('planning') in ('CRITICAL', 'EXHAUSTED'): - logger.info(' agentic: planning budget critical, ending BFS for current doc') - break - - doc_llm_fn = self._budgeted_doc_llm_fn( - state, - cast(LLMFn, llm_fn), - doc_id=doc.document_id, - depth=depth, - ) - - # ★ Unified navigate step (supports multi-scope batching) - try: - action, asset_tools, step_node, drill_paths = await tools.navigate_step( - db, - document_id=doc.document_id, - job_result_id=job_result_id, - query=query, - llm_fn=doc_llm_fn, - user_id=user_id, - namespace=namespace, - doc_name=doc_name, - scope_path=scope, - exclude_paths=doc_exclude, - revision_hint=revision_hint if depth == 0 else None, - budget_snapshot=state.ledger.snapshot() if state.ledger else None, - ) - except BudgetExceeded: - logger.info(' agentic: planning budget exhausted during navigation') - if trace_enabled: - trace.record_budget_stop('planning_exhausted') - break - state.step_count += 1 - - # ★ Asset collection (deferred reconcile) — runs if LLM selected tools. - # If this navigation call selected sections, bind asset tools to - # those selections; otherwise keep the current scope (STOP/root). - selected_asset_scopes = list(step_node.confidence.keys()) - asset_scope = selected_asset_scopes or scope - for asset_tool in asset_tools: - if asset_tool not in ('FIND_IMAGES', 'FIND_TABLES'): - continue - asset_type = 'image' if asset_tool == 'FIND_IMAGES' else 'table' - asset_chunks = await tools.asset_filter_step( - db, - document_id=doc.document_id, - job_result_id=job_result_id, - scope_path=asset_scope, - asset_type=asset_type, - ) - if asset_chunks: - doc_pending_assets.extend(asset_chunks) - - scope_display = ( - asset_scope if isinstance(asset_scope, list) - else (asset_scope or 'root') - ) - if trace_enabled: - trace.record_step( - 'asset_filter_step', ToolResult( - status='filtered' if asset_chunks else 'empty', - payload={ - 'document_id': doc.document_id, - 'scope': scope_display, - 'navigation_scope': scope if isinstance(scope, str) else (scope or 'root'), - 'asset_type': asset_type, - 'chunks_found': len(asset_chunks) if asset_chunks else 0, - }, - ), - decision_reason=f'asset_r{round_idx}_d{depth}_{doc.source_file_name}', - ) - - logger.info( - f' agentic step {state.step_count}: asset_filter_step ' - f'doc="{doc.source_file_name}" scope={scope_display} ' - f'type={asset_type} chunks={len(asset_chunks) if asset_chunks else 0}' - ) - - # Merge step result into parent node - parent_node.outline_items = step_node.outline_items - for leaf_path, chunks in step_node.leaf_content.items(): - parent_node.add_leaf_chunks(leaf_path, chunks) - parent_node.confidence = step_node.confidence - - # Accumulate hydrated leaf paths into doc_exclude - drill_path_set = {sel['path'] for sel in drill_paths} - for leaf_path in step_node.leaf_content: - if leaf_path not in drill_path_set: - doc_exclude.add(leaf_path) - - # Queue non-leaf selections as a SINGLE batched item - # (all drill paths expand simultaneously in the next call) - if drill_paths: - for sel in drill_paths: - child = DocTreeNode(scope_path=sel['path']) - parent_node.children[sel['path']] = child - batch_scope = [sel['path'] for sel in drill_paths] - pending.append((batch_scope, parent_node, depth + 1)) - - # Re-parent leaf paths that belong to a child's subtree - parent_node.reparent_leaf_content() - - if trace_enabled: - trace.record_step( - 'navigate_step', ToolResult( - status=f'{action.lower()}' + (' (content)' if step_node.has_content() else ''), - payload={ - 'document_id': doc.document_id, - 'scope': scope if isinstance(scope, str) else (scope or 'root'), - 'depth': depth, - 'action': action, - 'asset_tools': asset_tools, - 'outline_count': len(step_node.outline_items), - 'leaf_count': len(step_node.leaf_content), - 'pending_drills': len(drill_paths), - }, - ), - decision_reason=f'nav_r{round_idx}_d{depth}_{doc.source_file_name}', - ) - - scope_log = scope if isinstance(scope, str) else (', '.join(scope) if scope else 'root') - logger.info( - f' agentic step {state.step_count}: navigate_step ' - f'doc="{doc.source_file_name}" scope={scope_log} ' - f'depth={depth} action={action} tools={asset_tools} ' - f'outline={len(step_node.outline_items)} ' - f'leaves={len(step_node.leaf_content)} ' - f'drills={len(drill_paths)}' - ) - if state.ledger is not None: - state.ledger.mark_explored( - chunks=sum(len(chunks) for chunks in step_node.leaf_content.values()), - ) - else: - # B-class: no BFS, create empty root - root = DocTreeNode(scope_path=None) - - # ── Post-BFS: Discovery selection step ───────────────────── - doc_hints = discovery_by_doc.get(doc.document_id, []) - if doc_hints and planning_llm_fn is not None and state.elapsed_ms < config.latency_budget_ms: - discovery_exclude_paths = { - key.split('::', 1)[1] - for key in root.collect_all_paths(doc.document_id) - } - doc_discovery_llm_fn = self._budgeted_discovery_llm_fn( - state, - cast(LLMFn, llm_fn), - doc_id=doc.document_id, - low_priority=root.has_content(), - ) - try: - discovery_node = await tools.discovery_select_step( - db, - document_id=doc.document_id, - query=query, - llm_fn=doc_discovery_llm_fn, - user_id=user_id, - namespace=namespace, - doc_name=doc_name, - discovery_hints=doc_hints, - exclude_paths=discovery_exclude_paths, - revision_hint=revision_hint, - budget_snapshot=state.ledger.snapshot() if state.ledger else None, - ) - except BudgetExceeded: - logger.info(' agentic: planning budget exhausted during discovery selection') - if trace_enabled: - trace.record_budget_stop('planning_exhausted') - discovery_node = DocTreeNode(scope_path=None) - state.step_count += 1 - - if trace_enabled: - trace.record_step( - 'discovery_select_step', ToolResult( - status='selected' if discovery_node.has_content() else 'empty', - payload={ - 'document_id': doc.document_id, - 'hints_count': len(doc_hints), - 'hydrated_count': len(discovery_node.leaf_content), - }, - ), - decision_reason=f'discovery_r{round_idx}_{doc.source_file_name}', - ) - - # Merge discovery results into BFS tree - root.merge(discovery_node) - if state.ledger is not None: - state.ledger.mark_explored( - chunks=sum(len(chunks) for chunks in discovery_node.leaf_content.values()), - ) - - # ── Deferred asset reconcile ────────────────────────────── - # Assets were collected across all BFS depths but NOT placed - # into the tree yet. Now that the final navigated paths are - # known (BFS + discovery), filter and place only those assets - # whose owner path matches a navigated leaf. - if not is_b_class and doc_pending_assets: - # Inject doc file name as a visible root-level path, but - # ONLY when BFS stopped at root (no children = STOP action). - if doc_name and not root.children and not any( - item.get('path') == doc_name for item in root.outline_items - ): - root.outline_items.insert(0, {'path': doc_name, 'level': 0}) - _reconcile_deferred_assets(root, doc_pending_assets) - if trace_enabled: - trace.record_step( - 'deferred_asset_reconcile', ToolResult( - status='reconciled', - payload={ - 'document_id': doc.document_id, - 'pending_count': len(doc_pending_assets), - 'placed_count': sum( - 1 for a in doc_pending_assets - if str(a.get('chunk_id') or '') in { - str(r.get('chunk_id') or '') - for r in root.flatten_chunk_rows() - } - ), - }, - ), - decision_reason=f'deferred_reconcile_r{round_idx}_{doc.source_file_name}', - ) - - # Merge or store doc tree - if doc.document_id in state.doc_trees: - state.doc_trees[doc.document_id].merge(root) - else: - state.doc_trees[doc.document_id] = root - state.ever_explored_doc_ids.add(doc.document_id) - if state.ledger is not None: - state.ledger.mark_explored(docs=1) + await navigation_runner.navigate_selected_documents(revision_hint=revision_hint) # ── Phase 3: Render evidence + attempt_answer ──────────────── budget_snapshot_before_answer = state.ledger.snapshot() if state.ledger else None @@ -1201,9 +329,7 @@ async def _context_llm_call(prompt): ] async def vlm_context_call(prompt, _vlm_fn=vlm_fn): - return await self._call_llm_with_budget( - state, cast(LLMFn, _vlm_fn), prompt, pool='context' - ) + return await llm_budget.call(cast(LLMFn, _vlm_fn), prompt, pool='context') # Auto-trigger attempt_answer (VLM if images present) try: @@ -1270,36 +396,21 @@ async def vlm_context_call(prompt, _vlm_fn=vlm_fn): if bootstrap_llm_fn is None: stop_reason = 'no_llm' break - try: - kg_result = await tools.kg_document_select( - db, - user_id=user_id, - namespace=namespace, - query=query, - llm_fn=bootstrap_llm_fn, - exclude_document_ids=list(set(exclude_document_ids)), - revision_hint=revision_hint, - budget_snapshot=state.ledger.snapshot() if state.ledger else None, - ) - except BudgetExceeded: - logger.info(' agentic: bootstrap budget exhausted during revision doc selection') - if trace_enabled: - trace.record_budget_stop('bootstrap_exhausted') - stop_reason = 'bootstrap_budget' + revision_stop_reason = await select_revision_documents( + db, + state=state, + trace=trace, + trace_enabled=trace_enabled, + user_id=user_id, + namespace=namespace, + query=query, + exclude_document_ids=exclude_document_ids, + bootstrap_llm_fn=bootstrap_llm_fn, + revision_hint=revision_hint, + ) + if revision_stop_reason is not None: + stop_reason = revision_stop_reason break - state.step_count += 1 - - if kg_result.status == 'selected_docs': - for doc_data in kg_result.payload.get('candidate_docs', []): - state.selected_docs.append(CandidateDoc( - document_id=doc_data.get('document_id', ''), - source_file_name=doc_data.get('source_file_name', ''), - confidence=doc_data.get('confidence', 0.0), - reason=doc_data.get('reason', ''), - source=doc_data.get('source', ''), - )) - state.doc_id_to_name.update(kg_result.payload.get('doc_id_to_name', {})) - state.doc_job_map.update(kg_result.payload.get('doc_job_map', {})) if not state.selected_docs: logger.info(' agentic: revision found no new docs — stopping') diff --git a/packages/shared-python/shared/services/retrieval/agentic/policy.py b/packages/shared-python/shared/services/retrieval/agentic/policy.py index 25323580a..6c201b074 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/policy.py +++ b/packages/shared-python/shared/services/retrieval/agentic/policy.py @@ -18,8 +18,8 @@ from loguru import logger -from shared.services.retrieval.agentic.budget import BudgetExceeded -from shared.services.retrieval.agentic.types import AgentRunConfig, AgentState +from shared.services.retrieval.agentic.core.budget import BudgetExceeded +from shared.services.retrieval.agentic.core.types import AgentRunConfig, AgentState from shared.services.retrieval.llm_adapter import LLMFn from shared.utils.token_estimate import estimate_tokens @@ -27,19 +27,23 @@ def _parse_answer_response(text: str) -> dict[str, Any] | None: """Extract a JSON answer object from LLM response text.""" text = text.strip() - try: - return json.loads(text) - except (json.JSONDecodeError, ValueError): - pass + parsed = _load_json_object(text) + if parsed is not None: + return parsed match = re.search(r'\{.*\}', text, re.DOTALL) if match: - try: - return json.loads(match.group()) - except (json.JSONDecodeError, ValueError): - pass + return _load_json_object(match.group()) return None +def _load_json_object(raw_value: str) -> dict[str, Any] | None: + try: + parsed = json.loads(raw_value) + except (json.JSONDecodeError, ValueError): + return None + return parsed if isinstance(parsed, dict) else None + + def _looks_like_json_wrapper(text: str) -> bool: """Detect malformed JSON-ish answer wrappers without exposing them.""" stripped = text.strip() @@ -78,7 +82,7 @@ def _budget_line_parts(budget_snapshot: dict | None, pool_name: str) -> dict[str min(100, round(remaining_int * 100 / capacity_int)), ) except (TypeError, ValueError): - pass + remaining_pct = 'unknown' return { 'status': pool.get('status', 'HEALTHY'), 'remaining': remaining, diff --git a/packages/shared-python/shared/services/retrieval/agentic/prompts.py b/packages/shared-python/shared/services/retrieval/agentic/prompts.py new file mode 100644 index 000000000..e48e82901 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/prompts.py @@ -0,0 +1,224 @@ +"""Prompt templates and response parsers for agentic retrieval.""" +from __future__ import annotations + +import json +import re +from typing import Any + + +FILE_SELECT_PROMPT = """\ +You are a document routing assistant. + +{budget_block} +Below is a document corpus overview showing all available documents, +their navigation summaries, chunk counts, and media counts. + +=== Document Corpus Overview === +{overview} +=== End Overview === + +User query: {query} +{revision_context} +Based on the query, select documents that may contain relevant information. +If NO document in the corpus is relevant to the query, return an EMPTY array []. +Return ONLY a JSON array of document IDs, e.g.: ["doc_abc123", "doc_def456"] +Do not include any explanation. +""" + + +DISCOVERY_SELECT_PROMPT = """\ +You are a document navigation assistant. + +Document: "{doc_name}" + +{budget_block} +After navigating the document's section tree, the following section paths +were additionally discovered via keyword and semantic search. +They may contain relevant evidence not found through hierarchical navigation. + +=== Discovery Candidates === +{items} +=== End Discovery Candidates === + +User query: {query} +{revision_context} +Select section paths whose content is needed to answer the query. +If none are relevant, return an EMPTY list []. + +Return ONLY a JSON object: +{{"selections": [{{"path": "...", "confidence": }}, ...]}} +Do not include any explanation. +""" + + +ACTION_PROMPT = """\ +You are a document navigation agent. + +Document: "{doc_name}" (id: {doc_id}) + +{budget_block} +{scope_header} +Below is the document's section tree. +Sections tagged [SELECT] are within the current scope and may be selected. +Other sections are shown as structural context only (not selectable). +Nodes marked [Leaf] have no further sub-sections. + +=== Section Tree === +{items_overview} +=== End Section Tree === + +User query: {query} + +=== Available Actions === + +Choose ONE action: + +NAVIGATE — Drill into selected sections for detailed content. + Consider this when the query targets specific topics and you need deeper text evidence. + Select one or more [SELECT] sections. + +STOP — Current scope evidence is sufficient. No further drill-down. + Consider this when: + - The query asks for an outline, overview, or summary + - The query is broad/global, the tree section can fulfill it without drilling into individual sections. + - You have already collected enough evidence at this level. + +{tools_block} + +When action is NAVIGATE, provide selections: +- You may ONLY select sections marked with [SELECT]. + +When action is STOP, selections must be empty. + +Return ONLY a JSON object: +{{"action": "NAVIGATE", "tools": [...], "selections": [{{"path": "...", "confidence": }}, ...]}} +or +{{"action": "STOP", "tools": [...], "selections": []}} +Do not include any explanation. +""" + + +def parse_action_response(text: str) -> dict: + """Parse the unified navigation response from an LLM.""" + text = text.strip() + asset_tools = {"FIND_IMAGES", "FIND_TABLES"} + default = {"action": "NAVIGATE", "tools": [], "selections": []} + + def extract(data: dict) -> dict: + action = str(data.get("action", "NAVIGATE")).strip().upper() + if action not in ("NAVIGATE", "STOP"): + action = "NAVIGATE" + + tools_val = data.get("tools") or [] + if isinstance(tools_val, list): + tools = [ + str(tool).strip().upper() + for tool in tools_val + if str(tool).strip().upper() in asset_tools + ] + else: + tools = [] + + if action == "STOP": + return {"action": action, "tools": tools, "selections": []} + + selections_val = data.get("selections") or [] + selections = [] + if isinstance(selections_val, list): + for selection in selections_val: + if isinstance(selection, dict) and selection.get("path"): + confidence = normalize_confidence(selection.get("confidence", 0.7)) + selections.append({ + "path": str(selection["path"]), + "confidence": confidence or 0.7, + }) + + return {"action": action, "tools": tools, "selections": selections} + + data = _parse_json_object(text) + if data is not None: + return extract(data) + + fence_match = re.search(r"```(?:json)?\s*\n?(.*?)\n?```", text, re.DOTALL) + if fence_match: + data = _parse_json_object(fence_match.group(1).strip()) + if data is not None: + return extract(data) + + brace_match = re.search(r"\{.*\}", text, re.DOTALL) + if brace_match: + data = _parse_json_object(brace_match.group()) + if data is not None: + return extract(data) + + return default + + +def format_budget_block(snapshot: dict | None) -> str: + if not snapshot: + return "" + planning = snapshot.get("planning") or {} + context = snapshot.get("context") or {} + return ( + "=== Resource Status ===\n" + f"Planning Budget: {planning.get('status', 'HEALTHY')} " + f"({planning.get('used_pct', 0)}% used)\n" + f"Context Budget: {context.get('status', 'HEALTHY')} " + f"({context.get('used_pct', 0)}% used)\n" + f"KG Coverage: {snapshot.get('explored_chunks', 0)}/" + f"{snapshot.get('total_chunks', 0)} chunks explored\n" + f"Docs Explored: {snapshot.get('explored_docs', 0)}/" + f"{snapshot.get('total_docs', 0)}\n" + "When budget is TIGHT, prefer fewer high-confidence selections over broad exploration. " + "When CRITICAL, be very selective — only pick paths with strong relevance. Return empty if evidence suffices.\n" + "=== End Resource Status ===\n" + ) + + +def parse_json_array(text: str) -> list[str]: + """Best-effort extraction of a JSON array of strings from LLM response text.""" + result = extract_json_array_payload(text) + return [str(item) for item in result] + + +def extract_json_array_payload(text: str) -> list[Any]: + text = text.strip() + result = _parse_json_array(text) + if result is not None: + return result + match = re.search(r"\[.*?\]", text, re.DOTALL) + if match: + result = _parse_json_array(match.group()) + if result is not None: + return result + return [] + + +def _parse_json_object(raw_value: str) -> dict[str, Any] | None: + try: + result = json.loads(raw_value) + except (ValueError, json.JSONDecodeError): + return None + return result if isinstance(result, dict) else None + + +def _parse_json_array(raw_value: str) -> list[Any] | None: + try: + result = json.loads(raw_value) + except (ValueError, json.JSONDecodeError): + return None + return result if isinstance(result, list) else None + + +def normalize_confidence(value: Any) -> float | None: + if value is None: + return None + if isinstance(value, str): + value = value.strip().rstrip("%") + try: + parsed = float(value) + except (TypeError, ValueError): + return None + if parsed > 1.0: + parsed = parsed / 100.0 + return max(0.0, min(parsed, 1.0)) diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py index 10ab54757..bce60a71b 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py @@ -1,151 +1,22 @@ -"""Agentic retrieval tools — thin wrappers around existing retrieval components. +"""Agentic retrieval tool adapters. -Each tool: - 1. Calls existing functions from channels.py, agent_navigate.py, app_service.py - 2. Returns a unified ToolResult - 3. Never raises — errors are captured in ToolResult.error +Concrete tool implementations live in focused Modules. This file is the stable +adapter seam used by the workflow orchestrator and contract tests. """ from __future__ import annotations -import time from typing import Any -from loguru import logger -from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import Document -from shared.services.retrieval.agentic.budget import BudgetExceeded -from shared.services.retrieval.agentic.types import DocTreeNode, ToolResult -from shared.services.retrieval.agent_navigate import ( - _build_knowledge_map_overview, - _format_items_for_llm, - _load_child_sections, - _parse_json_array, - _parse_action_response, - _ACTION_PROMPT, - _DISCOVERY_SELECT_PROMPT, - _FILE_SELECT_PROMPT, - _format_budget_block, -) -from shared.services.retrieval.app_service import ( - _CHANNEL_WEIGHT_CONTENT, - _CHANNEL_WEIGHT_PATH, - _CHANNEL_WEIGHT_TERM, - _INTERNAL_RECALL_K_MULTIPLIER, - _merge_same_section_rows, - _normalize_row_scores, - _resolve_allowed_chunk_types, - hydrate_connected_target_rows, - merge_channels_rrf, -) -from shared.services.retrieval.channels import content_channel, path_channel, term_channel -from shared.services.retrieval.lexical_text import normalize_section_path +from shared.services.retrieval.agentic.core.types import DocTreeNode, ToolResult +from shared.services.retrieval.agentic.discovery import selection as discovery_selection +from shared.services.retrieval.agentic.discovery import tools as discovery_tools +from shared.services.retrieval.agentic.navigation import assets as asset_tools +from shared.services.retrieval.agentic.navigation import tools as navigation_tools from shared.services.retrieval.llm_adapter import LLMFn -# --------------------------------------------------------------------------- -# Helper: resolve connected asset → owner text chunk section_path -# --------------------------------------------------------------------------- - -def _build_connected_owner_map(text_chunks: list[dict[str, Any]]) -> dict[str, str]: - """Build target_chunk_id → owner text chunk section_path mapping. - - When text chunks reference images/tables via connect_to metadata, - the referenced assets live in Root section. This map lets us attribute - those assets back to the text chunk's section for correct tree placement. - """ - owner_map: dict[str, str] = {} - for chunk in text_chunks: - if (chunk.get('chunk_type') or 'text') != 'text': - continue - section_path = chunk.get('section_path') or '' - if not section_path: - continue - metadata = chunk.get('chunk_metadata') or {} - if not isinstance(metadata, dict): - continue - for conn in metadata.get('connect_to') or []: - if not isinstance(conn, dict): - continue - target_id = str(conn.get('target') or '').strip() - if target_id and target_id not in owner_map: - owner_map[target_id] = section_path - return owner_map - - -async def _resolve_root_asset_owners( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - chunks: list[dict[str, Any]], -) -> dict[str, str]: - """Resolve owner section_path for Root-stranded image/table chunks. - - When Root is hydrated directly (e.g. via discovery selection), the - batch contains standalone image/table chunks whose section_path is - 'Root'. ``_build_connected_owner_map`` cannot help because the - referencing text chunks live in other sections outside the batch. - - This function queries the *entire document* for text chunks with - connect_to metadata, using the same logic as - ``_build_connected_owner_map``, to resolve the true owner. - - Returns target_chunk_id → owner_section_path for Root assets only. - Returns empty dict when there are no Root assets (zero DB overhead). - """ - from shared.models.database.document import DocumentChunk, DocumentSection - - root_asset_ids = [ - str(c.get('chunk_id') or '') - for c in chunks - if not c.get('owner_section_path') # skip if already resolved by batch-level owner map - and (c.get('section_path') or '') == 'Root' - and (c.get('chunk_type') or '').lower() in ('image', 'table') - and c.get('chunk_id') - ] - if not root_asset_ids: - return {} - - root_asset_set = set(root_asset_ids) - - # Query all text chunks in this document for connect_to metadata - text_stmt = ( - select( - DocumentChunk.chunk_metadata, - DocumentSection.section_path, - ) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.chunk_type == 'text') - ) - result = await db.execute(text_stmt) - - owner_map: dict[str, str] = {} - for metadata, section_path in result.all(): - if not isinstance(metadata, dict) or not section_path: - continue - for conn in metadata.get('connect_to') or []: - if not isinstance(conn, dict): - continue - target_id = str(conn.get('target') or '').strip() - if target_id in root_asset_set and target_id not in owner_map: - owner_map[target_id] = section_path - - if owner_map: - logger.info( - f' _resolve_root_asset_owners: resolved {len(owner_map)}/{len(root_asset_ids)} ' - f'Root assets to their owner sections' - ) - return owner_map - - -# --------------------------------------------------------------------------- -# Tool: bottom_discovery -# --------------------------------------------------------------------------- - async def bottom_discovery( db: AsyncSession, *, @@ -157,108 +28,29 @@ async def bottom_discovery( exclude_sections: list[dict[str, str]], data_type: int = 1, signal_paths: list[str] | None = None, - filter_mode: str = 'delete', + filter_mode: str = "delete", channels: list[str] | None = None, channel_weights: dict[str, float] | None = None, internal_recall_k: int | None = None, - **_kwargs: Any, + **kwargs: Any, ) -> ToolResult: - """Run 3-channel BM25 discovery + RRF fusion.""" - t0 = time.monotonic() - try: - allowed_chunk_types = _resolve_allowed_chunk_types(data_type) - effective_recall_k = internal_recall_k if internal_recall_k is not None else top_k * _INTERNAL_RECALL_K_MULTIPLIER - active_channels = set(channels) if channels else {'path', 'content', 'term'} - - path_rows: list[dict[str, Any]] = [] - content_rows: list[dict[str, Any]] = [] - term_rows: list[dict[str, Any]] = [] - - if 'path' in active_channels: - path_rows = await path_channel( - db, user_id=user_id, namespace=namespace, query=query, - top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, filter_mode=filter_mode, - ) - - if 'content' in active_channels: - content_rows = await content_channel( - db, user_id=user_id, namespace=namespace, query=query, - top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, filter_mode=filter_mode, - ) - - if 'term' in active_channels: - term_rows = await term_channel( - db, user_id=user_id, namespace=namespace, query=query, - top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, filter_mode=filter_mode, - ) - - # RRF fusion - default_weights = { - 'path': _CHANNEL_WEIGHT_PATH, - 'content': _CHANNEL_WEIGHT_CONTENT, - 'term': _CHANNEL_WEIGHT_TERM, - } - effective_weights = {**default_weights, **(channel_weights or {})} - - channel_lists: list[list[dict[str, Any]]] = [] - weight_list: list[float] = [] - if path_rows: - channel_lists.append(path_rows) - weight_list.append(effective_weights.get('path', _CHANNEL_WEIGHT_PATH)) - if content_rows: - channel_lists.append(content_rows) - weight_list.append(effective_weights.get('content', _CHANNEL_WEIGHT_CONTENT)) - if term_rows: - channel_lists.append(term_rows) - weight_list.append(effective_weights.get('term', _CHANNEL_WEIGHT_TERM)) - - fused_rows = merge_channels_rrf(channel_lists, weight_list, top_k) if channel_lists else [] - fused_rows = _merge_same_section_rows(fused_rows) - - if fused_rows: - _normalize_row_scores(fused_rows, source_field='score', target_field='discovery_score', default=0.5) - - # Extract top document IDs as hints for KG selection - doc_id_counts: dict[str, int] = {} - for row in fused_rows: - did = row.get('document_id', '') - if did: - doc_id_counts[did] = doc_id_counts.get(did, 0) + 1 - top_doc_ids = sorted(doc_id_counts, key=lambda d: doc_id_counts[d], reverse=True)[:5] - - latency = int((time.monotonic() - t0) * 1000) - logger.info( - f' agentic.bottom_discovery: {len(fused_rows)} fused rows, ' - f'top_doc_ids={top_doc_ids}, {latency}ms' - ) - return ToolResult( - status='discovery_done', - payload={ - 'fused_rows': fused_rows, - 'top_doc_ids': top_doc_ids, - 'channel_counts': { - 'path': len(path_rows), - 'content': len(content_rows), - 'term': len(term_rows), - }, - }, - latency_ms=latency, - ) - except Exception as e: - latency = int((time.monotonic() - t0) * 1000) - logger.error(f' agentic.bottom_discovery failed: {e}') - return ToolResult(status='error', error=str(e), latency_ms=latency) - + return await discovery_tools.bottom_discovery( + db, + user_id=user_id, + namespace=namespace, + query=query, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + data_type=data_type, + signal_paths=signal_paths, + filter_mode=filter_mode, + channels=channels, + channel_weights=channel_weights, + internal_recall_k=internal_recall_k, + **kwargs, + ) -# --------------------------------------------------------------------------- -# Tool: kg_document_select -# --------------------------------------------------------------------------- async def kg_document_select( db: AsyncSession, @@ -269,104 +61,19 @@ async def kg_document_select( llm_fn: LLMFn | None, exclude_document_ids: list[str], revision_hint: str | None = None, - **_kwargs: Any, + **kwargs: Any, ) -> ToolResult: - """Select candidate documents from document-level KG.""" - t0 = time.monotonic() - try: - overview_text, doc_id_to_name = await _build_knowledge_map_overview( - db, user_id=user_id, namespace=namespace, - ) - if overview_text == '(empty)': - latency = int((time.monotonic() - t0) * 1000) - return ToolResult( - status='no_confident_doc', - payload={'reason': 'no active documents in namespace'}, - latency_ms=latency, - ) - - if llm_fn is None: - latency = int((time.monotonic() - t0) * 1000) - return ToolResult( - status='no_confident_doc', - payload={'reason': 'LLM not available'}, - latency_ms=latency, - ) - - revision_context = '' - if revision_hint: - revision_context = ( - f'\nIMPORTANT: This is a REVISION round. ' - f'The previous search attempt failed because:\n' - f'"{revision_hint}"\n' - f'Adjust your document selection accordingly. ' - f'If no document can address this, return an EMPTY array [].\n' - ) - - file_prompt = _FILE_SELECT_PROMPT.format( - overview=overview_text, query=query, - revision_context=revision_context, - budget_block=_format_budget_block(_kwargs.get('budget_snapshot')), - ) - file_response = await llm_fn(file_prompt) - selected_ids = _parse_json_array(file_response) - - exclude_set = set(exclude_document_ids) - valid_ids = [did for did in selected_ids if did in doc_id_to_name and did not in exclude_set] - - if not valid_ids: - latency = int((time.monotonic() - t0) * 1000) - logger.info(f' agentic.kg_document_select: LLM returned no valid docs, {latency}ms') - return ToolResult( - status='no_confident_doc', - payload={'reason': 'LLM returned no valid document IDs', 'raw_ids': selected_ids}, - latency_ms=latency, - ) - - # Load job_result_ids for selected documents - doc_job_map: dict[str, str] = {} - doc_stmt = ( - select(Document.document_id, Document.current_job_result_id) - .where(Document.document_id.in_(valid_ids)) - ) - doc_result = await db.execute(doc_stmt) - for did, jrid in doc_result.all(): - if jrid: - doc_job_map[did] = jrid - - candidate_docs = [] - for did in valid_ids: - candidate_docs.append({ - 'document_id': did, - 'source_file_name': doc_id_to_name.get(did, ''), - 'confidence': 1.0, - 'reason': 'LLM selected from KG overview', - 'source': 'kg_llm_select', - }) - - latency = int((time.monotonic() - t0) * 1000) - logger.info(f' agentic.kg_document_select: {len(candidate_docs)} docs selected, {latency}ms') - return ToolResult( - status='selected_docs', - payload={ - 'candidate_docs': candidate_docs, - 'doc_id_to_name': doc_id_to_name, - 'doc_job_map': doc_job_map, - }, - latency_ms=latency, - ) - except BudgetExceeded: - raise - except Exception as e: - latency = int((time.monotonic() - t0) * 1000) - logger.error(f' agentic.kg_document_select failed: {e}') - return ToolResult(status='error', error=str(e), latency_ms=latency) - - + return await discovery_tools.kg_document_select( + db, + user_id=user_id, + namespace=namespace, + query=query, + llm_fn=llm_fn, + exclude_document_ids=exclude_document_ids, + revision_hint=revision_hint, + **kwargs, + ) -# --------------------------------------------------------------------------- -# Tool: asset_filter_step (programmatic asset extraction) -# --------------------------------------------------------------------------- async def asset_filter_step( db: AsyncSession, @@ -374,218 +81,16 @@ async def asset_filter_step( document_id: str, job_result_id: str, scope_path: str | list[str] | None, - asset_type: str, # 'image' | 'table' + asset_type: str, ) -> list[dict[str, Any]]: - """Extract assets from all descendants under scope_path. - - Terminal action — no LLM involved. - Algorithm: load all text chunks under scope → parse connect_to metadata → - batch-load target image/table chunks → return directly. - - Also collects standalone asset chunks (image/table) that exist directly - under the scope but are not referenced via connect_to. - - scope_path can be: - - None: root scope (entire document) - - str: single scope path - - list[str]: multiple scope paths (queried simultaneously) - """ - from shared.models.database.document import DocumentChunk, DocumentSection - - t0 = time.monotonic() - try: - # 1. Find all section_ids under scope_path(s) - # Normalize scope to list for uniform handling - scope_list = ( - scope_path if isinstance(scope_path, list) - else [scope_path] if scope_path - else [] - ) - - section_stmt = ( - select(DocumentSection.section_id, DocumentSection.section_path) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - ) - if scope_list: - from sqlalchemy import or_ - scope_filters = [] - for sp in scope_list: - scope_filters.append(DocumentSection.section_path == sp) - scope_filters.append(DocumentSection.section_path.like(f'{sp} / %')) - section_stmt = section_stmt.where(or_(*scope_filters)) - section_result = await db.execute(section_stmt) - section_rows = section_result.all() - section_ids = {row[0] for row in section_rows} - - if not section_ids: - logger.info(f' asset_filter_step: no sections found under scope={scope_path}') - return [] - - # 2. Load target asset chunks directly (standalone assets in the scope) - asset_stmt = ( - select( - DocumentChunk.chunk_id, - DocumentChunk.chunk_type, - DocumentChunk.content, - DocumentChunk.file_path, - DocumentChunk.section_id, - DocumentChunk.source_chunk_path, - DocumentChunk.chunk_metadata, - DocumentChunk.sort_order, - DocumentChunk.job_result_id, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(list(section_ids))) - .where(DocumentChunk.chunk_type == asset_type) - .order_by(DocumentChunk.sort_order) - ) - asset_result = await db.execute(asset_stmt) - asset_rows = asset_result.all() - - section_path_by_id = {section_id: section_path for section_id, section_path in section_rows} - - # 3. Resolve media → owner text section via connect_to tracing - text_stmt = ( - select( - DocumentChunk.section_id, - DocumentChunk.chunk_type, - DocumentChunk.chunk_metadata, - DocumentChunk.source_chunk_path, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(list(section_ids))) - .where(DocumentChunk.chunk_type == 'text') - ) - text_result = await db.execute(text_stmt) - text_row_dicts = [ - { - 'chunk_type': chunk_type, - 'chunk_metadata': metadata or {}, - 'section_id': sid, - 'section_path': section_path_by_id.get(sid, ''), - 'source_chunk_path': scp, - } - for sid, chunk_type, metadata, scp in text_result.all() - ] - owner_by_target_id = _build_connected_owner_map(text_row_dicts) - - # Replace synthetic "Root" owner with the document's source_file_name. - # Root is a hybrid node whose real path is the file name (e.g. - # "32_安全大模型技术与市场研究报告_1.docx"); the DB stores the - # synthetic label "Root" which cannot match any outline node. - if any(v == 'Root' for v in owner_by_target_id.values()): - doc_stmt = select(Document.source_file_name).where( - Document.document_id == document_id - ) - doc_file_name = (await db.execute(doc_stmt)).scalar() or '' - if doc_file_name: - for tid in list(owner_by_target_id): - if owner_by_target_id[tid] == 'Root': - owner_by_target_id[tid] = doc_file_name - - # Collect connected target IDs for batch-loading - connected_target_ids: set[str] = set(owner_by_target_id.keys()) - - # Load connected targets that match asset_type - if connected_target_ids: - connected_stmt = ( - select( - DocumentChunk.chunk_id, - DocumentChunk.chunk_type, - DocumentChunk.content, - DocumentChunk.file_path, - DocumentChunk.section_id, - DocumentChunk.source_chunk_path, - DocumentChunk.chunk_metadata, - DocumentChunk.sort_order, - DocumentChunk.job_result_id, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.chunk_id.in_(list(connected_target_ids))) - .where(DocumentChunk.chunk_type == asset_type) - .order_by(DocumentChunk.sort_order) - ) - connected_result = await db.execute(connected_stmt) - connected_rows = connected_result.all() - else: - connected_rows = [] - - # 4. Merge and deduplicate - seen_ids: set[str] = set() - chunks: list[dict[str, Any]] = [] - - # Helper to look up job_id from job_result - from shared.models.database.job_result import JobResult - job_stmt = ( - select(JobResult.job_id) - .where(JobResult.id == job_result_id) - ) - job_result_row = await db.execute(job_stmt) - job_id = job_result_row.scalar() or '' - - for row in list(asset_rows) + list(connected_rows): - chunk_id = row[0] - if chunk_id in seen_ids: - continue - seen_ids.add(chunk_id) - - # Owner resolution: prefer connect_to-based owner - owner_section_path = owner_by_target_id.get(chunk_id) - - # Fallback: media's own section_id path, but guard against - # Root / top-level aggregation sections - if not owner_section_path: - own_section_path = section_path_by_id.get(row[4]) - if own_section_path and own_section_path == 'Root': - # Reject only the synthetic Root aggregation label; - # legitimate L1 sections (e.g. "前言") are valid owners. - logger.warning( - f' asset_filter_step: rejecting root-level owner fallback ' - f'chunk_id={chunk_id} section_path={own_section_path}' - ) - own_section_path = None - owner_section_path = own_section_path - - if not owner_section_path: - logger.warning( - f' asset_filter_step unresolved owner: chunk_id={chunk_id} ' - f'file_path={row[3]} scope={scope_path or "root"}' - ) - continue - chunks.append({ - 'document_id': document_id, - 'chunk_id': chunk_id, - 'chunk_type': row[1], - 'content': row[2], - 'file_path': row[3], - 'section_id': row[4], - 'section_path': owner_section_path, - 'owner_section_path': owner_section_path, - 'source_chunk_path': row[5], - 'chunk_metadata': row[6] or {}, - 'sort_order': row[7], - 'job_result_id': job_result_id, - 'job_id': job_id, - }) - - latency = int((time.monotonic() - t0) * 1000) - logger.info( - f' asset_filter_step scope={scope_path or "root"} ' - f'type={asset_type}: {len(chunks)} chunks found, {latency}ms' - ) - return chunks - - except Exception as e: - logger.error(f' asset_filter_step failed: {e}') - return [] + return await asset_tools.asset_filter_step( + db, + document_id=document_id, + job_result_id=job_result_id, + scope_path=scope_path, + asset_type=asset_type, + ) -# --------------------------------------------------------------------------- -# Tool: navigate_step (unified action — merges tool_select + scope_navigate) -# --------------------------------------------------------------------------- async def navigate_step( db: AsyncSession, @@ -596,226 +101,26 @@ async def navigate_step( llm_fn: LLMFn, user_id: str, namespace: str, - doc_name: str = '', + doc_name: str = "", scope_path: str | list[str] | None = None, exclude_paths: set[str] | None = None, revision_hint: str | None = None, budget_snapshot: dict | None = None, ) -> tuple[str, list[str], DocTreeNode, list[dict]]: - """Unified navigation step — one LLM call for action + tools + selections. - - scope_path can be: - - None: root scope - - str: single scope to drill into - - list[str]: multiple scopes to expand simultaneously - - Returns: - - action: 'STOP' | 'NAVIGATE' - - asset_tools: list of asset tools to run (FIND_IMAGES, FIND_TABLES) - - node: DocTreeNode with outline_items and leaf_content - - pending: list of {path, confidence} for non-leaf drill-downs (empty when STOP) - """ - from shared.services.retrieval.app_service import _hydrate_paths_to_rows - - # Normalize scope for internal use - scope_paths: list[str] = ( - scope_path if isinstance(scope_path, list) - else [scope_path] if scope_path - else [] + return await navigation_tools.navigate_step( + db, + document_id=document_id, + job_result_id=job_result_id, + query=query, + llm_fn=llm_fn, + user_id=user_id, + namespace=namespace, + doc_name=doc_name, + scope_path=scope_path, + exclude_paths=exclude_paths, + revision_hint=revision_hint, + budget_snapshot=budget_snapshot, ) - # Set of scope path strings (for filtering selections) - scope_path_set = set(scope_paths) - - empty = DocTreeNode.empty(scope_paths[0] if scope_paths else None) - - try: - # 1. Load continuous context tree (supports multi-scope) - items = await _load_child_sections( - db, document_id, job_result_id, scope_path, - exclude_paths=exclude_paths, - ) - if not items: - return 'STOP', [], empty, [] - - # 2. Build selectable index - selectable = {item['path']: item for item in items if item.get('selectable', False)} - - # 3. Count ALL image/table chunks under the scope subtree(s) - from shared.models.database.document import DocumentChunk, DocumentSection - from sqlalchemy import func as sa_func - - scope_section_stmt = ( - select(DocumentSection.section_id) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - ) - if scope_paths: - from sqlalchemy import or_ - scope_filters = [] - for sp in scope_paths: - scope_filters.append(DocumentSection.section_path == sp) - scope_filters.append(DocumentSection.section_path.like(f'{sp} / %')) - scope_section_stmt = scope_section_stmt.where(or_(*scope_filters)) - scope_section_ids = await db.execute(scope_section_stmt) - all_section_ids = [r[0] for r in scope_section_ids.all()] - - total_images = 0 - total_tables = 0 - if all_section_ids: - count_stmt = ( - select( - DocumentChunk.chunk_type, - sa_func.count(DocumentChunk.id), - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(all_section_ids)) - .where(DocumentChunk.chunk_type.in_(['image', 'table'])) - .group_by(DocumentChunk.chunk_type) - ) - count_result = await db.execute(count_stmt) - for chunk_type, cnt in count_result.all(): - if chunk_type == 'image': - total_images = cnt - elif chunk_type == 'table': - total_tables = cnt - - tools_block = '' - if total_images > 0 or total_tables > 0: - tools_lines = ['\nOptional asset tools (usable with NAVIGATE or STOP):\n'] - if total_images > 0: - tools_lines.append( - f' FIND_IMAGES — Extract image/chart assets under the current scope ({total_images} available).\n' - ) - if total_tables > 0: - tools_lines.append( - f' FIND_TABLES — Extract table/data assets under the current scope ({total_tables} available).\n' - ) - tools_lines.append( - ' Note: with NAVIGATE selections, asset tools are limited to the selected sections; ' - 'with STOP or no selections, they use the current scope.\n' - ) - tools_block = ''.join(tools_lines) - - # 4. Format tree and build prompt - text, overflowed = _format_items_for_llm(items) - if not scope_paths: - scope_header = 'Current scope: root (document top level)' - elif len(scope_paths) == 1: - scope_header = f'Current scope: navigating into "{scope_paths[0]}"' - else: - scope_header = f'Current scope: navigating into {len(scope_paths)} sections' - prompt = _ACTION_PROMPT.format( - doc_name=doc_name or document_id, - doc_id=document_id, - scope_header=scope_header, - budget_block=_format_budget_block(budget_snapshot), - items_overview=text, - query=query, - tools_block=tools_block, - ) - if revision_hint: - prompt += ( - f'\n\nIMPORTANT: Previous round feedback: ' - f'"{revision_hint}". Adjust your selections accordingly.' - ) - - # 5. Single LLM call - response = await llm_fn(prompt) - parsed = _parse_action_response(response) - action = parsed['action'] - asset_tools = parsed['tools'] - selections = parsed['selections'] - - scope_label = ', '.join(scope_paths) if scope_paths else 'root' - logger.info( - f' navigate_step scope={scope_label}: ' - f'action={action} tools={asset_tools} ' - f'selections={len(selections)} selectable={len(selectable)} ' - f'overflowed={overflowed}' - ) - - # 6. Build node with LOCAL items only - node = DocTreeNode(scope_path=scope_paths[0] if scope_paths else None) - local_items = [item for item in items if item.get('show_summary', True)] - node.outline_items = local_items - - # 7. Dispatch selections (only present when action == NAVIGATE) - valid_selections = [ - s for s in selections - if s['path'] in selectable and s['path'] not in scope_path_set - ] - - pending: list[dict] = [] - path_selections = [] - for sel in valid_selections: - path = sel['path'] - conf = sel.get('confidence', 0.7) - item = selectable[path] - node.confidence[path] = conf - - if item.get('is_leaf'): - path_selections.append({'path': path, 'confidence': conf, 'hydrate_mode': 'chunks'}) - else: - # Non-leaf → will be batched into a single next call - pending.append({'path': path, 'confidence': conf}) - path_selections.append({'path': path, 'confidence': conf, 'hydrate_mode': 'self_only'}) - - if path_selections: - chunks = await _hydrate_paths_to_rows( - db, - path_selections=path_selections, - user_id=user_id, - namespace=namespace, - document_id=document_id, - ) - if chunks: - connected = await hydrate_connected_target_rows( - db=db, - rows=chunks, - exclude_document_ids=[], - exclude_sections=[], - ) - if connected: - _owner_map = _build_connected_owner_map(chunks) - for c in connected: - if not c.get('owner_section_path'): - c['owner_section_path'] = _owner_map.get(str(c.get('chunk_id') or '')) - chunks = chunks + connected - - _root_map = await _resolve_root_asset_owners( - db, - document_id=document_id, - job_result_id=job_result_id, - chunks=chunks, - ) - if _root_map: - for c in chunks: - if c.get('owner_section_path'): - continue - cid = str(c.get('chunk_id') or '') - if cid in _root_map: - c['owner_section_path'] = _root_map[cid] - - for chunk in chunks: - real_path = chunk.get('owner_section_path') or chunk.get('section_path') or chunk.get('source_chunk_path') - if real_path: - node.add_leaf_chunks(str(real_path), [chunk]) - - return action, asset_tools, node, pending - - except BudgetExceeded: - raise - except Exception as e: - logger.error(f' navigate_step failed for doc={document_id}: {e}') - return 'STOP', [], empty, [] - - -# --------------------------------------------------------------------------- -# Tool: discovery_select_step (post-navigation discovery selection) -# --------------------------------------------------------------------------- - -_MAX_DISCOVERY_PER_DOC = 3 async def discovery_select_step( @@ -826,152 +131,22 @@ async def discovery_select_step( llm_fn: LLMFn, user_id: str, namespace: str, - doc_name: str = '', + doc_name: str = "", discovery_hints: list[dict[str, Any]], exclude_paths: set[str] | None = None, revision_hint: str | None = None, budget_snapshot: dict | None = None, ) -> DocTreeNode: - """Post-navigation discovery selection step. - - After BFS navigation exhausts for a document, present discovery-found - section paths (from bottom_discovery BM25) to the LLM for selection. - Selected paths are hydrated as leaf content. - - For B-class documents (discovery-only, not KG-selected), this is the - only navigation step — no prior BFS. - """ - from shared.services.retrieval.app_service import _hydrate_paths_to_rows - - node = DocTreeNode(scope_path=None) - if not discovery_hints: - return node - - # Limit hints per document - hints = discovery_hints[:_MAX_DISCOVERY_PER_DOC] - - t0 = time.monotonic() - try: - # 1. Format hints for LLM (deduplicate by section_path) - exclude_set = { - normalize_section_path(path) - for path in (exclude_paths or set()) - if path - } - hint_lines: list[str] = [] - hint_by_path: dict[str, dict] = {} - for h in hints: - sp = normalize_section_path(h.get('section_path', '')) - if not sp or sp == 'Root': - continue - if sp in exclude_set: - continue - if sp in hint_by_path: - continue # skip duplicate section_path - summary = h.get('summary', '') or '' - hint_lines.append(f'▸ path="{sp}"') - if summary: - clipped = summary[:300] - hint_lines.append(f' {clipped}') - hint_by_path[sp] = h - - if not hint_lines: - return node - - items_text = '\n'.join(hint_lines) - - revision_context = '' - if revision_hint: - revision_context = ( - f'\nIMPORTANT: This is a REVISION round. ' - f'The previous search attempt failed because:\n' - f'"{revision_hint}"\n' - f'Adjust your selection accordingly. ' - f'If no candidate is relevant, return an EMPTY list [].\n' - ) - - prompt = _DISCOVERY_SELECT_PROMPT.format( - doc_name=doc_name or document_id, - budget_block=_format_budget_block(budget_snapshot), - items=items_text, - query=query, - revision_context=revision_context, - ) - response = await llm_fn(prompt) - # Parse {"selections": [...]} response — reuse action parser's extraction - parsed = _parse_action_response(response) - selections = parsed.get('selections', []) - - logger.info( - f' discovery_select_step doc="{doc_name}": ' - f'hints={len(hints)} selections={len(selections)}' - ) - - # 2. Hydrate selected paths - valid_selections = [s for s in selections if s['path'] in hint_by_path] - path_selections = [] - for sel in valid_selections: - path = sel['path'] - conf = sel.get('confidence', 0.7) - node.confidence[path] = conf - path_selections.append({'path': path, 'confidence': conf}) - - if path_selections: - chunks = await _hydrate_paths_to_rows( - db, - path_selections=path_selections, - user_id=user_id, - namespace=namespace, - document_id=document_id, - ) - if chunks: - connected = await hydrate_connected_target_rows( - db=db, - rows=chunks, - exclude_document_ids=[], - exclude_sections=[], - ) - if connected: - _owner_map = _build_connected_owner_map(chunks) - for c in connected: - if not c.get('owner_section_path'): - c['owner_section_path'] = _owner_map.get(str(c.get('chunk_id') or '')) - chunks = chunks + connected - - # Resolve Root-stranded assets to their true owner sections - _disc_job_result_id = next( - (str(c['job_result_id']) for c in chunks if c.get('job_result_id')), - None, - ) - _root_map = await _resolve_root_asset_owners( - db, - document_id=document_id, - job_result_id=_disc_job_result_id, - chunks=chunks, - ) if _disc_job_result_id else {} - if _root_map: - for c in chunks: - if c.get('owner_section_path'): - continue # already resolved by batch-level owner map - cid = str(c.get('chunk_id') or '') - if cid in _root_map: - c['owner_section_path'] = _root_map[cid] - - for chunk in chunks: - # Distribute chunk to its real path or fallback to the selection path - real_path = chunk.get('owner_section_path') or chunk.get('section_path') or chunk.get('source_chunk_path') - if real_path: - node.add_leaf_chunks(str(real_path), [chunk]) - - latency = int((time.monotonic() - t0) * 1000) - logger.info( - f' discovery_select_step done: hydrated={len(node.leaf_content)} ' - f'latency={latency}ms' - ) - return node - - except BudgetExceeded: - raise - except Exception as e: - logger.error(f' discovery_select_step failed for doc={document_id}: {e}') - return node + return await discovery_selection.discovery_select_step( + db, + document_id=document_id, + query=query, + llm_fn=llm_fn, + user_id=user_id, + namespace=namespace, + doc_name=doc_name, + discovery_hints=discovery_hints, + exclude_paths=exclude_paths, + revision_hint=revision_hint, + budget_snapshot=budget_snapshot, + ) diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py index 96c125e49..d8ec2ff22 100644 --- a/packages/shared-python/shared/services/retrieval/app_service.py +++ b/packages/shared-python/shared/services/retrieval/app_service.py @@ -1,966 +1,16 @@ from __future__ import annotations -import asyncio -import os -import re -import time from typing import Any -from loguru import logger -from sqlalchemy import and_, func, or_, select from sqlalchemy.ext.asyncio import AsyncSession -from shared.core.database import get_db_context -from shared.models.database.document import Document, DocumentChunk, DocumentSection, RetrievalHitStat -from shared.services.retrieval.graph_service import GraphQueryService, is_excluded_section -from shared.services.retrieval.lexical_text import normalize_section_path -from shared.services.retrieval.cache_service import get_cached_retrieval_query_result, set_cached_retrieval_query_result -from shared.services.retrieval.hit_stats_service import compute_importance_score, record_retrieval_hits -from shared.services.retrieval.channels import path_channel, content_channel, term_channel -from shared.services.storage.result_storage import get_result_storage -from shared.models.database.job_result import JobResult +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace +from shared.services.retrieval.execution.plan import ( + run_retrieval_query as execute_retrieval_query, +) +from shared.services.retrieval.search.scoring import merge_channels_rrf - -_MEDIA_CHUNK_TYPES = {'image', 'table'} - -_RRF_K = 60 -_CHANNEL_WEIGHT_PATH = 1.0 -_CHANNEL_WEIGHT_CONTENT = 2.0 -_CHANNEL_WEIGHT_TERM = 1.5 -_INTERNAL_RECALL_K_MULTIPLIER = 2 -_pending_retrieval_hit_stat_tasks: set[asyncio.Task[None]] = set() - -_DATA_TYPE_ALLOWED_CHUNK_TYPES: dict[int, set[str] | None] = { - 1: None, - 2: {'text'}, - 3: {'image'}, - 4: {'table'}, - 5: {'text', 'image'}, - 6: {'text', 'table'}, -} - - -def _resolve_allowed_chunk_types(data_type: int) -> set[str] | None: - return _DATA_TYPE_ALLOWED_CHUNK_TYPES.get(data_type) - - -_PATH_REF_RE = re.compile(r'\[(?:images|tables)/[^\]\n]+\]') - - -def _clean_content(content: str) -> str: - return _PATH_REF_RE.sub('', content).strip() - -_PUBLIC_RESULT_FIELDS = { - 'chunk_type', 'content', 'score', 'asset_url', -} - -_PUBLIC_SOURCE_FIELDS = { - 'document_id', 'source_file_name', 'section_path', -} - - -def _normalize_chunk_type(raw: str | None) -> str: - return str(raw or '').strip().split('\n', 1)[0].lower() - - -def _filter_excluded_rows( - rows: list[dict[str, Any]], - *, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], -) -> list[dict[str, Any]]: - filtered: list[dict[str, Any]] = [] - excluded_documents = set(exclude_document_ids) - for row in rows: - document_id = row.get('document_id') - if document_id in excluded_documents: - continue - if is_excluded_section( - document_id=document_id, - section_path=row.get('section_path'), - exclude_sections=exclude_sections, - ): - continue - filtered.append(row) - return filtered - - -def _iter_connected_target_ids(row: dict[str, Any]) -> list[str]: - metadata = row.get('chunk_metadata') or {} - if not isinstance(metadata, dict): - return [] - - target_ids: list[str] = [] - for item in metadata.get('connect_to') or []: - if not isinstance(item, dict): - continue - target_id = str(item.get('target') or '').strip() - if target_id: - target_ids.append(target_id) - return target_ids - - -async def hydrate_connected_target_rows( - *, - db: AsyncSession | None, - rows: list[dict[str, Any]], - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], -) -> list[dict[str, Any]]: - if db is None: - return [] - - existing_chunk_ids = { - str(row.get('chunk_id') or '').strip() - for row in rows - if row.get('chunk_id') - } - target_ids_by_revision: dict[tuple[str, str], set[str]] = {} - for row in rows: - if _normalize_chunk_type(row.get('chunk_type')) != 'text': - continue - document_id = str(row.get('document_id') or '').strip() - job_result_id = str(row.get('job_result_id') or '').strip() - if not document_id or not job_result_id: - continue - for target_id in _iter_connected_target_ids(row): - if target_id in existing_chunk_ids: - continue - target_ids_by_revision.setdefault((document_id, job_result_id), set()).add(target_id) - - if not target_ids_by_revision: - return [] - - revision_filters = [ - and_( - DocumentChunk.document_id == document_id, - DocumentChunk.job_result_id == job_result_id, - DocumentChunk.chunk_id.in_(sorted(target_ids)), - ) - for (document_id, job_result_id), target_ids in target_ids_by_revision.items() - if target_ids - ] - if not revision_filters: - return [] - - stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, DocumentChunk.document_id == Document.document_id) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(or_(*revision_filters)) - .order_by(DocumentChunk.sort_order) - ) - result = await db.execute(stmt) - - hydrated_rows: list[dict[str, Any]] = [] - for document, chunk, section, job_result in result.all(): - section_path = section.section_path if section else None - hydrated_rows.append( - { - 'document_id': document.document_id, - 'chunk_id': chunk.chunk_id, - 'section_id': chunk.section_id, - 'section_path': section_path, - 'source_file_name': document.source_file_name, - 'chunk_type': chunk.chunk_type, - 'content': chunk.content, - 'score': 0.0, - 'file_path': chunk.file_path, - 'chunk_metadata': chunk.chunk_metadata or {}, - 'job_result_id': chunk.job_result_id, - 'job_id': job_result.job_id if job_result else None, - 'sort_order': chunk.sort_order, - } - ) - - return _filter_excluded_rows( - hydrated_rows, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - - -async def assemble_retrieval_results( - *, - db: AsyncSession | None = None, - rows: list[dict[str, Any]], - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - allowed_chunk_types: set[str] | None = None, -) -> list[dict[str, Any]]: - filtered_rows = _filter_excluded_rows( - rows, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - if allowed_chunk_types is not None: - filtered_rows = [ - row for row in filtered_rows - if _normalize_chunk_type(row.get('chunk_type')) in allowed_chunk_types - ] - hydrated_rows = await hydrate_connected_target_rows( - db=db, - rows=filtered_rows, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - rows_by_chunk_id = { - str(row.get('chunk_id') or ''): row - for row in [*filtered_rows, *hydrated_rows] - if row.get('chunk_id') - } - - embedded_targets: set[str] = set() - for row in filtered_rows: - for target_id in _iter_connected_target_ids(row): - if target_id in rows_by_chunk_id: - embedded_targets.add(target_id) - - assembled: list[dict[str, Any]] = [] - for row in filtered_rows: - if row.get('chunk_id') in embedded_targets: - continue - metadata = row.get('chunk_metadata') or {} - if not isinstance(metadata, dict): - metadata = {} - assembled_row = dict(row) - base_content = str(row.get('content') or '') - if _normalize_chunk_type(row.get('chunk_type')) == 'text': - connected_targets: list[tuple[int, str]] = [] - for target_id in _iter_connected_target_ids(row): - target_row = rows_by_chunk_id.get(target_id) - if not target_row: - continue - if _normalize_chunk_type(target_row.get('chunk_type')) != 'table': - continue - target_content = str(target_row.get('content') or '').strip() - if target_content: - sort_key = int(target_row.get('sort_order', 0) or 0) - connected_targets.append((sort_key, target_content)) - connected_targets.sort(key=lambda x: x[0]) - related_parts = [content for _, content in connected_targets] - if base_content and related_parts: - assembled_row['content'] = '\n\n'.join([base_content, *related_parts]) - else: - assembled_row['content'] = base_content - else: - assembled_row['content'] = base_content - assembled_row['content'] = _clean_content(assembled_row['content']) - assembled.append(assembled_row) - return assembled - - - - - -def _merge_same_section_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: - if not rows: - return rows - groups: dict[str, list[dict[str, Any]]] = {} - order: list[str] = [] - for row in rows: - sp = row.get('section_path') - if sp: - key = f"{row.get('document_id', '')}::{sp}" - else: - key = row.get('chunk_id', '') - if key not in groups: - groups[key] = [] - order.append(key) - groups[key].append(row) - - merged: list[dict[str, Any]] = [] - for key in order: - group = groups[key] - if len(group) == 1: - merged.append(group[0]) - continue - base = dict(group[0]) - base['content'] = '\n'.join(str(r.get('content', '')) for r in group) - base['score'] = max(r.get('score', 0.0) for r in group) - merged.append(base) - return merged - - -def merge_channels_rrf( - channels: list[list[dict[str, Any]]], - weights: list[float], - top_k: int, - k: int = _RRF_K, -) -> list[dict[str, Any]]: - """Reciprocal Rank Fusion across multiple retrieval channels.""" - score_dict: dict[str, float] = {} - row_by_chunk_id: dict[str, dict[str, Any]] = {} - - for channel_idx, channel_rows in enumerate(channels): - w = weights[channel_idx] if channel_idx < len(weights) else 1.0 - for rank, row in enumerate(channel_rows): - chunk_id = str(row.get('chunk_id') or '') - if not chunk_id: - continue - rrf_score = w / (k + rank + 1) - score_dict[chunk_id] = score_dict.get(chunk_id, 0.0) + rrf_score - if chunk_id not in row_by_chunk_id: - row_by_chunk_id[chunk_id] = row - - ranked = sorted(score_dict.items(), key=lambda x: x[1], reverse=True) - results: list[dict[str, Any]] = [] - for chunk_id, fused_score in ranked[:top_k]: - row = row_by_chunk_id[chunk_id] - results.append(dict(row, score=round(fused_score, 6))) - return results - - -async def list_graph_routed_chunks( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], -) -> list[dict[str, Any]]: - service = GraphQueryService() - entry_document_ids = await service.find_entry_documents( - db, - user_id=user_id, - namespace=namespace, - query=query, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - return await service.collect_candidate_chunks( - db, - user_id=user_id, - namespace=namespace, - entry_document_ids=entry_document_ids, - query=query, - top_k=top_k * _INTERNAL_RECALL_K_MULTIPLIER, - exclude_sections=exclude_sections, - ) - - -def _finalize_retrieval_hit_stats_task(task: asyncio.Task[None]) -> None: - _pending_retrieval_hit_stat_tasks.discard(task) - - try: - task.result() - except asyncio.CancelledError: - pass - except Exception as e: - logger.warning(f'Failed to record retrieval hit stats (ignored): {e}') - - -def schedule_retrieval_hit_stats_update(*, user_id: str, namespace: str, results: list[dict[str, Any]]) -> None: - try: - task = asyncio.create_task( - _record_retrieval_hit_stats_best_effort( - user_id=user_id, - namespace=namespace, - results=results, - ), - name=f'retrieval_hit_stats:{user_id}:{namespace}', - ) - _pending_retrieval_hit_stat_tasks.add(task) - task.add_done_callback(_finalize_retrieval_hit_stats_task) - except Exception as e: - logger.warning(f'Failed to schedule retrieval hit stats update (ignored): {e}') - - -async def drain_retrieval_hit_stats_updates(timeout_seconds: float = 2.0) -> None: - if not _pending_retrieval_hit_stat_tasks: - return - - pending_tasks = tuple(_pending_retrieval_hit_stat_tasks) - - try: - await asyncio.wait_for( - asyncio.gather(*pending_tasks, return_exceptions=True), - timeout=timeout_seconds, - ) - except asyncio.TimeoutError: - for task in pending_tasks: - if not task.done(): - task.cancel() - - await asyncio.gather(*pending_tasks, return_exceptions=True) - - -async def _record_retrieval_hit_stats_best_effort(*, user_id: str, namespace: str, results: list[dict[str, Any]]) -> None: - try: - async with get_db_context() as db: - await record_retrieval_hits(db, user_id=user_id, namespace=namespace, results=results) - await db.commit() - except Exception as e: - logger.warning(f'Failed to record retrieval hit stats (ignored): {e}') - - -def _with_citation(row: dict[str, Any]) -> dict[str, Any]: - citation = { - 'document_id': row.get('document_id'), - 'chunk_id': row.get('chunk_id'), - 'source_file_name': row.get('source_file_name'), - 'section_path': row.get('section_path'), - } - return {**row, 'citation': citation} - - -def _to_public_source(row: dict[str, Any]) -> dict[str, Any]: - return {field: row.get(field) for field in _PUBLIC_SOURCE_FIELDS} - - -def _is_media_chunk(row: dict[str, Any]) -> bool: - return _normalize_chunk_type(row.get('chunk_type')) in _MEDIA_CHUNK_TYPES - - -async def generate_retrieval_asset_url(*, job_id: str, artifact_ref: str) -> str | None: - return get_result_storage().generate_artifact_url(job_id=job_id, artifact_ref=artifact_ref) - - -def _is_client_result_artifact_ref(asset_ref: str | None) -> bool: - return get_result_storage().normalize_artifact_ref(asset_ref) is not None - - -async def _to_public_response(response: dict[str, Any]) -> dict[str, Any]: - public_response = { - 'namespace': response.get('namespace'), - 'query': response.get('query'), - 'router_used': response.get('router_used'), - 'results': [], - } - - # Forward agentic evidence fields when present - if response.get('answer_text') is not None: - public_response['answer_text'] = response['answer_text'] - if response.get('referenced_chunks') is not None: - public_response['referenced_chunks'] = response['referenced_chunks'] - - public_results: list[dict[str, Any]] = [] - for row in response.get('results', []): - artifact_ref = row.get('file_path') - asset_url = None - if _is_media_chunk(row) and _is_client_result_artifact_ref(artifact_ref) and row.get('job_id'): - try: - asset_url = await generate_retrieval_asset_url( - job_id=str(row['job_id']), - artifact_ref=str(artifact_ref), - ) - except Exception as e: - logger.warning(f'Failed to generate retrieval asset URL (ignored): {e}') - - public_row: dict[str, Any] = {} - for field in _PUBLIC_RESULT_FIELDS: - if field == 'asset_url': - if asset_url: - public_row['asset_url'] = asset_url - elif field in row: - public_row[field] = row[field] - if 'source' in row: - public_row['source'] = row['source'] - else: - public_row['source'] = _to_public_source(row) - public_results.append(public_row) - - public_response['results'] = public_results - return public_response - - -def _get_row_path(row: dict[str, Any]) -> str: - """Extract the canonical path from a row for deduplication.""" - return str(row.get('section_path') or row.get('source_chunk_path') or '') - - -def _get_candidate_key(row: dict[str, Any]) -> str: - path = _get_row_path(row) - if path: - return f'path:{path}' - chunk_id = str(row.get('chunk_id') or '').strip() - return f'chunk:{chunk_id}' if chunk_id else '' - - -def _normalize_row_scores( - rows: list[dict[str, Any]], - *, - source_field: str, - target_field: str, - default: float, -) -> None: - if not rows: - return - values = [float(row.get(source_field, 0.0) or 0.0) for row in rows] - min_score = min(values) - max_score = max(values) - if max_score <= 0.0 and min_score <= 0.0: - for row in rows: - row[target_field] = 0.0 - return - if max_score == min_score: - for row in rows: - row[target_field] = default - return - denominator = max_score - min_score - for row in rows: - raw_score = float(row.get(source_field, 0.0) or 0.0) - row[target_field] = round((raw_score - min_score) / denominator, 6) - - -def _importance_multiplier( - rows: list[dict[str, Any]], - *, - raw_field: str = 'importance_raw_score', - low: float = 0.1, - high: float = 2.0, -) -> None: - """Apply adaptive sigmoid-based importance boost to agent/discovery scores. - - Uses median of ``raw_field`` as center and IQR as spread so the curve - adapts to any KB size without hard-coded thresholds. When all values - are identical (IQR ≈ 0) the multiplier is 1.0 (neutral). - - Output range ``[low, high]`` — default [0.1, 2.0] — is the only - configured constant: max 2× boost, min 10%. The function modifies - ``agent_score`` and ``discovery_score`` **in place**. - """ - import math - - if not rows: - return - - values = sorted(float(r.get(raw_field, 0.0) or 0.0) for r in rows) - n = len(values) - median = values[n // 2] if n % 2 else (values[n // 2 - 1] + values[n // 2]) / 2 - q1 = values[n // 4] if n >= 4 else values[0] - q3 = values[3 * n // 4] if n >= 4 else values[-1] - iqr = q3 - q1 - - for row in rows: - raw = float(row.get(raw_field, 0.0) or 0.0) - if iqr <= 1e-9: - mult = 1.0 - else: - z = (raw - median) / iqr - s = 1.0 / (1.0 + math.exp(-z)) - mult = low + (high - low) * s - row['importance_multiplier'] = round(mult, 4) - row['agent_score'] = round( - float(row.get('agent_score', 0.0) or 0.0) * mult, 6, - ) - row['discovery_score'] = round( - float(row.get('discovery_score', 0.0) or 0.0) * mult, 6, - ) - - -async def _load_chunk_importance_scores( - db: AsyncSession, - *, - user_id: str, - namespace: str, - rows: list[dict[str, Any]], -) -> dict[str, float]: - chunk_ids = sorted({ - str(row.get('chunk_id') or '').strip() - for row in rows - if row.get('chunk_id') - }) - if not chunk_ids: - return {} - stmt = ( - select( - RetrievalHitStat.chunk_id, - RetrievalHitStat.hit_count, - RetrievalHitStat.last_hit_at, - RetrievalHitStat.created_at, - ) - .where(RetrievalHitStat.user_id == user_id) - .where(RetrievalHitStat.namespace == namespace) - .where(RetrievalHitStat.hit_kind == 'chunk') - .where(RetrievalHitStat.chunk_id.in_(chunk_ids)) - ) - result = await db.execute(stmt) - importance_scores: dict[str, float] = {} - for chunk_id, hit_count, last_hit_at, created_at in result.all(): - if not chunk_id: - continue - importance_scores[str(chunk_id)] = compute_importance_score(hit_count, last_hit_at, created_at) - return importance_scores - - -def _rank_candidates_by_path( - discovery_rows: list[dict[str, Any]], - routed_rows: list[dict[str, Any]], - top_k: int, -) -> list[dict[str, Any]]: - """Rank discovery and routed candidates in one comparable path space.""" - merged: dict[str, dict[str, Any]] = {} - insertion_order: dict[str, int] = {} - counter = 0 - - for row in discovery_rows: - key = _get_candidate_key(row) - if not key: - continue - candidate = dict(row) - candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0) - candidate['agent_score'] = 0.0 - candidate.setdefault('hydrate_mode', 'chunks') - merged[key] = candidate - insertion_order[key] = counter - counter += 1 - - for row in routed_rows: - key = _get_candidate_key(row) - if not key: - continue - routed_agent_score = float(row.get('agent_score', 0.0) or 0.0) - if key not in merged: - candidate = dict(row) - candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0) - candidate['agent_score'] = routed_agent_score - merged[key] = candidate - insertion_order[key] = counter - counter += 1 - continue - candidate = merged[key] - candidate['agent_score'] = max(float(candidate.get('agent_score', 0.0) or 0.0), routed_agent_score) - if not candidate.get('source_chunk_path') and row.get('source_chunk_path'): - candidate['source_chunk_path'] = row.get('source_chunk_path') - if not candidate.get('section_path') and row.get('section_path'): - candidate['section_path'] = row.get('section_path') - - # ── Dual-priority ranking ──────────────────────────────────────────── - # When the agent produced results (routed_rows non-empty), rows with - # agent_score=0 are demoted to a fallback pool. Primary sort is by - # agent_score (includes importance boost from _importance_multiplier), - # with discovery_score as tiebreaker. - has_agent_results = len(routed_rows) > 0 - - primary_rows: list[dict[str, Any]] = [] - fallback_rows: list[dict[str, Any]] = [] - - for key, row in merged.items(): - agent_score = float(row.get('agent_score', 0.0) or 0.0) - discovery_score = float(row.get('discovery_score', 0.0) or 0.0) - row['evidence_score'] = round(agent_score if has_agent_results else max(discovery_score, agent_score), 6) - row['score'] = row['evidence_score'] - row['_candidate_order'] = insertion_order[key] - - if has_agent_results and agent_score <= 0.0: - fallback_rows.append(row) - else: - primary_rows.append(row) - - def _sort_key(row): - return ( - float(row.get('agent_score', 0.0) or 0.0), - float(row.get('discovery_score', 0.0) or 0.0), - -int(row.get('_candidate_order', 0) or 0), - ) - - primary_rows.sort(key=_sort_key, reverse=True) - ranked_rows = primary_rows[:top_k] - - # Back-fill from fallback if primary results are insufficient - if len(ranked_rows) < top_k and fallback_rows: - fallback_rows.sort(key=_sort_key, reverse=True) - ranked_rows.extend(fallback_rows[:top_k - len(ranked_rows)]) - - for row in ranked_rows: - row.pop('_candidate_order', None) - return ranked_rows - - -async def _count_scoped_chunks( - db: AsyncSession, - *, - user_id: str, - namespace: str, - exclude_document_ids: list[str], - allowed_chunk_types: set[str] | None, -) -> int: - stmt = ( - select(func.count(DocumentChunk.id)) - .join(Document, (Document.document_id == DocumentChunk.document_id) & (Document.current_job_result_id == DocumentChunk.job_result_id)) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - ) - if exclude_document_ids: - stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) - if allowed_chunk_types is not None: - stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types))) - result = await db.execute(stmt) - return result.scalar() or 0 - - -async def _load_all_scoped_chunks( - db: AsyncSession, - *, - user_id: str, - namespace: str, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - allowed_chunk_types: set[str] | None, - signal_paths: list[str], - filter_mode: str, -) -> list[dict[str, Any]]: - stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) & (DocumentChunk.job_result_id == Document.current_job_result_id)) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .order_by(DocumentChunk.sort_order) - ) - if exclude_document_ids: - stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) - if allowed_chunk_types is not None: - stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types))) - - result = await db.execute(stmt) - rows: list[dict[str, Any]] = [] - for document, chunk, section, job_result in result.all(): - section_path = section.section_path if section else None - if is_excluded_section(document_id=document.document_id, section_path=section_path, exclude_sections=exclude_sections): - continue - if signal_paths and section_path: - path_lower = section_path.lower() - matches_any = any(kw.lower() in path_lower for kw in signal_paths) - if filter_mode == 'keep' and not matches_any: - continue - if filter_mode == 'delete' and matches_any: - continue - rows.append({ - 'document_id': document.document_id, - 'chunk_id': chunk.chunk_id, - 'section_id': chunk.section_id, - 'section_path': section_path, - 'source_file_name': document.source_file_name, - 'chunk_type': chunk.chunk_type, - 'content': chunk.content, - 'score': 1.0, - 'file_path': chunk.file_path, - 'chunk_metadata': chunk.chunk_metadata or {}, - 'job_result_id': chunk.job_result_id, - 'job_id': job_result.job_id if job_result else None, - 'sort_order': chunk.sort_order, - }) - return rows - - -async def _hydrate_paths_to_rows( - db: AsyncSession, - *, - path_selections: list[dict[str, Any]], - user_id: str, - namespace: str, - document_id: str | None = None, -) -> list[dict[str, Any]]: - """Load full chunk rows by section_path or source_chunk_path. - - When *document_id* is provided the query is scoped to that single - document, preventing cross-document collisions on generic paths - such as ``Root``. - - Supports hydrate_mode branching: - - 'chunks' (default): all chunk types under the section subtree - - 'outline': synthetic row from section metadata, no real chunks - - 'assets_only': only image + table chunks - - 'image_only': only image chunks - - 'table_only': only table chunks - """ - if not path_selections: - return [] - - # Group selections by hydrate_mode - confidence_by_path: dict[str, float] = {} - mode_by_path: dict[str, str] = {} - ordered_paths: list[str] = [] - for item in path_selections: - raw_path = str(item.get('path') or '').strip() - path = normalize_section_path(raw_path) if raw_path and '/' in raw_path else raw_path - if not path: - continue - confidence = float(item.get('confidence', 0.0) or 0.0) - hydrate_mode = str(item.get('hydrate_mode') or 'chunks').strip().lower() - if path not in confidence_by_path: - ordered_paths.append(path) - confidence_by_path[path] = confidence - mode_by_path[path] = hydrate_mode - else: - confidence_by_path[path] = max(confidence_by_path[path], confidence) - if not ordered_paths: - return [] - - # Separate outline paths from chunk-loading paths - outline_paths = [p for p in ordered_paths if mode_by_path.get(p) == 'outline'] - chunk_paths = [p for p in ordered_paths if mode_by_path.get(p) != 'outline'] - - rows: list[dict[str, Any]] = [] - - # ── Outline mode: synthesize rows from section metadata ────────────── - if outline_paths: - outline_section_filters = [] - for path in outline_paths: - outline_section_filters.append(DocumentSection.section_path == path) - - outline_stmt = ( - select(Document, DocumentSection) - .join(DocumentSection, (DocumentSection.document_id == Document.document_id) - & (DocumentSection.job_result_id == Document.current_job_result_id)) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where(or_(*outline_section_filters)) - ) - if document_id: - outline_stmt = outline_stmt.where(Document.document_id == document_id) - outline_result = await db.execute(outline_stmt) - for document, section in outline_result.all(): - agent_score = confidence_by_path.get(section.section_path, 0.0) - summary_text = (section.summary or '').strip() - title_text = (section.section_title or '').strip() - content = f'[Outline] {title_text}' - if summary_text: - content += f'\n{summary_text}' - rows.append({ - 'document_id': document.document_id, - 'chunk_id': f'outline_{section.section_id}', - 'section_id': section.section_id, - 'section_path': section.section_path, - 'source_file_name': document.source_file_name, - 'chunk_type': 'outline', - 'content': content, - 'score': agent_score, - 'agent_score': agent_score, - 'file_path': None, - 'chunk_metadata': {}, - 'job_result_id': section.job_result_id, - 'job_id': None, - 'source_chunk_path': None, - 'sort_order': section.sort_order, - 'hydrate_mode': 'outline', - }) - - # ── Chunk modes: load real chunks with optional type filters ───────── - if chunk_paths: - section_path_filters = [] - # Separate self_only paths (exact match only, no descendant LIKE) - # from regular chunk paths (exact + descendant subtree match) - self_only_paths = {p for p in chunk_paths if mode_by_path.get(p) == 'self_only'} - for path in chunk_paths: - section_path_filters.append(DocumentSection.section_path == path) - if path not in self_only_paths: - section_path_filters.append(DocumentSection.section_path.like(f'{path} / %')) - - stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id)) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where( - or_( - *section_path_filters, - DocumentChunk.source_chunk_path.in_(chunk_paths), - ) - ) - ) - if document_id: - stmt = stmt.where(Document.document_id == document_id) - result = await db.execute(stmt) - - # Build a map of path → allowed chunk_types based on hydrate_mode - _MODE_ALLOWED_TYPES: dict[str, set[str] | None] = { - 'chunks': None, # all types - 'self_only': None, # all types, but without descendant filtering - 'assets_only': {'image', 'table'}, - 'image_only': {'image'}, - 'table_only': {'table'}, - } - - seen_paths: set[str] = set() - for document, chunk, section, job_result in result.all(): - row_path = (section.section_path if section else None) or chunk.source_chunk_path or '' - if row_path in seen_paths: - continue - - # Find which ordered path this row belongs to - matched_path = row_path - if section and section.section_path not in confidence_by_path: - matched_path = next( - ( - path for path in chunk_paths - if section.section_path == path or section.section_path.startswith(f'{path} / ') - ), - row_path, - ) - - # Check chunk_type filter based on hydrate_mode - path_mode = mode_by_path.get(matched_path, 'chunks') - allowed_types = _MODE_ALLOWED_TYPES.get(path_mode) - if allowed_types is not None: - chunk_type_lower = (chunk.chunk_type or '').strip().lower() - if chunk_type_lower not in allowed_types: - continue - - seen_paths.add(row_path) - agent_score = confidence_by_path.get(matched_path, 0.0) - rows.append({ - 'document_id': document.document_id, - 'chunk_id': chunk.chunk_id, - 'section_id': chunk.section_id, - 'section_path': section.section_path if section else None, - 'source_file_name': document.source_file_name, - 'chunk_type': chunk.chunk_type, - 'content': chunk.content, - 'score': agent_score, - 'agent_score': agent_score, - 'file_path': chunk.file_path, - 'chunk_metadata': chunk.chunk_metadata or {}, - 'job_result_id': chunk.job_result_id, - 'job_id': job_result.job_id if job_result else None, - 'source_chunk_path': chunk.source_chunk_path, - 'sort_order': chunk.sort_order, - 'hydrate_mode': path_mode, - }) - - # ── Sort by agent-selected order ───────────────────────────────────── - path_order = {p: idx for idx, p in enumerate(ordered_paths)} - - def _row_sort_key(row: dict[str, Any]) -> int: - row_path = _get_row_path(row) - if row_path in path_order: - return path_order[row_path] - for path, idx in path_order.items(): - if row_path.startswith(f'{path} / '): - return idx - return 10**9 - - rows.sort(key=_row_sort_key) - hydrated_paths = {_get_row_path(r) for r in rows} - resolved_inputs = { - path for path in ordered_paths - if path in hydrated_paths or any(row_path.startswith(f'{path} / ') for row_path in hydrated_paths) - } - # Outline paths are always resolved (synthesized) - resolved_inputs |= set(outline_paths) - missed = len(ordered_paths) - len(resolved_inputs) - if missed > 0: - missing_paths = [p for p in ordered_paths if p not in resolved_inputs] - logger.warning( - f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved (missed={missed}); ' - f'missing[:5]={missing_paths[:5]}' - ) - else: - logger.info(f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved') - return rows +__all__ = ["merge_channels_rrf", "run_retrieval_query"] async def run_retrieval_query( @@ -974,7 +24,7 @@ async def run_retrieval_query( exclude_sections: list[dict[str, str]], data_type: int = 1, signal_paths: list[str] | None = None, - filter_mode: str = 'delete', + filter_mode: str = "delete", channels: list[str] | None = None, channel_weights: dict[str, float] | None = None, rerank: bool = False, @@ -982,30 +32,14 @@ async def run_retrieval_query( internal_recall_k: int | None = None, use_agentic: bool | None = None, ) -> dict[str, Any]: - """Checkerboard retrieval: 3 independent channels -> RRF -> agent/graph union -> assembly.""" - t_start = time.monotonic() - query = query.strip() - logger.info('\n' + '█' * 70) - logger.info(' 🚀 RETRIEVAL PIPELINE START') - logger.info(f' query="{query}"') - logger.info(f' user={user_id} ns={namespace} top_k={top_k} data_type={data_type}') - logger.info(f' exclude_docs={exclude_document_ids} exclude_secs={len(exclude_sections)}') - logger.info('█' * 70) - - if not query: - logger.info(' ⛔ Empty query filtered, skipping retrieval pipeline') - return { - "namespace": namespace, - "query": query, - "router_used": "empty_query_filtered", - "results": [], - } - - allowed_chunk_types = _resolve_allowed_chunk_types(data_type) - effective_recall_k = internal_recall_k if internal_recall_k is not None else top_k * _INTERNAL_RECALL_K_MULTIPLIER - logger.info(f' allowed_chunk_types={allowed_chunk_types} effective_recall_k={effective_recall_k} signal_paths={signal_paths} filter_mode={filter_mode} rerank={rerank} threshold={threshold}') - - cache_extra = dict( + return await execute_retrieval_query( + db=db, + user_id=user_id, + namespace=normalize_retrieval_namespace(namespace), + query=query, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, data_type=data_type, signal_paths=signal_paths, filter_mode=filter_mode, @@ -1014,383 +48,5 @@ async def run_retrieval_query( rerank=rerank, threshold=threshold, internal_recall_k=internal_recall_k, - # Always True: agentic mode now always routes through workflow - decomposition_enabled=True, - ) - - cache_version: int | None = None - try: - cache_version, cached = await get_cached_retrieval_query_result( - user_id=user_id, - namespace=namespace, - query=query, - top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - **cache_extra, - ) - if cached: - logger.info(f'retrieval: cache_hit=True version={cache_version}') - try: - schedule_retrieval_hit_stats_update( - user_id=user_id, - namespace=namespace, - results=cached.get("results", []), - ) - except Exception as e: - logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") - return await _to_public_response(cached) - except Exception as e: - logger.warning(f"Failed to read retrieval cache (ignored): {e}") - - logger.debug(f' 📦 Cache miss (version={cache_version}), running full pipeline') - - # ── Small KB optimization ── - try: - total_chunk_count = await _count_scoped_chunks( - db, user_id=user_id, namespace=namespace, - exclude_document_ids=exclude_document_ids, - allowed_chunk_types=allowed_chunk_types, - ) - except Exception as e: - logger.warning(f"Failed to count scoped chunks, skipping small KB optimization: {e}") - total_chunk_count = top_k + 1 - logger.info(f'\n 📊 Total chunks in scope: {total_chunk_count}') - if total_chunk_count <= top_k: - logger.info(f' ⚡ Small KB optimization: {total_chunk_count} chunks <= top_k={top_k}, returning all') - all_rows = await _load_all_scoped_chunks( - db, user_id=user_id, namespace=namespace, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths or [], - filter_mode=filter_mode, - ) - logger.info(f' small_kb load: loaded={len(all_rows)} rows after signal/exclude filters') - assembled_rows = await assemble_retrieval_results( - db=db, rows=all_rows, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - ) - results = [_with_citation(row) for row in assembled_rows] - response = { - "namespace": namespace, "query": query, - "router_used": "small_kb_all", "results": results, - } - if cache_version is not None: - try: - await set_cached_retrieval_query_result( - user_id=user_id, namespace=namespace, version=cache_version, - query=query, top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - response=response, **cache_extra, - ) - except Exception as e: - logger.warning(f"Failed to write retrieval cache (ignored): {e}") - try: - schedule_retrieval_hit_stats_update(user_id=user_id, namespace=namespace, results=results) - except Exception as e: - logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") - elapsed_total = round((time.monotonic() - t_start) * 1000) - logger.info(f' ✅ Small KB: {len(results)} results in {elapsed_total}ms') - return await _to_public_response(response) - - # ══ Route: agentic (unified workflow) vs legacy ══ - if use_agentic is not None: - _agentic_enabled = use_agentic - else: - _agentic_enabled = os.environ.get('RETRIEVAL_AGENTIC_ENABLED', 'true') == 'true' - if _agentic_enabled: - # ── Unified agentic path via WorkflowOrchestrator ── - # Simple queries: planner returns a single-step plan (no decomposition). - # Complex queries: planner returns a multi-step plan with synthesize. - # Both go through the same code path. - from shared.services.retrieval.workflow.orchestrator import WorkflowOrchestrator - - workflow = WorkflowOrchestrator() - workflow_result = await workflow.run( - db, - user_id=user_id, - namespace=namespace, - query=query, - top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - data_type=data_type, - signal_paths=signal_paths, - filter_mode=filter_mode, - channels=channels, - channel_weights=channel_weights, - ) - - # Enrich referenced_chunks with asset URLs (images/tables) - enriched_refs: list[dict[str, Any]] = [] - for ref in workflow_result.referenced_chunks: - enriched = dict(ref) - chunk_type = _normalize_chunk_type(ref.get('chunk_type')) - artifact_ref = ref.get('file_path', '') - job_id = ref.get('job_id', '') - if chunk_type in _MEDIA_CHUNK_TYPES and job_id and _is_client_result_artifact_ref(artifact_ref): - try: - asset_url = await generate_retrieval_asset_url( - job_id=str(job_id), artifact_ref=str(artifact_ref), - ) - if asset_url: - enriched['asset_url'] = asset_url - except Exception as e: - logger.warning(f'Failed to generate agentic asset URL (ignored): {e}') - enriched_refs.append(enriched) - - response = workflow_result.to_api_response() - # Override referenced_chunks with enriched versions - response['referenced_chunks'] = enriched_refs - - if cache_version is not None: - try: - await set_cached_retrieval_query_result( - user_id=user_id, namespace=namespace, version=cache_version, - query=query, top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - response=response, **cache_extra, - ) - except Exception as e: - logger.warning(f"Failed to write retrieval cache (ignored): {e}") - - try: - schedule_retrieval_hit_stats_update( - user_id=user_id, namespace=namespace, - results=enriched_refs, - ) - except Exception as e: - logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") - - elapsed_total = round((time.monotonic() - t_start) * 1000) - logger.info( - f'\n{"█" * 70}\n' - f' ✅ AGENTIC RETRIEVAL COMPLETE: ' - f'{len(enriched_refs)} chunks | ' - f'answer={len(workflow_result.answer_text)} chars | ' - f'router={workflow_result.router_used} | {elapsed_total}ms\n' - f'{"█" * 70}' - ) - - return await _to_public_response(response) - - else: - - # ── LEGACY path (existing code, unchanged) ── - - # ── Channel execution ── - active_channels = set(channels) if channels else {'path', 'content', 'term'} - logger.info(f'\n 📡 PHASE 1: Bottom-Layer Discovery (channels={sorted(active_channels)})') - logger.info(f' effective_recall_k={effective_recall_k}') - - path_rows: list[dict[str, Any]] = [] - content_rows: list[dict[str, Any]] = [] - term_rows: list[dict[str, Any]] = [] - - if 'path' in active_channels: - t_ch = time.monotonic() - path_rows = await path_channel( - db, user_id=user_id, namespace=namespace, query=query, - top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, filter_mode=filter_mode, - ) - elapsed_ch = round((time.monotonic() - t_ch) * 1000) - logger.info(f'\n 📡 path_channel: {len(path_rows)} rows in {elapsed_ch}ms') - for i, r in enumerate(path_rows[:5]): - logger.info(f' [{i}] score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")} type={r.get("chunk_type","?")}') - if len(path_rows) > 5: - logger.info(f' ... and {len(path_rows) - 5} more') - - if 'content' in active_channels: - t_ch = time.monotonic() - content_rows = await content_channel( - db, user_id=user_id, namespace=namespace, query=query, - top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, filter_mode=filter_mode, - ) - elapsed_ch = round((time.monotonic() - t_ch) * 1000) - logger.info(f'\n 📡 content_channel: {len(content_rows)} rows in {elapsed_ch}ms') - for i, r in enumerate(content_rows[:5]): - logger.info(f' [{i}] score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")} content={str(r.get("content",""))[:80]}') - if len(content_rows) > 5: - logger.info(f' ... and {len(content_rows) - 5} more') - - if 'term' in active_channels: - t_ch = time.monotonic() - term_rows = await term_channel( - db, user_id=user_id, namespace=namespace, query=query, - top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, filter_mode=filter_mode, - ) - elapsed_ch = round((time.monotonic() - t_ch) * 1000) - logger.info(f'\n 📡 term_channel: {len(term_rows)} rows in {elapsed_ch}ms') - for i, r in enumerate(term_rows[:5]): - logger.info(f' [{i}] score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")} type={r.get("chunk_type","?")}') - if len(term_rows) > 5: - logger.info(f' ... and {len(term_rows) - 5} more') - - # ── RRF fusion with configurable weights ── - default_weights = { - 'path': _CHANNEL_WEIGHT_PATH, - 'content': _CHANNEL_WEIGHT_CONTENT, - 'term': _CHANNEL_WEIGHT_TERM, - } - effective_weights = {**default_weights, **(channel_weights or {})} - - channel_lists: list[list[dict[str, Any]]] = [] - weight_list: list[float] = [] - - if path_rows: - channel_lists.append(path_rows) - weight_list.append(effective_weights.get('path', _CHANNEL_WEIGHT_PATH)) - if content_rows: - channel_lists.append(content_rows) - weight_list.append(effective_weights.get('content', _CHANNEL_WEIGHT_CONTENT)) - if term_rows: - channel_lists.append(term_rows) - weight_list.append(effective_weights.get('term', _CHANNEL_WEIGHT_TERM)) - - fused_rows = merge_channels_rrf(channel_lists, weight_list, top_k) if channel_lists else [] - logger.info(f'\n 🔀 RRF Fusion: {len(fused_rows)} rows from {len(channel_lists)} channels (weights={dict(zip(["path","content","term"][:len(weight_list)], weight_list))})') - for i, r in enumerate(fused_rows[:5]): - logger.info(f' [{i}] rrf_score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")}') - if len(fused_rows) > 5: - logger.info(f' ... and {len(fused_rows) - 5} more') - - # ── Section merging ── - pre_merge = len(fused_rows) - fused_rows = _merge_same_section_rows(fused_rows) - if len(fused_rows) != pre_merge: - logger.info(f'retrieval: section_merge={pre_merge}->{len(fused_rows)}') - - # ── Threshold filtering ── - if threshold > 0.0 and fused_rows: - pre_count = len(fused_rows) - fused_rows = [row for row in fused_rows if row.get('score', 0.0) >= threshold] - logger.info(f'retrieval: threshold_filter={pre_count}->{len(fused_rows)} (threshold={threshold})') - - if fused_rows: - _normalize_row_scores( - fused_rows, - source_field='score', - target_field='discovery_score', - default=0.5, - ) - - # ── Legacy graph routing ── - logger.info('\n 🧭 PHASE 2: Legacy Graph Routing') - router_used = 'discovery_only' - agent_rows: list[dict[str, Any]] = [] - - try: - agent_rows = await list_graph_routed_chunks( - db, user_id=user_id, namespace=namespace, query=query, - top_k=top_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - if agent_rows: - router_used = 'discovery+graph' - logger.info(f' 📊 Graph routing: {len(agent_rows)} rows') - except Exception as exc: - logger.error(f' ❌ Graph routing failed (ignored): {exc}') - agent_rows = [] - - if agent_rows: - _normalize_row_scores( - agent_rows, - source_field='score', - target_field='agent_score', - default=0.5, - ) - - combined_rows = [*fused_rows, *agent_rows] - if combined_rows: - try: - chunk_importance_scores = await _load_chunk_importance_scores( - db, - user_id=user_id, - namespace=namespace, - rows=combined_rows, - ) - except Exception as exc: - logger.warning(f'Failed to load chunk importance scores, continuing without importance: {exc}') - chunk_importance_scores = {} - for row in combined_rows: - row['importance_raw_score'] = float(chunk_importance_scores.get(str(row.get('chunk_id') or ''), 0.0) or 0.0) - - ranked_rows = _rank_candidates_by_path(fused_rows, agent_rows, top_k) - if ranked_rows: - logger.info(f'\n 🧮 Unified candidate ranking: {len(ranked_rows)} rows') - for i, row in enumerate(ranked_rows[:10]): - logger.info( - ' ' - f'[{i}] evidence={row.get("evidence_score", 0.0):.4f} ' - f'discovery={row.get("discovery_score", 0.0):.4f} ' - f'agent={row.get("agent_score", 0.0):.4f} ' - f'path={_get_row_path(row)}' - ) - - assembled_rows = await assemble_retrieval_results( - db=db, - rows=ranked_rows, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, + use_agentic=use_agentic, ) - results = [_with_citation(row) for row in assembled_rows] - - response = { - "namespace": namespace, - "query": query, - "router_used": router_used, - "results": results, - } - - if cache_version is not None: - try: - await set_cached_retrieval_query_result( - user_id=user_id, - namespace=namespace, - version=cache_version, - query=query, - top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - response=response, - **cache_extra, - ) - except Exception as e: - logger.warning(f"Failed to write retrieval cache (ignored): {e}") - - try: - schedule_retrieval_hit_stats_update( - user_id=user_id, - namespace=namespace, - results=results, - ) - except Exception as e: - logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") - - elapsed_total = round((time.monotonic() - t_start) * 1000) - logger.info(f'\n{"█" * 70}') - logger.info(f' ✅ RETRIEVAL COMPLETE: {len(results)} results | router={router_used} | {elapsed_total}ms') - for i, r in enumerate(results[:10]): - src = r.get('source', {}) - logger.info( - f' [{i+1}] type={r.get("chunk_type","?")} score={r.get("score",0):.4f}' - f' path={src.get("section_path","")}' - f' file={src.get("source_file_name","")}' - ) - if len(results) > 10: - logger.info(f' ... and {len(results) - 10} more') - logger.info(f'{"█" * 70}') - - return await _to_public_response(response) diff --git a/packages/shared-python/shared/services/retrieval/cache_service.py b/packages/shared-python/shared/services/retrieval/cache_service.py index bb165eb41..e081b2a8c 100644 --- a/packages/shared-python/shared/services/retrieval/cache_service.py +++ b/packages/shared-python/shared/services/retrieval/cache_service.py @@ -5,6 +5,7 @@ from loguru import logger +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.redis import RedisServiceFactory _RETRIEVAL_CACHE_TTL_SECONDS = 300 @@ -13,6 +14,7 @@ def _namespace_version_key(*, user_id: str, namespace: str) -> str: + namespace = normalize_retrieval_namespace(namespace) return f"retrieval:version:{user_id}:{namespace}" @@ -75,6 +77,7 @@ def _query_cache_key( exclude_sections: list[dict[str, str]], **extra_params: Any, ) -> str: + namespace = normalize_retrieval_namespace(namespace) digest = _cache_shape_digest( query=query, top_k=top_k, @@ -110,7 +113,8 @@ async def invalidate_retrieval_cache_namespaces( *, user_id: str, namespaces: list[str] ) -> None: seen: set[str] = set() - for namespace in namespaces: + for raw_namespace in namespaces: + namespace = normalize_retrieval_namespace(raw_namespace) if not namespace or namespace in seen: continue seen.add(namespace) @@ -184,6 +188,7 @@ async def set_cached_retrieval_query_result( def _workflow_plan_cache_key(*, user_id: str, namespace: str, query: str) -> str: + namespace = normalize_retrieval_namespace(namespace) digest = hashlib.sha256(query.encode("utf-8")).hexdigest() return f"retrieval:workflow:plan:{user_id}:{namespace}:{digest}" diff --git a/packages/shared-python/shared/services/retrieval/execution/__init__.py b/packages/shared-python/shared/services/retrieval/execution/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/execution/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/execution/legacy_route.py b/packages/shared-python/shared/services/retrieval/execution/legacy_route.py new file mode 100644 index 000000000..6b1976f53 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/execution/legacy_route.py @@ -0,0 +1,347 @@ +from __future__ import annotations + +import time +from typing import Any + +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.search.channels import content_channel, path_channel, term_channel +from shared.services.retrieval.graph.query_service import GraphQueryService +from shared.services.retrieval.search.ranking import rank_retrieval_candidates +from shared.services.retrieval.execution.response_projection import attach_citation +from shared.services.retrieval.hydration.result_assembly import assemble_retrieval_results +from shared.services.retrieval.execution.route_types import ( + RetrievalRouteContext, + RetrievalRouteOutcome, +) +from shared.services.retrieval.search.scoring import ( + get_row_path, + merge_channels_rrf, + merge_same_section_rows, + normalize_row_scores, +) +from shared.services.retrieval.settings import ( + CHANNEL_WEIGHT_CONTENT, + CHANNEL_WEIGHT_PATH, + CHANNEL_WEIGHT_TERM, + INTERNAL_RECALL_K_MULTIPLIER, +) + + +async def run_legacy_retrieval_route( + context: RetrievalRouteContext, +) -> RetrievalRouteOutcome: + active_channels = set(context.channels) if context.channels else { + "path", + "content", + "term", + } + logger.info( + f"\n PHASE 1: Bottom-Layer Discovery " + f"(channels={sorted(active_channels)})" + ) + logger.info(f" effective_recall_k={context.effective_recall_k}") + + path_rows = await _load_path_rows(context, active_channels) + content_rows = await _load_content_rows(context, active_channels) + term_rows = await _load_term_rows(context, active_channels) + + fused_rows = _fuse_legacy_rows( + context=context, + path_rows=path_rows, + content_rows=content_rows, + term_rows=term_rows, + ) + router_used, graph_rows = await _run_legacy_graph_routing(context) + + ranked_rows = await rank_retrieval_candidates( + context.db, + user_id=context.user_id, + namespace=context.namespace, + discovery_rows=fused_rows, + routed_rows=graph_rows, + top_k=context.top_k, + ) + _log_ranked_rows(ranked_rows) + + assembled_rows = await assemble_retrieval_results( + db=context.db, + rows=ranked_rows, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + ) + results = [attach_citation(row) for row in assembled_rows] + response = { + "namespace": context.namespace, + "query": context.query, + "router_used": router_used, + "results": results, + } + return RetrievalRouteOutcome( + response=response, + hit_stats_results=results, + completion_label="RETRIEVAL", + completion_count=len(results), + completion_detail=f"results | router={router_used}", + ) + + +async def list_graph_routed_chunks( + db: AsyncSession, + *, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], +) -> list[dict[str, Any]]: + service = GraphQueryService() + entry_document_ids = await service.find_entry_documents( + db, + user_id=user_id, + namespace=namespace, + query=query, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + ) + return await service.collect_candidate_chunks( + db, + user_id=user_id, + namespace=namespace, + entry_document_ids=entry_document_ids, + query=query, + top_k=top_k * INTERNAL_RECALL_K_MULTIPLIER, + exclude_sections=exclude_sections, + ) + + +async def _run_legacy_graph_routing( + context: RetrievalRouteContext, +) -> tuple[str, list[dict[str, Any]]]: + logger.info("\n PHASE 2: Legacy Graph Routing") + try: + graph_rows = await list_graph_routed_chunks( + context.db, + user_id=context.user_id, + namespace=context.namespace, + query=context.query, + top_k=context.top_k, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + ) + if graph_rows: + logger.info(f" Graph routing: {len(graph_rows)} rows") + normalize_row_scores( + graph_rows, + source_field="score", + target_field="agent_score", + default=0.5, + ) + return "discovery+graph", graph_rows + except Exception as exc: + logger.error(f" Graph routing failed (ignored): {exc}") + + return "discovery_only", [] + + +async def _load_path_rows( + context: RetrievalRouteContext, + active_channels: set[str], +) -> list[dict[str, Any]]: + if "path" not in active_channels: + return [] + + start_time = time.monotonic() + rows = await path_channel( + context.db, + user_id=context.user_id, + namespace=context.namespace, + query=context.query, + top_k=context.effective_recall_k, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + signal_paths=context.signal_paths, + filter_mode=context.filter_mode, + ) + elapsed_ms = round((time.monotonic() - start_time) * 1000) + logger.info(f"\n path_channel: {len(rows)} rows in {elapsed_ms}ms") + for index, row in enumerate(rows[:5]): + logger.info( + f" [{index}] score={row.get('score', 0):.4f} " + f"path={row.get('section_path', '') or row.get('source_chunk_path', '')} " + f"type={row.get('chunk_type', '?')}" + ) + if len(rows) > 5: + logger.info(f" ... and {len(rows) - 5} more") + return rows + + +async def _load_content_rows( + context: RetrievalRouteContext, + active_channels: set[str], +) -> list[dict[str, Any]]: + if "content" not in active_channels: + return [] + + start_time = time.monotonic() + rows = await content_channel( + context.db, + user_id=context.user_id, + namespace=context.namespace, + query=context.query, + top_k=context.effective_recall_k, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + signal_paths=context.signal_paths, + filter_mode=context.filter_mode, + ) + elapsed_ms = round((time.monotonic() - start_time) * 1000) + logger.info(f"\n content_channel: {len(rows)} rows in {elapsed_ms}ms") + for index, row in enumerate(rows[:5]): + logger.info( + f" [{index}] score={row.get('score', 0):.4f} " + f"path={row.get('section_path', '') or row.get('source_chunk_path', '')} " + f"content={str(row.get('content', ''))[:80]}" + ) + if len(rows) > 5: + logger.info(f" ... and {len(rows) - 5} more") + return rows + + +async def _load_term_rows( + context: RetrievalRouteContext, + active_channels: set[str], +) -> list[dict[str, Any]]: + if "term" not in active_channels: + return [] + + start_time = time.monotonic() + rows = await term_channel( + context.db, + user_id=context.user_id, + namespace=context.namespace, + query=context.query, + top_k=context.effective_recall_k, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + signal_paths=context.signal_paths, + filter_mode=context.filter_mode, + ) + elapsed_ms = round((time.monotonic() - start_time) * 1000) + logger.info(f"\n term_channel: {len(rows)} rows in {elapsed_ms}ms") + for index, row in enumerate(rows[:5]): + logger.info( + f" [{index}] score={row.get('score', 0):.4f} " + f"path={row.get('section_path', '') or row.get('source_chunk_path', '')} " + f"type={row.get('chunk_type', '?')}" + ) + if len(rows) > 5: + logger.info(f" ... and {len(rows) - 5} more") + return rows + + +def _fuse_legacy_rows( + *, + context: RetrievalRouteContext, + path_rows: list[dict[str, Any]], + content_rows: list[dict[str, Any]], + term_rows: list[dict[str, Any]], +) -> list[dict[str, Any]]: + default_weights = { + "path": CHANNEL_WEIGHT_PATH, + "content": CHANNEL_WEIGHT_CONTENT, + "term": CHANNEL_WEIGHT_TERM, + } + effective_weights = {**default_weights, **(context.channel_weights or {})} + + channel_lists: list[list[dict[str, Any]]] = [] + weight_list: list[float] = [] + + if path_rows: + channel_lists.append(path_rows) + weight_list.append(effective_weights.get("path", CHANNEL_WEIGHT_PATH)) + if content_rows: + channel_lists.append(content_rows) + weight_list.append(effective_weights.get("content", CHANNEL_WEIGHT_CONTENT)) + if term_rows: + channel_lists.append(term_rows) + weight_list.append(effective_weights.get("term", CHANNEL_WEIGHT_TERM)) + + if channel_lists: + fused_rows = merge_channels_rrf( + channel_lists, + weight_list, + context.effective_recall_k, + ) + else: + fused_rows = [] + logger.info( + f"\n RRF Fusion: {len(fused_rows)} rows from " + f"{len(channel_lists)} channels " + f"(weights={dict(zip(['path', 'content', 'term'][:len(weight_list)], weight_list))})" + ) + for index, row in enumerate(fused_rows[:5]): + logger.info( + f" [{index}] rrf_score={row.get('score', 0):.4f} " + f"path={row.get('section_path', '') or row.get('source_chunk_path', '')}" + ) + if len(fused_rows) > 5: + logger.info(f" ... and {len(fused_rows) - 5} more") + + pre_merge = len(fused_rows) + fused_rows = merge_same_section_rows(fused_rows) + if len(fused_rows) != pre_merge: + logger.info(f"retrieval: section_merge={pre_merge}->{len(fused_rows)}") + + if context.channel_weights is not None: + logger.debug(f"retrieval: channel_weights={context.channel_weights}") + + fused_rows = _filter_rows_by_threshold(fused_rows, context) + if fused_rows: + normalize_row_scores( + fused_rows, + source_field="score", + target_field="discovery_score", + default=0.5, + ) + + return fused_rows + + +def _filter_rows_by_threshold( + rows: list[dict[str, Any]], + context: RetrievalRouteContext, +) -> list[dict[str, Any]]: + if context.threshold <= 0.0 or not rows: + return rows + + pre_count = len(rows) + filtered_rows = [ + row for row in rows if row.get("score", 0.0) >= context.threshold + ] + logger.info( + f"retrieval: threshold_filter={pre_count}->{len(filtered_rows)} " + f"(threshold={context.threshold})" + ) + return filtered_rows + + +def _log_ranked_rows(ranked_rows: list[dict[str, Any]]) -> None: + if not ranked_rows: + return + + logger.info(f"\n Unified candidate ranking: {len(ranked_rows)} rows") + for index, row in enumerate(ranked_rows[:10]): + logger.info( + " " + f"[{index}] evidence={row.get('evidence_score', 0.0):.4f} " + f"discovery={row.get('discovery_score', 0.0):.4f} " + f"agent={row.get('agent_score', 0.0):.4f} " + f"path={get_row_path(row)}" + ) diff --git a/packages/shared-python/shared/services/retrieval/execution/plan.py b/packages/shared-python/shared/services/retrieval/execution/plan.py new file mode 100644 index 000000000..89fc85048 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/execution/plan.py @@ -0,0 +1,269 @@ +from __future__ import annotations + +import time +from typing import Any + +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.cache_service import ( + get_cached_retrieval_query_result, + set_cached_retrieval_query_result, +) +from shared.services.retrieval.execution.routes import run_retrieval_route +from shared.services.retrieval.stats.recorder import ( + schedule_retrieval_hit_stats_update, +) +from shared.services.retrieval.execution.response_projection import ( + project_public_retrieval_response, +) +from shared.services.retrieval.execution.query_request import RetrievalQuery + + +async def run_retrieval_query( + *, + db: AsyncSession, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + data_type: int = 1, + signal_paths: list[str] | None = None, + filter_mode: str = "delete", + channels: list[str] | None = None, + channel_weights: dict[str, float] | None = None, + rerank: bool = False, + threshold: float = 0.0, + internal_recall_k: int | None = None, + use_agentic: bool | None = None, +) -> dict[str, Any]: + """Run retrieval through the plan module.""" + return await RetrievalExecutionPlan( + RetrievalQuery.from_parameters( + db=db, + user_id=user_id, + namespace=namespace, + query=query, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + data_type=data_type, + signal_paths=signal_paths, + filter_mode=filter_mode, + channels=channels, + channel_weights=channel_weights, + rerank=rerank, + threshold=threshold, + internal_recall_k=internal_recall_k, + use_agentic=use_agentic, + ) + ).execute() + + +class RetrievalExecutionPlan: + def __init__(self, request: RetrievalQuery) -> None: + self.request = request + + async def execute(self) -> dict[str, Any]: + request = self.request + + start_time = time.monotonic() + _log_retrieval_start( + query=request.query, + user_id=request.user_id, + namespace=request.namespace, + top_k=request.top_k, + data_type=request.data_type, + exclude_document_ids=request.exclude_document_ids, + exclude_sections=request.exclude_sections, + ) + + if not request.query: + logger.info(" ⛔ Empty query filtered, skipping retrieval pipeline") + return { + "namespace": request.namespace, + "query": request.query, + "router_used": "empty_query_filtered", + "results": [], + } + + allowed_chunk_types = request.resolve_allowed_chunk_types() + effective_recall_k = request.resolve_effective_recall_k() + logger.info( + f" allowed_chunk_types={allowed_chunk_types} " + f"effective_recall_k={effective_recall_k} " + f"signal_paths={request.signal_paths} filter_mode={request.filter_mode} " + f"rerank={request.rerank} threshold={request.threshold}" + ) + + cache_extra = request.build_cache_extra() + cache_version, cached_response = await _read_cached_response( + user_id=request.user_id, + namespace=request.namespace, + query=request.query, + top_k=request.top_k, + exclude_document_ids=request.exclude_document_ids, + exclude_sections=request.exclude_sections, + cache_extra=cache_extra, + ) + if cached_response is not None: + return cached_response + + logger.debug(f" 📦 Cache miss (version={cache_version}), running full pipeline") + + outcome = await run_retrieval_route(request.build_route_context()) + + if cache_version is not None: + await _write_cached_response( + user_id=request.user_id, + namespace=request.namespace, + version=cache_version, + query=request.query, + top_k=request.top_k, + exclude_document_ids=request.exclude_document_ids, + exclude_sections=request.exclude_sections, + response=outcome.response, + cache_extra=cache_extra, + ) + + _schedule_hit_stats_update( + user_id=request.user_id, + namespace=request.namespace, + results=outcome.hit_stats_results, + ) + _log_retrieval_complete( + outcome=outcome.response, + label=outcome.completion_label, + count=outcome.completion_count, + detail=outcome.completion_detail, + elapsed_ms=round((time.monotonic() - start_time) * 1000), + ) + return await project_public_retrieval_response(outcome.response) + + +async def _read_cached_response( + *, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + cache_extra: dict[str, Any], +) -> tuple[int | None, dict[str, Any] | None]: + cache_version: int | None = None + try: + cache_version, cached = await get_cached_retrieval_query_result( + user_id=user_id, + namespace=namespace, + query=query, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + **cache_extra, + ) + if cached: + logger.info(f"retrieval: cache_hit=True version={cache_version}") + _schedule_hit_stats_update( + user_id=user_id, + namespace=namespace, + results=cached.get("results", []), + ) + return cache_version, await project_public_retrieval_response(cached) + except Exception as exc: + logger.warning(f"Failed to read retrieval cache (ignored): {exc}") + return cache_version, None + + +async def _write_cached_response( + *, + user_id: str, + namespace: str, + version: int, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + response: dict[str, Any], + cache_extra: dict[str, Any], +) -> None: + try: + await set_cached_retrieval_query_result( + user_id=user_id, + namespace=namespace, + version=version, + query=query, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + response=response, + **cache_extra, + ) + except Exception as exc: + logger.warning(f"Failed to write retrieval cache (ignored): {exc}") + + +def _schedule_hit_stats_update( + *, + user_id: str, + namespace: str, + results: list[dict[str, Any]], +) -> None: + try: + schedule_retrieval_hit_stats_update( + user_id=user_id, + namespace=namespace, + results=results, + ) + except Exception as exc: + logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {exc}") + + +def _log_retrieval_start( + *, + query: str, + user_id: str, + namespace: str, + top_k: int, + data_type: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], +) -> None: + logger.info("\n" + "█" * 70) + logger.info(" 🚀 RETRIEVAL PIPELINE START") + logger.info(f' query="{query}"') + logger.info( + f" user={user_id} ns={namespace} top_k={top_k} data_type={data_type}" + ) + logger.info( + f" exclude_docs={exclude_document_ids} " + f"exclude_secs={len(exclude_sections)}" + ) + logger.info("█" * 70) + + +def _log_retrieval_complete( + *, + outcome: dict[str, Any], + label: str, + count: int, + detail: str, + elapsed_ms: int, +) -> None: + logger.info(f'\n{"█" * 70}') + logger.info(f" ✅ {label} COMPLETE: {count} {detail} | {elapsed_ms}ms") + results = outcome.get("results", []) + if isinstance(results, list): + for index, result in enumerate(results[:10]): + source = result.get("source", {}) + logger.info( + f" [{index + 1}] type={result.get('chunk_type', '?')} " + f"score={result.get('score', 0):.4f}" + f" path={source.get('section_path', '')}" + f" file={source.get('source_file_name', '')}" + ) + if len(results) > 10: + logger.info(f" ... and {len(results) - 10} more") + logger.info(f'{"█" * 70}') diff --git a/packages/shared-python/shared/services/retrieval/execution/query_request.py b/packages/shared-python/shared/services/retrieval/execution/query_request.py new file mode 100644 index 000000000..493a31d53 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/execution/query_request.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace +from shared.services.retrieval.execution.route_types import RetrievalRouteContext +from shared.services.retrieval.settings import ( + INTERNAL_RECALL_K_MULTIPLIER, + resolve_allowed_chunk_types, +) + + +@dataclass(frozen=True) +class RetrievalQuery: + db: AsyncSession + user_id: str + namespace: str + query: str + top_k: int + exclude_document_ids: list[str] + exclude_sections: list[dict[str, str]] + data_type: int = 1 + signal_paths: list[str] | None = None + filter_mode: str = "delete" + channels: list[str] | None = None + channel_weights: dict[str, float] | None = None + rerank: bool = False + threshold: float = 0.0 + internal_recall_k: int | None = None + use_agentic: bool | None = None + + @classmethod + def from_parameters( + cls, + *, + db: AsyncSession, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + data_type: int = 1, + signal_paths: list[str] | None = None, + filter_mode: str = "delete", + channels: list[str] | None = None, + channel_weights: dict[str, float] | None = None, + rerank: bool = False, + threshold: float = 0.0, + internal_recall_k: int | None = None, + use_agentic: bool | None = None, + ) -> "RetrievalQuery": + return cls( + db=db, + user_id=user_id, + namespace=normalize_retrieval_namespace(namespace), + query=str(query).strip(), + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + data_type=data_type, + signal_paths=signal_paths, + filter_mode=filter_mode, + channels=channels, + channel_weights=channel_weights, + rerank=rerank, + threshold=threshold, + internal_recall_k=internal_recall_k, + use_agentic=use_agentic, + ) + + def build_cache_extra(self) -> dict[str, Any]: + return { + "data_type": self.data_type, + "signal_paths": self.signal_paths, + "filter_mode": self.filter_mode, + "channels": self.channels, + "channel_weights": self.channel_weights, + "rerank": self.rerank, + "threshold": self.threshold, + "internal_recall_k": self.internal_recall_k, + "decomposition_enabled": True, + } + + def resolve_allowed_chunk_types(self) -> set[str] | None: + return resolve_allowed_chunk_types(self.data_type) + + def resolve_effective_recall_k(self) -> int: + if self.internal_recall_k is not None: + return self.internal_recall_k + return self.top_k * INTERNAL_RECALL_K_MULTIPLIER + + def build_route_context(self) -> RetrievalRouteContext: + return RetrievalRouteContext( + db=self.db, + user_id=self.user_id, + namespace=self.namespace, + query=self.query, + top_k=self.top_k, + exclude_document_ids=self.exclude_document_ids, + exclude_sections=self.exclude_sections, + allowed_chunk_types=self.resolve_allowed_chunk_types(), + data_type=self.data_type, + signal_paths=self.signal_paths, + filter_mode=self.filter_mode, + channels=self.channels, + channel_weights=self.channel_weights, + rerank=self.rerank, + threshold=self.threshold, + internal_recall_k=self.internal_recall_k, + effective_recall_k=self.resolve_effective_recall_k(), + use_agentic=self.use_agentic, + ) diff --git a/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py new file mode 100644 index 000000000..c44095f50 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.hydration.reference import hydrate_referenced_chunk_rows +from shared.services.retrieval.execution.response_projection import ( + enrich_referenced_chunks_with_asset_urls, +) +from shared.services.retrieval.hydration.row_utils import build_reference_lookup_key + + +@dataclass(frozen=True) +class ResolvedWorkflowReferences: + refs: list[dict[str, Any]] + rows: list[dict[str, Any]] + + +async def resolve_workflow_references( + *, + db: AsyncSession, + user_id: str, + namespace: str, + refs: list[dict[str, Any]], +) -> ResolvedWorkflowReferences: + enriched_refs = await enrich_referenced_chunks_with_asset_urls(refs) + hydrated_rows = await hydrate_referenced_chunk_rows( + db=db, + user_id=user_id, + namespace=namespace, + refs=enriched_refs, + ) + return _select_matching_references(enriched_refs, hydrated_rows) + + +def _select_matching_references( + refs: list[dict[str, Any]], + rows: list[dict[str, Any]], +) -> ResolvedWorkflowReferences: + selected_refs: list[dict[str, Any]] = [] + selected_rows: list[dict[str, Any]] = [] + seen_row_keys: set[tuple[str, str, str, str]] = set() + + for ref in refs: + matching_row = next( + ( + row + for row in rows + if _matches_reference(ref, row) + and _row_key(row) not in seen_row_keys + ), + None, + ) + if matching_row is None: + continue + + selected_refs.append(ref) + selected_rows.append(matching_row) + seen_row_keys.add(_row_key(matching_row)) + + return ResolvedWorkflowReferences(refs=selected_refs, rows=selected_rows) + + +def _matches_reference(ref: dict[str, Any], row: dict[str, Any]) -> bool: + ref_key = build_reference_lookup_key( + document_id=ref.get("document_id"), + chunk_id=ref.get("chunk_id"), + section_path=ref.get("section_path"), + file_path=ref.get("file_path"), + ) + row_key = _row_key(row) + if ref_key[:2] != row_key[:2]: + return False + if ref_key[2] and ref_key[2] != row_key[2] and not _matches_root_alias(ref, row): + return False + if ref_key[3] and ref_key[3] != row_key[3]: + return False + return True + + +def _row_key(row: dict[str, Any]) -> tuple[str, str, str, str]: + return build_reference_lookup_key( + document_id=row.get("document_id"), + chunk_id=row.get("chunk_id"), + section_path=row.get("section_path"), + file_path=row.get("file_path"), + ) + + +def _matches_root_alias(ref: dict[str, Any], row: dict[str, Any]) -> bool: + ref_section_path = str(ref.get("section_path") or "").strip() + row_section_path = str(row.get("section_path") or "").strip() + source_file_name = str(row.get("source_file_name") or "").strip() + return bool( + source_file_name + and row_section_path == "Root" + and ref_section_path == source_file_name + ) diff --git a/packages/shared-python/shared/services/retrieval/execution/response_projection.py b/packages/shared-python/shared/services/retrieval/execution/response_projection.py new file mode 100644 index 000000000..14e84bf4b --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/execution/response_projection.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from typing import Any + +from shared.services.retrieval.hydration.assets import enrich_rows_with_retrieval_asset_urls +from shared.services.retrieval.hydration.row_utils import ( + PUBLIC_RESULT_FIELDS, + PUBLIC_SOURCE_FIELDS, +) + + +def attach_citation(row: dict[str, Any]) -> dict[str, Any]: + citation = { + 'document_id': row.get('document_id'), + 'chunk_id': row.get('chunk_id'), + 'source_file_name': row.get('source_file_name'), + 'section_path': row.get('section_path'), + } + return {**row, 'citation': citation} + + +def to_public_source(row: dict[str, Any]) -> dict[str, Any]: + return {field: row.get(field) for field in PUBLIC_SOURCE_FIELDS} + + +async def enrich_referenced_chunks_with_asset_urls(refs: list[dict[str, Any]]) -> list[dict[str, Any]]: + return await enrich_rows_with_retrieval_asset_urls( + refs, + log_context='agentic referenced chunk', + ) + + +async def project_public_retrieval_response(response: dict[str, Any]) -> dict[str, Any]: + public_response = { + 'namespace': response.get('namespace'), + 'query': response.get('query'), + 'router_used': response.get('router_used'), + 'results': [], + } + + if response.get('answer_text') is not None: + public_response['answer_text'] = response['answer_text'] + if response.get('referenced_chunks') is not None: + public_response['referenced_chunks'] = response['referenced_chunks'] + + projected_rows = await enrich_rows_with_retrieval_asset_urls( + response.get('results', []), + log_context='retrieval result', + ) + public_results: list[dict[str, Any]] = [] + for row in projected_rows: + public_row: dict[str, Any] = {} + for field in PUBLIC_RESULT_FIELDS: + if field in row: + public_row[field] = row[field] + if 'source' in row: + public_row['source'] = row['source'] + else: + public_row['source'] = to_public_source(row) + public_results.append(public_row) + + public_response['results'] = public_results + return public_response diff --git a/packages/shared-python/shared/services/retrieval/execution/route_types.py b/packages/shared-python/shared/services/retrieval/execution/route_types.py new file mode 100644 index 000000000..1e7e53d79 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/execution/route_types.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncSession + + +@dataclass(frozen=True) +class RetrievalRouteContext: + db: AsyncSession + user_id: str + namespace: str + query: str + top_k: int + exclude_document_ids: list[str] + exclude_sections: list[dict[str, str]] + allowed_chunk_types: set[str] | None + data_type: int + signal_paths: list[str] | None + filter_mode: str + channels: list[str] | None + channel_weights: dict[str, float] | None + rerank: bool + threshold: float + internal_recall_k: int | None + effective_recall_k: int + use_agentic: bool | None + + +@dataclass(frozen=True) +class RetrievalRouteOutcome: + response: dict[str, Any] + hit_stats_results: list[dict[str, Any]] + completion_label: str + completion_count: int + completion_detail: str diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py new file mode 100644 index 000000000..b47c7351b --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -0,0 +1,142 @@ +from __future__ import annotations + +import os + +from loguru import logger + +from shared.services.retrieval.execution.legacy_route import run_legacy_retrieval_route +from shared.services.retrieval.execution.reference_resolver import resolve_workflow_references +from shared.services.retrieval.hydration.result_assembly import assemble_retrieval_results +from shared.services.retrieval.execution.response_projection import ( + attach_citation, +) +from shared.services.retrieval.execution.route_types import ( + RetrievalRouteContext, + RetrievalRouteOutcome, +) +from shared.services.retrieval.search.scoped_corpus import ( + count_scoped_chunks, + load_all_scoped_chunks, +) + + +async def run_retrieval_route( + context: RetrievalRouteContext, +) -> RetrievalRouteOutcome: + small_corpus_outcome = await _try_run_small_corpus_route(context) + if small_corpus_outcome is not None: + return small_corpus_outcome + + if _should_use_agentic_route(context.use_agentic): + return await _run_agentic_route(context) + + return await run_legacy_retrieval_route(context) + + +async def _try_run_small_corpus_route( + context: RetrievalRouteContext, +) -> RetrievalRouteOutcome | None: + try: + total_chunk_count = await count_scoped_chunks( + context.db, + user_id=context.user_id, + namespace=context.namespace, + exclude_document_ids=context.exclude_document_ids, + allowed_chunk_types=context.allowed_chunk_types, + ) + except Exception as exc: + logger.warning( + f"Failed to count scoped chunks, skipping small corpus optimization: {exc}" + ) + total_chunk_count = context.top_k + 1 + + logger.info(f"\n Total chunks in scope: {total_chunk_count}") + if total_chunk_count > context.top_k: + return None + + logger.info( + f" Small corpus optimization: {total_chunk_count} chunks " + f"<= top_k={context.top_k}, returning all" + ) + all_rows = await load_all_scoped_chunks( + context.db, + user_id=context.user_id, + namespace=context.namespace, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + signal_paths=context.signal_paths or [], + filter_mode=context.filter_mode, + ) + logger.info( + f" small_corpus load: loaded={len(all_rows)} rows after signal/exclude filters" + ) + assembled_rows = await assemble_retrieval_results( + db=context.db, + rows=all_rows, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + ) + results = [attach_citation(row) for row in assembled_rows] + response = { + "namespace": context.namespace, + "query": context.query, + "router_used": "small_corpus_all", + "results": results, + } + return RetrievalRouteOutcome( + response=response, + hit_stats_results=results, + completion_label="Small corpus", + completion_count=len(results), + completion_detail="results", + ) + + +def _should_use_agentic_route(use_agentic: bool | None) -> bool: + if use_agentic is not None: + return use_agentic + return os.environ.get("RETRIEVAL_AGENTIC_ENABLED", "true") == "true" + + +async def _run_agentic_route( + context: RetrievalRouteContext, +) -> RetrievalRouteOutcome: + from shared.services.retrieval.workflow.orchestrator import WorkflowOrchestrator + from shared.services.retrieval.workflow.run_request import WorkflowRunRequest + + workflow = WorkflowOrchestrator() + workflow_result = await workflow.run_request( + context.db, + request=WorkflowRunRequest.from_route_context(context), + ) + + resolved_references = await resolve_workflow_references( + db=context.db, + user_id=context.user_id, + namespace=context.namespace, + refs=workflow_result.referenced_chunks, + ) + assembled_workflow_rows = await assemble_retrieval_results( + db=context.db, + rows=resolved_references.rows, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + ) + response = workflow_result.to_api_response() + response["referenced_chunks"] = resolved_references.refs + response["results"] = [attach_citation(row) for row in assembled_workflow_rows] + + completion_detail = ( + f"chunks | answer={len(workflow_result.answer_text)} chars | " + f"router={workflow_result.router_used}" + ) + return RetrievalRouteOutcome( + response=response, + hit_stats_results=resolved_references.refs, + completion_label="AGENTIC RETRIEVAL", + completion_count=len(resolved_references.refs), + completion_detail=completion_detail, + ) diff --git a/packages/shared-python/shared/services/retrieval/graph/__init__.py b/packages/shared-python/shared/services/retrieval/graph/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/graph/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/graph/keywords.py b/packages/shared-python/shared/services/retrieval/graph/keywords.py new file mode 100644 index 000000000..5769bd355 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/graph/keywords.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +import math +import re +from typing import Any + +MIN_KEYWORD_OVERLAP = 3 +KEYWORD_SCORE_WEIGHT = 1.0 +MIN_SCORE_THRESHOLD = 0.8 + + +def normalize_keyword(keyword: str) -> str: + """Normalize a keyword: lowercase, strip, collapse spaces.""" + keyword = keyword.lower().strip() + return re.sub(r'\s+', ' ', keyword) + + +def extract_keywords_from_chunk_metadata(meta: dict) -> list[str]: + """Extract keywords from chunk metadata.""" + if not isinstance(meta, dict): + return [] + + keywords = meta.get('keywords', []) + if isinstance(keywords, list) and keywords: + return [str(keyword) for keyword in keywords if keyword] + + tokens = meta.get('tokens', []) + if isinstance(tokens, list) and tokens: + return [str(token) for token in tokens if token and len(str(token)) > 1] + + return [] + + +def compute_tfidf_keywords( + chunk_metadata_list: list[dict[str, Any]], + top_k: int = 10, +) -> list[str]: + """Compute TF-IDF keywords from chunk metadata.""" + df_count: dict[str, int] = {} + tf_count: dict[str, int] = {} + total = len(chunk_metadata_list) or 1 + for meta in chunk_metadata_list: + keywords = extract_keywords_from_chunk_metadata(meta) + seen: set[str] = set() + for keyword in keywords: + if len(str(keyword)) <= 1 or re.match(r'^\d+[.,%]*$', str(keyword)): + continue + normalized = normalize_keyword(str(keyword)) + if not normalized: + continue + tf_count[normalized] = tf_count.get(normalized, 0) + 1 + if normalized not in seen: + df_count[normalized] = df_count.get(normalized, 0) + 1 + seen.add(normalized) + scored = [ + (term, freq * (math.log(total / (df_count.get(term, 1))) + 1)) + for term, freq in tf_count.items() + ] + scored.sort(key=lambda item: item[1], reverse=True) + return [term for term, _ in scored[:top_k]] + + +def compute_keyword_score( + shared_keywords: set[str], + keywords_a: set[str], + keywords_b: set[str], + weight: float = 1.0, +) -> float: + """Character-length-weighted keyword overlap score.""" + weighted_a = sum(len(keyword) for keyword in keywords_a) + weighted_b = sum(len(keyword) for keyword in keywords_b) + denominator = min(weighted_a, weighted_b) + if denominator == 0: + return 0.0 + weighted_shared = sum(len(keyword) for keyword in shared_keywords) + return weight * weighted_shared / denominator + + +def get_normalized_keyword_set(chunk_metadata_list: list[dict[str, Any]]) -> set[str]: + """Collect all normalized keywords from chunk metadata for a document.""" + result: set[str] = set() + for meta in chunk_metadata_list: + for keyword in extract_keywords_from_chunk_metadata(meta): + normalized = normalize_keyword(str(keyword)) + if normalized and len(normalized) > 1 and not re.match( + r'^\d+[.,%]*$', normalized + ): + result.add(normalized) + return result + + +def extract_document_top_summary(chunk_metadata_list: list[dict[str, Any]]) -> str: + """Read the parser-injected top summary from chunk metadata.""" + for meta in chunk_metadata_list: + if not isinstance(meta, dict): + continue + summary = str(meta.get('document_top_summary') or '').strip() + if summary: + return summary + return '' diff --git a/packages/shared-python/shared/services/retrieval/graph/query_service.py b/packages/shared-python/shared/services/retrieval/graph/query_service.py new file mode 100644 index 000000000..819d99469 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/graph/query_service.py @@ -0,0 +1,208 @@ +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from typing import Any + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.job_result import JobResult +from shared.services.retrieval.search.section_filters import is_excluded_section + +_SECTION_EXCLUSION_PAGE_MULTIPLIER = 2 + + +def _build_lexical_match_predicate(query: str): + like = f'%{query}%' + return ( + DocumentChunk.content_lexical_text.ilike(like) + | DocumentChunk.path_lexical_text.ilike(like) + ) + + +class GraphQueryService: + """Read-side graph routing before canonical chunk hydration.""" + + async def find_entry_documents( + self, + db: AsyncSession, + *, + user_id: str, + namespace: str, + query: str, + exclude_document_ids: Iterable[str] = (), + exclude_sections: Iterable[dict[str, str]] = (), + ) -> list[str]: + query_lc = query.lower().strip() + excluded_document_ids = set(exclude_document_ids) + + if query_lc: + section_matches = await self._find_documents_by_section( + db, + user_id=user_id, + namespace=namespace, + query=query_lc, + exclude_document_ids=excluded_document_ids, + exclude_sections=exclude_sections, + ) + if section_matches: + return section_matches + + return await self._find_documents_by_content( + db, + user_id=user_id, + namespace=namespace, + query=query_lc, + exclude_document_ids=excluded_document_ids, + ) + + async def _find_documents_by_section( + self, + db: AsyncSession, + *, + user_id: str, + namespace: str, + query: str, + exclude_document_ids: set[str], + exclude_sections: Iterable[dict[str, str]], + ) -> list[str]: + like = f'%{query}%' + stmt = ( + select(DocumentSection.document_id) + .join( + Document, + (Document.document_id == DocumentSection.document_id) + & (Document.current_job_result_id == DocumentSection.job_result_id), + ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where( + DocumentSection.section_title.ilike(like) + | DocumentSection.section_path.ilike(like) + ) + .distinct() + ) + if exclude_document_ids: + stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) + for item in exclude_sections or (): + if not isinstance(item, dict): + continue + excluded_document_id = str(item.get('document_id') or '').strip() + excluded_path = str(item.get('section_path') or '').strip() + if excluded_document_id and excluded_path: + stmt = stmt.where( + ~( + (DocumentSection.document_id == excluded_document_id) + & ( + (DocumentSection.section_path == excluded_path) + | DocumentSection.section_path.like(f'{excluded_path} / %') + ) + ) + ) + + result = await db.execute(stmt) + return [document_id for (document_id,) in result.all()] + + async def _find_documents_by_content( + self, + db: AsyncSession, + *, + user_id: str, + namespace: str, + query: str, + exclude_document_ids: set[str], + ) -> list[str]: + like = f'%{query}%' + stmt = ( + select(Document.document_id) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where(DocumentChunk.content_lexical_text.ilike(like)) + ) + if exclude_document_ids: + stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) + result = await db.execute(stmt) + seen: list[str] = [] + for (document_id,) in result.all(): + if document_id and document_id not in seen: + seen.append(document_id) + return seen + + async def collect_candidate_chunks( + self, + db: AsyncSession, + *, + user_id: str, + namespace: str, + entry_document_ids: Sequence[str], + query: str, + top_k: int, + exclude_sections: Iterable[dict[str, str]] = (), + ) -> list[dict[str, Any]]: + if not entry_document_ids: + return [] + page_size = top_k + if exclude_sections: + page_size = max(top_k, top_k * _SECTION_EXCLUSION_PAGE_MULTIPLIER) + base_stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) + .outerjoin( + DocumentSection, + DocumentSection.section_id == DocumentChunk.section_id, + ) + .join(JobResult, JobResult.id == DocumentChunk.job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where(Document.document_id.in_(list(entry_document_ids))) + .where(_build_lexical_match_predicate(query)) + .order_by(DocumentChunk.sort_order) + ) + rows: list[dict[str, Any]] = [] + offset = 0 + while len(rows) < top_k: + result = await db.execute(base_stmt.limit(page_size).offset(offset)) + result_rows = result.all() + if not result_rows: + break + for document, chunk, section, job_result in result_rows: + section_path = section.section_path if section else None + if is_excluded_section( + document_id=document.document_id, + section_path=section_path, + exclude_sections=exclude_sections, + ): + continue + rows.append({ + 'document_id': document.document_id, + 'chunk_id': chunk.chunk_id, + 'section_id': chunk.section_id, + 'section_path': section_path, + 'source_file_name': document.source_file_name, + 'chunk_type': chunk.chunk_type, + 'content': chunk.content, + 'score': 2.0, + 'file_path': chunk.file_path, + 'chunk_metadata': chunk.chunk_metadata or {}, + 'job_result_id': chunk.job_result_id, + 'job_id': job_result.job_id if job_result else None, + }) + if len(rows) >= top_k: + break + if len(result_rows) < page_size: + break + offset += page_size + return rows diff --git a/packages/shared-python/shared/services/retrieval/graph/service.py b/packages/shared-python/shared/services/retrieval/graph/service.py new file mode 100644 index 000000000..0d1cd1456 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/graph/service.py @@ -0,0 +1,203 @@ +from __future__ import annotations + +import logging +from collections import defaultdict +from dataclasses import dataclass + +from sqlalchemy import delete, or_, select +from sqlalchemy.orm import Session + +from shared.models.database.document import ( + Document, + DocumentChunk, + GraphEdge, + GraphNode, +) +from shared.services.retrieval.graph.keywords import ( + KEYWORD_SCORE_WEIGHT, + MIN_KEYWORD_OVERLAP, + MIN_SCORE_THRESHOLD, + compute_keyword_score, + compute_tfidf_keywords, + extract_document_top_summary, + get_normalized_keyword_set, + normalize_keyword, +) + +logger = logging.getLogger(__name__) + + +@dataclass +class GraphScope: + user_id: str + namespace: str + + +class DocumentGraphService: + """Write-side graph publication over persisted graph_nodes/graph_edges. + + Aligned with KB's knowledge_graph.json structure: + - Only document-level nodes (no section nodes) + - Document nodes carry rich metadata: top_keywords, chunks_count, types, top_summary + - Edges are keyword-overlap-based cross-document connections with meaningful scores + """ + + def publish_document_graph( + self, + db: Session, + *, + user_id: str, + namespace: str, + document_id: str, + job_result_id: str, + ) -> None: + document = db.execute( + select(Document).where(Document.document_id == document_id) + ).scalar_one_or_none() + if document is None: + return + + # ── Gather chunk metadata for keyword extraction ── + chunk_meta_rows = list( + db.execute( + select(DocumentChunk.chunk_type, DocumentChunk.chunk_metadata) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + ).all() + ) + chunk_metadata_list = [row[1] or {} for row in chunk_meta_rows] + + # Compute document-level metadata (aligned with KB knowledge_graph.json files dict) + top_keywords = compute_tfidf_keywords(chunk_metadata_list) + new_doc_kws = get_normalized_keyword_set(chunk_metadata_list) + + types_breakdown: dict[str, int] = defaultdict(int) + for chunk_type, _ in chunk_meta_rows: + types_breakdown[chunk_type or 'text'] += 1 + chunks_count = len(chunk_meta_rows) + + top_summary = extract_document_top_summary(chunk_metadata_list) + + # ── Clean up old graph data for this document ── + self.remove_document_graph( + db, + scope=GraphScope(user_id=user_id, namespace=namespace), + document_id=document_id, + ) + + # ── Create document-level node (no section nodes — aligned with KB KG) ── + document_node_id = f"doc:{document_id}" + db.add( + GraphNode( + node_id=document_node_id, + user_id=user_id, + namespace=namespace, + node_kind='document', + owner_document_id=document_id, + job_result_id=job_result_id, + ref_document_id=document_id, + ref_section_id=None, + properties={ + 'source_file_name': document.source_file_name, + 'top_keywords': top_keywords, + 'chunks_count': chunks_count, + 'types': dict(types_breakdown), + 'top_summary': top_summary, + }, + ) + ) + db.flush() + + # ── Keyword-overlap-based cross-document edges ── + # Only create edges where keyword overlap score >= threshold. + other_doc_nodes = list( + db.execute( + select(GraphNode) + .where(GraphNode.user_id == user_id) + .where(GraphNode.namespace == namespace) + .where(GraphNode.node_kind == 'document') + .where(GraphNode.owner_document_id != document_id) + ).scalars() + ) + + for peer_node in other_doc_nodes: + peer_props = peer_node.properties or {} + peer_keywords = peer_props.get('top_keywords', []) + + # Build normalized keyword sets for comparison + peer_kws: set[str] = set() + for k in peer_keywords: + normalized = normalize_keyword(str(k)) + if normalized: + peer_kws.add(normalized) + + if not peer_kws or not new_doc_kws: + continue + + # Find shared keywords + shared_kws = new_doc_kws & peer_kws + if len(shared_kws) < MIN_KEYWORD_OVERLAP: + continue + + # Compute character-length-weighted score + score = compute_keyword_score( + shared_keywords=shared_kws, + keywords_a=new_doc_kws, + keywords_b=peer_kws, + weight=KEYWORD_SCORE_WEIGHT, + ) + if score < MIN_SCORE_THRESHOLD: + continue + + # Create edge with meaningful weight and metadata + peer_doc_id = peer_node.owner_document_id + edge_pair = tuple(sorted([document_id, peer_doc_id])) + db.add( + GraphEdge( + edge_id=f"related:{edge_pair[0]}<->{edge_pair[1]}", + user_id=user_id, + namespace=namespace, + edge_kind='related', + source_node_id=document_node_id, + target_node_id=peer_node.node_id, + owner_document_id=document_id, + job_result_id=job_result_id, + is_directed=False, + weight=round(score, 4), + properties={ + 'shared_keywords': sorted(shared_kws), + 'connection_count': len(shared_kws), + }, + ) + ) + + db.flush() + logger.info( + f"publish_document_graph: doc={document_id} " + f"keywords={len(top_keywords)} chunks={chunks_count}" + ) + + def remove_document_graph( + self, db: Session, *, scope: GraphScope | None, document_id: str + ) -> None: + document_node_id = f"doc:{document_id}" + edge_delete = delete(GraphEdge).where( + or_( + GraphEdge.owner_document_id == document_id, + GraphEdge.source_node_id == document_node_id, + GraphEdge.target_node_id == document_node_id, + ) + ) + node_delete = delete(GraphNode).where(GraphNode.owner_document_id == document_id) + if scope is not None: + edge_delete = edge_delete.where( + GraphEdge.user_id == scope.user_id, + GraphEdge.namespace == scope.namespace, + ) + node_delete = node_delete.where( + GraphNode.user_id == scope.user_id, + GraphNode.namespace == scope.namespace, + ) + db.execute(edge_delete) + db.execute(node_delete) + db.flush() diff --git a/packages/shared-python/shared/services/retrieval/graph_service.py b/packages/shared-python/shared/services/retrieval/graph_service.py deleted file mode 100644 index 8244a8a72..000000000 --- a/packages/shared-python/shared/services/retrieval/graph_service.py +++ /dev/null @@ -1,480 +0,0 @@ -from __future__ import annotations - -import logging -import math -import re -from collections import defaultdict -from dataclasses import dataclass -from typing import Any, Iterable, Sequence - -from sqlalchemy import delete, or_, select -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.orm import Session - -from shared.models.database.document import Document, DocumentChunk, DocumentSection, GraphEdge, GraphNode -from shared.models.database.job_result import JobResult - -logger = logging.getLogger(__name__) - -_SECTION_EXCLUSION_PAGE_MULTIPLIER = 2 - -# ── Keyword overlap config (aligned with connect_builder DEFAULT_CONFIG) ── -_MIN_KEYWORD_OVERLAP = 3 -_KEYWORD_SCORE_WEIGHT = 1.0 -_MIN_SCORE_THRESHOLD = 0.8 -_CROSS_FILE_ONLY = True -_MAX_CONTENT_OVERLAP = 0.8 - - -def _build_lexical_match_predicate(query: str): - like = f'%{query}%' - return ( - DocumentChunk.content_lexical_text.ilike(like) - | DocumentChunk.path_lexical_text.ilike(like) - ) - - -def is_excluded_section( - *, - document_id: str | None, - section_path: str | None, - exclude_sections: Iterable[dict[str, str]], -) -> bool: - document_id = str(document_id or '').strip() - section_path = str(section_path or '').strip() - if not document_id or not section_path: - return False - for item in exclude_sections: - if not isinstance(item, dict): - continue - exc_doc = str(item.get('document_id') or '').strip() - exc_path = str(item.get('section_path') or '').strip() - if document_id == exc_doc and (section_path == exc_path or section_path.startswith(exc_path + ' / ')): - return True - return False - - -# ── Keyword extraction & scoring (aligned with connect_builder/builder.py) ── - -def _normalize_keyword(keyword: str) -> str: - """Normalize a keyword: lowercase, strip, collapse spaces.""" - kw = keyword.lower().strip() - return re.sub(r'\s+', ' ', kw) - - -def _extract_keywords_from_chunk_metadata(meta: dict) -> list[str]: - """Extract keywords from chunk metadata, same logic as builder._get_keywords.""" - if not isinstance(meta, dict): - return [] - # Try metadata.keywords - kws = meta.get('keywords', []) - if isinstance(kws, list) and kws: - return [str(k) for k in kws if k] - # Fallback: tokens - tokens = meta.get('tokens', []) - if isinstance(tokens, list) and tokens: - return [str(t) for t in tokens if t and len(str(t)) > 1] - return [] - - -def _compute_tfidf_keywords( - chunk_metadata_list: list[dict[str, Any]], - top_k: int = 10, -) -> list[str]: - """Compute TF-IDF keywords from chunk metadata, aligned with graph_builder.""" - df_count: dict[str, int] = {} - tf_count: dict[str, int] = {} - total = len(chunk_metadata_list) or 1 - for meta in chunk_metadata_list: - kws = _extract_keywords_from_chunk_metadata(meta) - seen: set[str] = set() - for k in kws: - if len(str(k)) <= 1 or re.match(r'^\d+[.,%]*$', str(k)): - continue - lower = _normalize_keyword(str(k)) - if not lower: - continue - tf_count[lower] = tf_count.get(lower, 0) + 1 - if lower not in seen: - df_count[lower] = df_count.get(lower, 0) + 1 - seen.add(lower) - scored = [ - (term, freq * (math.log(total / (df_count.get(term, 1))) + 1)) - for term, freq in tf_count.items() - ] - scored.sort(key=lambda x: x[1], reverse=True) - return [s[0] for s in scored[:top_k]] - - -def _compute_keyword_score( - shared_kws: set[str], - kws_a: set[str], - kws_b: set[str], - weight: float = 1.0, -) -> float: - """Character-length-weighted keyword overlap score (aligned with builder.py). - - Longer tokens contribute more: '施工现场'(4) has 2x weight of '交底'(2). - Formula: score = weight * sum(len(kw) for shared) / min(sum(len) for A, sum(len) for B) - """ - weighted_a = sum(len(k) for k in kws_a) - weighted_b = sum(len(k) for k in kws_b) - denominator = min(weighted_a, weighted_b) - if denominator == 0: - return 0.0 - weighted_shared = sum(len(k) for k in shared_kws) - return weight * weighted_shared / denominator - - -def _get_normalized_keyword_set(chunk_metadata_list: list[dict[str, Any]]) -> set[str]: - """Collect all normalized keywords from chunk metadata for a document.""" - result: set[str] = set() - for meta in chunk_metadata_list: - for k in _extract_keywords_from_chunk_metadata(meta): - normalized = _normalize_keyword(str(k)) - if normalized and len(normalized) > 1 and not re.match(r'^\d+[.,%]*$', normalized): - result.add(normalized) - return result - - -def _extract_document_top_summary( - chunk_metadata_list: list[dict[str, Any]], - section_titles: Sequence[str], -) -> str: - """Extract document_top_summary from chunk metadata. - - The summary is injected by kb_tasks.py via load_nav_top_summary() - at parse time, so it should always be present. If missing, return empty - string rather than fabricating a low-quality fallback. - """ - for meta in chunk_metadata_list: - if not isinstance(meta, dict): - continue - summary = str(meta.get('document_top_summary') or '').strip() - if summary: - return summary - return '' - - -@dataclass -class GraphScope: - user_id: str - namespace: str - - -class DocumentGraphService: - """Write-side graph publication over persisted graph_nodes/graph_edges. - - Aligned with KB's knowledge_graph.json structure: - - Only document-level nodes (no section nodes) - - Document nodes carry rich metadata: top_keywords, chunks_count, types, top_summary - - Edges are keyword-overlap-based cross-document connections with meaningful scores - - Edge scoring uses connect_builder DEFAULT_CONFIG thresholds - """ - - def publish_document_graph(self, db: Session, *, user_id: str, namespace: str, document_id: str, job_result_id: str) -> None: - document = db.execute( - select(Document).where(Document.document_id == document_id) - ).scalar_one_or_none() - if document is None: - return - - # ── Gather chunk metadata for keyword extraction ── - chunk_meta_rows = list( - db.execute( - select(DocumentChunk.chunk_type, DocumentChunk.chunk_metadata) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - ).all() - ) - chunk_metadata_list = [row[1] or {} for row in chunk_meta_rows] - - # Compute document-level metadata (aligned with KB knowledge_graph.json files dict) - top_keywords = _compute_tfidf_keywords(chunk_metadata_list) - new_doc_kws = _get_normalized_keyword_set(chunk_metadata_list) - - types_breakdown: dict[str, int] = defaultdict(int) - for chunk_type, _ in chunk_meta_rows: - types_breakdown[chunk_type or 'text'] += 1 - chunks_count = len(chunk_meta_rows) - - sections = [ - section_title - for section_title in db.execute( - select(DocumentSection.section_title) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - .where(DocumentSection.section_level <= 2) - .order_by(DocumentSection.sort_order) - ).scalars() - if section_title is not None - ] - top_summary = _extract_document_top_summary(chunk_metadata_list, sections) - - # ── Clean up old graph data for this document ── - self.remove_document_graph(db, scope=GraphScope(user_id=user_id, namespace=namespace), document_id=document_id) - - # ── Create document-level node (no section nodes — aligned with KB KG) ── - document_node_id = f"doc:{document_id}" - db.add( - GraphNode( - node_id=document_node_id, - user_id=user_id, - namespace=namespace, - node_kind='document', - owner_document_id=document_id, - job_result_id=job_result_id, - ref_document_id=document_id, - ref_section_id=None, - properties={ - 'source_file_name': document.source_file_name, - 'top_keywords': top_keywords, - 'chunks_count': chunks_count, - 'types': dict(types_breakdown), - 'top_summary': top_summary, - }, - ) - ) - db.flush() - - # ── Keyword-overlap-based cross-document edges (aligned with KB edges) ── - # Only create edges where keyword overlap score >= threshold, - # matching connect_builder DEFAULT_CONFIG parameters. - other_doc_nodes = list( - db.execute( - select(GraphNode) - .where(GraphNode.user_id == user_id) - .where(GraphNode.namespace == namespace) - .where(GraphNode.node_kind == 'document') - .where(GraphNode.owner_document_id != document_id) - ).scalars() - ) - - for peer_node in other_doc_nodes: - peer_props = peer_node.properties or {} - peer_keywords = peer_props.get('top_keywords', []) - - # Build normalized keyword sets for comparison - peer_kws: set[str] = set() - for k in peer_keywords: - normalized = _normalize_keyword(str(k)) - if normalized: - peer_kws.add(normalized) - - if not peer_kws or not new_doc_kws: - continue - - # Find shared keywords - shared_kws = new_doc_kws & peer_kws - if len(shared_kws) < _MIN_KEYWORD_OVERLAP: - continue - - # Compute character-length-weighted score - score = _compute_keyword_score( - shared_kws=shared_kws, - kws_a=new_doc_kws, - kws_b=peer_kws, - weight=_KEYWORD_SCORE_WEIGHT, - ) - if score < _MIN_SCORE_THRESHOLD: - continue - - # Create edge with meaningful weight and metadata - peer_doc_id = peer_node.owner_document_id - edge_pair = tuple(sorted([document_id, peer_doc_id])) - db.add( - GraphEdge( - edge_id=f"related:{edge_pair[0]}<->{edge_pair[1]}", - user_id=user_id, - namespace=namespace, - edge_kind='related', - source_node_id=document_node_id, - target_node_id=peer_node.node_id, - owner_document_id=document_id, - job_result_id=job_result_id, - is_directed=False, - weight=round(score, 4), - properties={ - 'shared_keywords': sorted(shared_kws), - 'connection_count': len(shared_kws), - }, - ) - ) - - db.flush() - logger.info( - f"publish_document_graph: doc={document_id} " - f"keywords={len(top_keywords)} chunks={chunks_count}" - ) - - def remove_document_graph(self, db: Session, *, scope: GraphScope | None, document_id: str) -> None: - document_node_id = f"doc:{document_id}" - edge_delete = delete(GraphEdge).where( - or_( - GraphEdge.owner_document_id == document_id, - GraphEdge.source_node_id == document_node_id, - GraphEdge.target_node_id == document_node_id, - ) - ) - node_delete = delete(GraphNode).where(GraphNode.owner_document_id == document_id) - if scope is not None: - edge_delete = edge_delete.where( - GraphEdge.user_id == scope.user_id, - GraphEdge.namespace == scope.namespace, - ) - node_delete = node_delete.where( - GraphNode.user_id == scope.user_id, - GraphNode.namespace == scope.namespace, - ) - db.execute(edge_delete) - db.execute(node_delete) - db.flush() - - -class GraphQueryService: - """Read-side graph service for document routing before canonical chunk hydration.""" - - async def find_entry_documents( - self, - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - exclude_document_ids: Iterable[str] = (), - exclude_sections: Iterable[dict[str, str]] = (), - ) -> list[str]: - query_lc = query.lower().strip() - exclude_document_ids = set(exclude_document_ids) - - if query_lc: - like = f'%{query_lc}%' - stmt = ( - select(DocumentSection.document_id) - .join(Document, (Document.document_id == DocumentSection.document_id) & (Document.current_job_result_id == DocumentSection.job_result_id)) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where( - DocumentSection.section_title.ilike(like) - | DocumentSection.section_path.ilike(like) - ) - .distinct() - ) - if exclude_document_ids: - stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) - for exc in (exclude_sections or ()): - if not isinstance(exc, dict): - continue - exc_doc = str(exc.get('document_id') or '').strip() - exc_path = str(exc.get('section_path') or '').strip() - if exc_doc and exc_path: - stmt = stmt.where( - ~((DocumentSection.document_id == exc_doc) & ( - (DocumentSection.section_path == exc_path) | - DocumentSection.section_path.like(f'{exc_path} / %') - )) - ) - result = await db.execute(stmt) - seen = [row[0] for row in result.all()] - if seen: - return seen - - return await self._find_documents_by_content( - db, - user_id=user_id, - namespace=namespace, - query=query_lc, - exclude_document_ids=exclude_document_ids, - ) - - async def _find_documents_by_content( - self, - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - exclude_document_ids: set[str], - ) -> list[str]: - like = f'%{query}%' - stmt = ( - select(Document.document_id) - .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) & (DocumentChunk.job_result_id == Document.current_job_result_id)) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where(DocumentChunk.content_lexical_text.ilike(like)) - ) - if exclude_document_ids: - stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) - result = await db.execute(stmt) - seen: list[str] = [] - for (doc_id,) in result.all(): - if doc_id and doc_id not in seen: - seen.append(doc_id) - return seen - - async def collect_candidate_chunks( - self, - db: AsyncSession, - *, - user_id: str, - namespace: str, - entry_document_ids: Sequence[str], - query: str, - top_k: int, - exclude_sections: Iterable[dict[str, str]] = (), - ) -> list[dict[str, Any]]: - if not entry_document_ids: - return [] - page_size = top_k - if exclude_sections: - page_size = max(top_k, top_k * _SECTION_EXCLUSION_PAGE_MULTIPLIER) - base_stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) & (DocumentChunk.job_result_id == Document.current_job_result_id)) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where(Document.document_id.in_(list(entry_document_ids))) - .where(_build_lexical_match_predicate(query)) - .order_by(DocumentChunk.sort_order) - ) - rows = [] - offset = 0 - while len(rows) < top_k: - result = await db.execute(base_stmt.limit(page_size).offset(offset)) - result_rows = result.all() - if not result_rows: - break - for document, chunk, section, job_result in result_rows: - section_path = section.section_path if section else None - if is_excluded_section( - document_id=document.document_id, - section_path=section_path, - exclude_sections=exclude_sections, - ): - continue - rows.append({ - 'document_id': document.document_id, - 'chunk_id': chunk.chunk_id, - 'section_id': chunk.section_id, - 'section_path': section_path, - 'source_file_name': document.source_file_name, - 'chunk_type': chunk.chunk_type, - 'content': chunk.content, - 'score': 2.0, - 'file_path': chunk.file_path, - 'chunk_metadata': chunk.chunk_metadata or {}, - 'job_result_id': chunk.job_result_id, - 'job_id': job_result.job_id if job_result else None, - }) - if len(rows) >= top_k: - break - if len(result_rows) < page_size: - break - offset += page_size - return rows diff --git a/packages/shared-python/shared/services/retrieval/hydration/__init__.py b/packages/shared-python/shared/services/retrieval/hydration/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/hydration/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/hydration/assets.py b/packages/shared-python/shared/services/retrieval/hydration/assets.py new file mode 100644 index 000000000..99d1bb4a2 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/hydration/assets.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +from typing import Any + +from loguru import logger + +from shared.services.retrieval.hydration.row_utils import MEDIA_CHUNK_TYPES, normalize_chunk_type +from shared.services.storage.result_storage import get_result_storage + + +def _normalize_artifact_ref(asset_ref: object) -> str | None: + return get_result_storage().normalize_artifact_ref( + None if asset_ref is None else str(asset_ref) + ) + + +def _is_retrieval_media_row(row: dict[str, Any]) -> bool: + raw_chunk_type = row.get("chunk_type") or row.get("type") + return normalize_chunk_type(raw_chunk_type) in MEDIA_CHUNK_TYPES + + +def _resolve_asset_request(row: dict[str, Any]) -> tuple[str, str] | None: + job_id = str(row.get("job_id") or "").strip() + if not job_id or not _is_retrieval_media_row(row): + return None + + artifact_ref = _normalize_artifact_ref(row.get("file_path")) + if artifact_ref is None: + return None + + return job_id, artifact_ref + + +async def _generate_retrieval_asset_url( + *, + row: dict[str, Any], + log_context: str, +) -> str | None: + request = _resolve_asset_request(row) + if request is None: + return None + + job_id, artifact_ref = request + try: + return get_result_storage().generate_artifact_url( + job_id=job_id, + artifact_ref=artifact_ref, + ) + except Exception as exc: + logger.warning(f"Failed to generate {log_context} asset URL (ignored): {exc}") + return None + + +async def enrich_rows_with_retrieval_asset_urls( + rows: list[dict[str, Any]], + *, + log_context: str, +) -> list[dict[str, Any]]: + enriched_rows: list[dict[str, Any]] = [] + for row in rows: + enriched = dict(row) + asset_url = await _generate_retrieval_asset_url( + row=row, + log_context=log_context, + ) + if asset_url: + enriched["asset_url"] = asset_url + enriched_rows.append(enriched) + return enriched_rows + + +async def build_retrieval_asset_url_map( + rows: list[dict[str, Any]], + *, + log_context: str, +) -> dict[str, str]: + url_map: dict[str, str] = {} + for row in rows: + chunk_id = str(row.get("chunk_id") or "").strip() + if not chunk_id: + continue + + asset_url = await _generate_retrieval_asset_url( + row=row, + log_context=log_context, + ) + if asset_url: + url_map[chunk_id] = asset_url + return url_map + + +def is_client_result_artifact_ref(asset_ref: str | None) -> bool: + return _normalize_artifact_ref(asset_ref) is not None diff --git a/packages/shared-python/shared/services/retrieval/hydration/connected.py b/packages/shared-python/shared/services/retrieval/hydration/connected.py new file mode 100644 index 000000000..f7c299b77 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/hydration/connected.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +from typing import Any + +from sqlalchemy import and_, or_, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.job_result import JobResult +from shared.services.retrieval.hydration.row_utils import ( + filter_excluded_rows, + iter_connected_target_ids, + normalize_chunk_type, +) + + +async def hydrate_connected_target_rows( + *, + db: AsyncSession | None, + rows: list[dict[str, Any]], + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], +) -> list[dict[str, Any]]: + if db is None: + return [] + + existing_chunk_ids = { + str(row.get('chunk_id') or '').strip() + for row in rows + if row.get('chunk_id') + } + target_ids_by_revision: dict[tuple[str, str], set[str]] = {} + for row in rows: + if normalize_chunk_type(row.get('chunk_type')) != 'text': + continue + document_id = str(row.get('document_id') or '').strip() + job_result_id = str(row.get('job_result_id') or '').strip() + if not document_id or not job_result_id: + continue + for target_id in iter_connected_target_ids(row): + if target_id in existing_chunk_ids: + continue + target_ids_by_revision.setdefault((document_id, job_result_id), set()).add( + target_id + ) + + if not target_ids_by_revision: + return [] + + revision_filters = [ + and_( + DocumentChunk.document_id == document_id, + DocumentChunk.job_result_id == job_result_id, + DocumentChunk.chunk_id.in_(sorted(target_ids)), + ) + for (document_id, job_result_id), target_ids in target_ids_by_revision.items() + if target_ids + ] + if not revision_filters: + return [] + + stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join(DocumentChunk, DocumentChunk.document_id == Document.document_id) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .join(JobResult, JobResult.id == DocumentChunk.job_result_id) + .where(or_(*revision_filters)) + .order_by(DocumentChunk.sort_order) + ) + result = await db.execute(stmt) + + hydrated_rows: list[dict[str, Any]] = [] + for document, chunk, section, job_result in result.all(): + section_path = section.section_path if section else None + hydrated_rows.append( + { + 'document_id': document.document_id, + 'chunk_id': chunk.chunk_id, + 'section_id': chunk.section_id, + 'section_path': section_path, + 'source_file_name': document.source_file_name, + 'chunk_type': chunk.chunk_type, + 'content': chunk.content, + 'score': 0.0, + 'file_path': chunk.file_path, + 'chunk_metadata': chunk.chunk_metadata or {}, + 'job_result_id': chunk.job_result_id, + 'job_id': job_result.job_id if job_result else None, + 'sort_order': chunk.sort_order, + } + ) + + return filter_excluded_rows( + hydrated_rows, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + ) diff --git a/packages/shared-python/shared/services/retrieval/hydration/path.py b/packages/shared-python/shared/services/retrieval/hydration/path.py new file mode 100644 index 000000000..de2ec15a2 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/hydration/path.py @@ -0,0 +1,271 @@ +from __future__ import annotations + +from typing import Any + +from loguru import logger +from sqlalchemy import or_, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.job_result import JobResult +from shared.services.retrieval.search.lexical_text import normalize_section_path +from shared.services.retrieval.search.scoring import get_row_path + + +async def hydrate_paths_to_rows( + db: AsyncSession, + *, + path_selections: list[dict[str, Any]], + user_id: str, + namespace: str, + document_id: str | None = None, +) -> list[dict[str, Any]]: + """Load full chunk rows by section_path or source_chunk_path.""" + if not path_selections: + return [] + + confidence_by_path: dict[str, float] = {} + mode_by_path: dict[str, str] = {} + ordered_paths: list[str] = [] + for item in path_selections: + raw_path = str(item.get('path') or '').strip() + path = normalize_section_path(raw_path) if raw_path and '/' in raw_path else raw_path + if not path: + continue + confidence = float(item.get('confidence', 0.0) or 0.0) + hydrate_mode = str(item.get('hydrate_mode') or 'chunks').strip().lower() + if path not in confidence_by_path: + ordered_paths.append(path) + confidence_by_path[path] = confidence + mode_by_path[path] = hydrate_mode + else: + confidence_by_path[path] = max(confidence_by_path[path], confidence) + if not ordered_paths: + return [] + + outline_paths = [path for path in ordered_paths if mode_by_path.get(path) == 'outline'] + chunk_paths = [path for path in ordered_paths if mode_by_path.get(path) != 'outline'] + + rows: list[dict[str, Any]] = [] + + if outline_paths: + rows.extend( + await _hydrate_outline_paths( + db, + outline_paths=outline_paths, + confidence_by_path=confidence_by_path, + user_id=user_id, + namespace=namespace, + document_id=document_id, + ) + ) + + if chunk_paths: + rows.extend( + await _hydrate_chunk_paths( + db, + chunk_paths=chunk_paths, + confidence_by_path=confidence_by_path, + mode_by_path=mode_by_path, + user_id=user_id, + namespace=namespace, + document_id=document_id, + ) + ) + + _sort_rows_by_selection_order(rows, ordered_paths) + _log_hydration_resolution(rows=rows, ordered_paths=ordered_paths, outline_paths=outline_paths) + return rows + + +async def _hydrate_outline_paths( + db: AsyncSession, + *, + outline_paths: list[str], + confidence_by_path: dict[str, float], + user_id: str, + namespace: str, + document_id: str | None, +) -> list[dict[str, Any]]: + outline_section_filters = [ + DocumentSection.section_path == path + for path in outline_paths + ] + outline_stmt = ( + select(Document, DocumentSection) + .join( + DocumentSection, + (DocumentSection.document_id == Document.document_id) + & (DocumentSection.job_result_id == Document.current_job_result_id), + ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where(or_(*outline_section_filters)) + ) + if document_id: + outline_stmt = outline_stmt.where(Document.document_id == document_id) + + rows: list[dict[str, Any]] = [] + outline_result = await db.execute(outline_stmt) + for document, section in outline_result.all(): + agent_score = confidence_by_path.get(section.section_path, 0.0) + summary_text = (section.summary or '').strip() + title_text = (section.section_title or '').strip() + content = f'[Outline] {title_text}' + if summary_text: + content += f'\n{summary_text}' + rows.append({ + 'document_id': document.document_id, + 'chunk_id': f'outline_{section.section_id}', + 'section_id': section.section_id, + 'section_path': section.section_path, + 'source_file_name': document.source_file_name, + 'chunk_type': 'outline', + 'content': content, + 'score': agent_score, + 'agent_score': agent_score, + 'file_path': None, + 'chunk_metadata': {}, + 'job_result_id': section.job_result_id, + 'job_id': None, + 'source_chunk_path': None, + 'sort_order': section.sort_order, + 'hydrate_mode': 'outline', + }) + return rows + + +async def _hydrate_chunk_paths( + db: AsyncSession, + *, + chunk_paths: list[str], + confidence_by_path: dict[str, float], + mode_by_path: dict[str, str], + user_id: str, + namespace: str, + document_id: str | None, +) -> list[dict[str, Any]]: + section_path_filters = [] + self_only_paths = {path for path in chunk_paths if mode_by_path.get(path) == 'self_only'} + for path in chunk_paths: + section_path_filters.append(DocumentSection.section_path == path) + if path not in self_only_paths: + section_path_filters.append(DocumentSection.section_path.like(f'{path} / %')) + + stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .join(JobResult, JobResult.id == DocumentChunk.job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where( + or_( + *section_path_filters, + DocumentChunk.source_chunk_path.in_(chunk_paths), + ) + ) + ) + if document_id: + stmt = stmt.where(Document.document_id == document_id) + result = await db.execute(stmt) + + rows: list[dict[str, Any]] = [] + seen_paths: set[str] = set() + for document, chunk, section, job_result in result.all(): + row_path = (section.section_path if section else None) or chunk.source_chunk_path or '' + if row_path in seen_paths: + continue + + matched_path = row_path + if section and section.section_path not in confidence_by_path: + matched_path = next( + ( + path for path in chunk_paths + if section.section_path == path + or section.section_path.startswith(f'{path} / ') + ), + row_path, + ) + + path_mode = mode_by_path.get(matched_path, 'chunks') + allowed_types = _get_allowed_types_for_mode(path_mode) + if allowed_types is not None: + chunk_type_lower = (chunk.chunk_type or '').strip().lower() + if chunk_type_lower not in allowed_types: + continue + + seen_paths.add(row_path) + agent_score = confidence_by_path.get(matched_path, 0.0) + rows.append({ + 'document_id': document.document_id, + 'chunk_id': chunk.chunk_id, + 'section_id': chunk.section_id, + 'section_path': section.section_path if section else None, + 'source_file_name': document.source_file_name, + 'chunk_type': chunk.chunk_type, + 'content': chunk.content, + 'score': agent_score, + 'agent_score': agent_score, + 'file_path': chunk.file_path, + 'chunk_metadata': chunk.chunk_metadata or {}, + 'job_result_id': chunk.job_result_id, + 'job_id': job_result.job_id if job_result else None, + 'source_chunk_path': chunk.source_chunk_path, + 'sort_order': chunk.sort_order, + 'hydrate_mode': path_mode, + }) + return rows + + +def _get_allowed_types_for_mode(path_mode: str) -> set[str] | None: + mode_allowed_types: dict[str, set[str] | None] = { + 'chunks': None, + 'self_only': None, + 'assets_only': {'image', 'table'}, + 'image_only': {'image'}, + 'table_only': {'table'}, + } + return mode_allowed_types.get(path_mode) + + +def _sort_rows_by_selection_order(rows: list[dict[str, Any]], ordered_paths: list[str]) -> None: + path_order = {path: index for index, path in enumerate(ordered_paths)} + + def row_sort_key(row: dict[str, Any]) -> int: + row_path = get_row_path(row) + if row_path in path_order: + return path_order[row_path] + for path, index in path_order.items(): + if row_path.startswith(f'{path} / '): + return index + return 10**9 + + rows.sort(key=row_sort_key) + + +def _log_hydration_resolution( + *, rows: list[dict[str, Any]], ordered_paths: list[str], outline_paths: list[str] +) -> None: + hydrated_paths = {get_row_path(row) for row in rows} + resolved_inputs = { + path for path in ordered_paths + if path in hydrated_paths + or any(row_path.startswith(f'{path} / ') for row_path in hydrated_paths) + } + resolved_inputs |= set(outline_paths) + missed = len(ordered_paths) - len(resolved_inputs) + if missed > 0: + missing_paths = [path for path in ordered_paths if path not in resolved_inputs] + logger.warning( + f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved (missed={missed}); ' + f'missing[:5]={missing_paths[:5]}' + ) + else: + logger.info(f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved') diff --git a/packages/shared-python/shared/services/retrieval/hydration/reference.py b/packages/shared-python/shared/services/retrieval/hydration/reference.py new file mode 100644 index 000000000..5c6a0556e --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/hydration/reference.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +from typing import Any + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.job_result import JobResult +from shared.services.retrieval.hydration.row_utils import ( + ReferenceLookupKey, + build_reference_lookup_key, +) + + +async def hydrate_referenced_chunk_rows( + *, + db: AsyncSession | None, + user_id: str, + namespace: str, + refs: list[dict[str, Any]], +) -> list[dict[str, Any]]: + if db is None or not refs: + return [] + + ref_keys = [ + build_reference_lookup_key( + document_id=ref.get('document_id'), + chunk_id=ref.get('chunk_id'), + section_path=ref.get('section_path'), + file_path=ref.get('file_path'), + ) + for ref in refs + ] + ref_keys = [key for key in ref_keys if key[0] and key[1]] + if not ref_keys: + return [] + + document_ids = sorted({document_id for document_id, _, _, _ in ref_keys}) + chunk_ids = sorted({chunk_id for _, chunk_id, _, _ in ref_keys}) + stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .join(JobResult, JobResult.id == DocumentChunk.job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where(Document.document_id.in_(document_ids)) + .where(DocumentChunk.chunk_id.in_(chunk_ids)) + .order_by(DocumentChunk.sort_order) + ) + result = await db.execute(stmt) + + rows_by_key: dict[ReferenceLookupKey, dict[str, Any]] = {} + rows_by_base_key: dict[tuple[str, str], list[dict[str, Any]]] = {} + for document, chunk, section, job_result in result.all(): + row = { + 'document_id': document.document_id, + 'chunk_id': chunk.chunk_id, + 'section_id': chunk.section_id, + 'section_path': section.section_path if section else None, + 'source_file_name': document.source_file_name, + 'chunk_type': chunk.chunk_type, + 'content': chunk.content, + 'score': 1.0, + 'file_path': chunk.file_path, + 'chunk_metadata': chunk.chunk_metadata or {}, + 'job_result_id': chunk.job_result_id, + 'job_id': job_result.job_id if job_result else None, + 'source_chunk_path': chunk.source_chunk_path, + 'sort_order': chunk.sort_order, + } + key = build_reference_lookup_key( + document_id=row['document_id'], + chunk_id=row['chunk_id'], + section_path=row['section_path'], + file_path=row['file_path'], + ) + rows_by_key[key] = row + rows_by_base_key.setdefault((key[0], key[1]), []).append(row) + + rows: list[dict[str, Any]] = [] + seen_keys: set[ReferenceLookupKey] = set() + for key in ref_keys: + row = rows_by_key.get(key) + if row is None: + candidates = rows_by_base_key.get((key[0], key[1]), []) + row = next( + ( + candidate + for candidate in candidates + if key[2] + and str(candidate.get('section_path') or '').strip() == key[2] + ), + None, + ) + if row is None: + row = next( + ( + candidate + for candidate in candidates + if build_reference_lookup_key( + document_id=candidate.get('document_id'), + chunk_id=candidate.get('chunk_id'), + section_path=candidate.get('section_path'), + file_path=candidate.get('file_path'), + ) + not in seen_keys + ), + None, + ) + if row is not None: + row_key = build_reference_lookup_key( + document_id=row.get('document_id'), + chunk_id=row.get('chunk_id'), + section_path=row.get('section_path'), + file_path=row.get('file_path'), + ) + if row_key in seen_keys: + continue + seen_keys.add(row_key) + rows.append(row) + return rows diff --git a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py new file mode 100644 index 000000000..17a26a43a --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows +from shared.services.retrieval.hydration.row_utils import ( + clean_content, + filter_excluded_rows, + iter_connected_target_ids, + normalize_chunk_type, +) + + +async def assemble_retrieval_results( + *, + db: AsyncSession | None = None, + rows: list[dict[str, Any]], + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + allowed_chunk_types: set[str] | None = None, +) -> list[dict[str, Any]]: + filtered_rows = filter_excluded_rows( + rows, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + ) + if allowed_chunk_types is not None: + filtered_rows = [ + row for row in filtered_rows + if normalize_chunk_type(row.get('chunk_type')) in allowed_chunk_types + ] + hydrated_rows = await hydrate_connected_target_rows( + db=db, + rows=filtered_rows, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + ) + rows_by_chunk_id = { + str(row.get('chunk_id') or ''): row + for row in [*filtered_rows, *hydrated_rows] + if row.get('chunk_id') + } + + embedded_targets: set[str] = set() + for row in filtered_rows: + for target_id in iter_connected_target_ids(row): + if target_id in rows_by_chunk_id: + embedded_targets.add(target_id) + + assembled: list[dict[str, Any]] = [] + for row in filtered_rows: + if row.get('chunk_id') in embedded_targets: + continue + assembled_row = dict(row) + base_content = str(row.get('content') or '') + if normalize_chunk_type(row.get('chunk_type')) == 'text': + connected_targets: list[tuple[int, str]] = [] + for target_id in iter_connected_target_ids(row): + target_row = rows_by_chunk_id.get(target_id) + if not target_row: + continue + if normalize_chunk_type(target_row.get('chunk_type')) != 'table': + continue + target_content = str(target_row.get('content') or '').strip() + if target_content: + sort_key = int(target_row.get('sort_order', 0) or 0) + connected_targets.append((sort_key, target_content)) + connected_targets.sort(key=lambda item: item[0]) + related_parts = [content for _, content in connected_targets] + if base_content and related_parts: + assembled_row['content'] = '\n\n'.join([base_content, *related_parts]) + else: + assembled_row['content'] = base_content + else: + assembled_row['content'] = base_content + assembled_row['content'] = clean_content(assembled_row['content']) + assembled.append(assembled_row) + return assembled diff --git a/packages/shared-python/shared/services/retrieval/hydration/row_utils.py b/packages/shared-python/shared/services/retrieval/hydration/row_utils.py new file mode 100644 index 000000000..d163836ca --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/hydration/row_utils.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +import re +from typing import Any + +from shared.services.retrieval.search.section_filters import is_excluded_section + +MEDIA_CHUNK_TYPES = {'image', 'table'} +PUBLIC_RESULT_FIELDS = { + 'chunk_type', 'content', 'score', 'asset_url', +} +PUBLIC_SOURCE_FIELDS = { + 'document_id', 'source_file_name', 'section_path', +} + +ReferenceLookupKey = tuple[str, str, str, str] + +_PATH_REF_RE = re.compile(r'\[(?:images|tables)/[^\]\n]+\]') + + +def clean_content(content: str) -> str: + return _PATH_REF_RE.sub('', content).strip() + + +def normalize_chunk_type(raw: object) -> str: + return str(raw or '').strip().split('\n', 1)[0].lower() + + +def is_media_chunk(row: dict[str, Any]) -> bool: + return normalize_chunk_type(row.get('chunk_type')) in MEDIA_CHUNK_TYPES + + +def build_reference_lookup_key( + *, + document_id: object, + chunk_id: object, + section_path: object = '', + file_path: object = '', +) -> ReferenceLookupKey: + return ( + str(document_id or '').strip(), + str(chunk_id or '').strip(), + str(section_path or '').strip(), + str(file_path or '').strip(), + ) + + +def filter_excluded_rows( + rows: list[dict[str, Any]], + *, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], +) -> list[dict[str, Any]]: + filtered: list[dict[str, Any]] = [] + excluded_documents = set(exclude_document_ids) + for row in rows: + document_id = row.get('document_id') + if document_id in excluded_documents: + continue + if is_excluded_section( + document_id=document_id, + section_path=row.get('section_path'), + exclude_sections=exclude_sections, + ): + continue + filtered.append(row) + return filtered + + +def iter_connected_target_ids(row: dict[str, Any]) -> list[str]: + metadata = row.get('chunk_metadata') or {} + if not isinstance(metadata, dict): + return [] + + target_ids: list[str] = [] + for item in metadata.get('connect_to') or []: + if not isinstance(item, dict): + continue + target_id = str(item.get('target') or '').strip() + if target_id: + target_ids.append(target_id) + return target_ids diff --git a/packages/shared-python/shared/services/retrieval/llm_adapter.py b/packages/shared-python/shared/services/retrieval/llm_adapter.py index c5e592b5e..9c8bd9263 100644 --- a/packages/shared-python/shared/services/retrieval/llm_adapter.py +++ b/packages/shared-python/shared/services/retrieval/llm_adapter.py @@ -87,7 +87,7 @@ def create_retrieval_llm_fn( effective_model = model or _resolve_default_model() async def llm_fn(prompt: LLMFnInput) -> str: - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client client = get_openai_client(model=effective_model) current_llm_usage.set(None) @@ -118,7 +118,7 @@ def create_retrieval_planner_fn( effective_model = model or _resolve_planner_model(thinking=thinking) async def llm_fn(prompt: LLMFnInput) -> str: - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client client = get_openai_client(model=effective_model) current_llm_usage.set(None) @@ -159,7 +159,7 @@ def create_retrieval_vlm_fn( return None async def vlm_fn(prompt: LLMFnInput) -> str: - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client client = get_openai_client(model=effective_model) current_llm_usage.set(None) diff --git a/packages/shared-python/shared/services/retrieval/publication_content.py b/packages/shared-python/shared/services/retrieval/publication_content.py new file mode 100644 index 000000000..1c5c3b4ef --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/publication_content.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +from typing import Any +from uuid import uuid4 + +from sqlalchemy import delete +from sqlalchemy.orm import Session + +from shared.models.database.document import DocumentChunk, DocumentSection +from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.retrieval.search.lexical_text import ( + build_content_lexical_text, + build_content_search_text, + build_path_lexical_text, + build_path_search_text, + build_term_search_text, + section_path_from_chunk_path, +) + + +def replace_document_revision_content( + db: Session, + *, + scope: DocumentPublicationScope, + chunks: list[dict[str, Any]], +) -> None: + """Replace retrieval sections and chunks for one published document revision.""" + _delete_existing_revision_content(db, scope=scope) + section_publisher = DocumentSectionPublisher(db=db, scope=scope) + for index, chunk in enumerate(chunks): + chunk_metadata = _get_chunk_metadata(chunk) + source_path = _get_source_path(chunk=chunk, chunk_metadata=chunk_metadata) + section_path = section_path_from_chunk_path( + source_path, + source_file_name=scope.source_file_name, + ) + section = section_publisher.ensure_section(section_path) + db.add( + _build_document_chunk( + chunk=chunk, + chunk_metadata=chunk_metadata, + source_path=source_path, + section=section, + scope=scope, + fallback_sort_order=index, + ) + ) + + +class DocumentSectionPublisher: + def __init__(self, *, db: Session, scope: DocumentPublicationScope) -> None: + self._db = db + self._scope = scope + self._sections_by_path: dict[str, DocumentSection] = {} + + def ensure_section(self, section_path: str) -> DocumentSection: + existing_section = self._sections_by_path.get(section_path) + if existing_section is not None: + return existing_section + + path_parts = [part for part in section_path.split(" / ") if part] + for depth in range(1, len(path_parts) + 1): + ancestor_path = " / ".join(path_parts[:depth]) + if ancestor_path in self._sections_by_path: + continue + + ancestor_section = DocumentSection( + user_id=self._scope.user_id, + namespace=self._scope.namespace, + document_id=self._scope.document_id, + job_result_id=self._scope.job_result_id, + parent_section_id=self._get_parent_section_id(path_parts, depth), + section_path=ancestor_path, + section_title=path_parts[depth - 1], + section_level=depth, + section_metadata={}, + sort_order=len(self._sections_by_path), + ) + self._db.add(ancestor_section) + self._db.flush() + self._sections_by_path[ancestor_path] = ancestor_section + + return self._sections_by_path[section_path] + + def _get_parent_section_id( + self, + path_parts: list[str], + depth: int, + ) -> str | None: + if depth <= 1: + return None + parent_path = " / ".join(path_parts[: depth - 1]) + parent = self._sections_by_path.get(parent_path) + return parent.section_id if parent is not None else None + + +def _delete_existing_revision_content( + db: Session, + *, + scope: DocumentPublicationScope, +) -> None: + db.execute( + delete(DocumentChunk) + .where(DocumentChunk.document_id == scope.document_id) + .where(DocumentChunk.job_result_id == scope.job_result_id) + ) + db.execute( + delete(DocumentSection) + .where(DocumentSection.document_id == scope.document_id) + .where(DocumentSection.job_result_id == scope.job_result_id) + ) + + +def _build_document_chunk( + *, + chunk: dict[str, Any], + chunk_metadata: dict[str, Any], + source_path: str | None, + section: DocumentSection, + scope: DocumentPublicationScope, + fallback_sort_order: int, +) -> DocumentChunk: + section_summary = section.summary + section_path = section.section_path + section_title = section.section_title + path_text = f"{scope.source_file_name or ''} {section_path}".strip() + return DocumentChunk( + id=f"dchk_{uuid4().hex[:12]}", + chunk_id=str(chunk.get("chunk_id") or f"chunk_{uuid4().hex[:12]}"), + user_id=scope.user_id, + namespace=scope.namespace, + document_id=scope.document_id, + job_result_id=scope.job_result_id, + section_id=section.section_id, + chunk_type=chunk.get("type") or chunk.get("chunk_type") or "text", + content=chunk.get("content") or chunk.get("text"), + content_lexical_text=build_content_lexical_text(chunk), + path_lexical_text=build_path_lexical_text( + source_path, + source_file_name=scope.source_file_name, + ), + content_search_text=build_content_search_text( + chunk, + section_summary=section_summary, + ), + path_search_text=build_path_search_text( + source_file_name=scope.source_file_name, + section_path=section_path, + section_title=section_title, + section_summary=section_summary, + ), + term_search_text=build_term_search_text(chunk, path_text=path_text), + source_chunk_path=source_path, + file_path=chunk_metadata.get("file_path") or chunk.get("file_path"), + chunk_metadata=chunk_metadata, + sort_order=_get_sort_order(chunk, fallback_sort_order), + ) + + +def _get_chunk_metadata(chunk: dict[str, Any]) -> dict[str, Any]: + metadata = chunk.get("metadata") + return metadata if isinstance(metadata, dict) else {} + + +def _get_source_path( + *, + chunk: dict[str, Any], + chunk_metadata: dict[str, Any], +) -> str | None: + source_path = chunk_metadata.get("path") or chunk.get("path") + return str(source_path) if source_path is not None else None + + +def _get_sort_order(chunk: dict[str, Any], fallback_sort_order: int) -> int: + try: + return int(chunk.get("order", fallback_sort_order)) + except (TypeError, ValueError): + return fallback_sort_order diff --git a/packages/shared-python/shared/services/retrieval/publication_models.py b/packages/shared-python/shared/services/retrieval/publication_models.py new file mode 100644 index 000000000..8c653212c --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/publication_models.py @@ -0,0 +1,26 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class ExistingDocumentScope: + document_id: str + namespace: str + + +@dataclass(frozen=True) +class PublishedDocumentState: + user_id: str + namespace: str + document_id: str | None + skipped_all_duplicate: bool = False + + +@dataclass(frozen=True) +class DocumentPublicationScope: + user_id: str + namespace: str + document_id: str + job_result_id: str + source_file_name: str | None diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index 80aaa4ae0..aa252ef8e 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -9,24 +9,25 @@ from __future__ import annotations from datetime import datetime, timezone -from typing import Any, Dict, List, Optional +from typing import Any from uuid import uuid4 from loguru import logger -from sqlalchemy import delete, select +from sqlalchemy import select from sqlalchemy.orm import Session -from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.document import Document from shared.models.database.job import Job from shared.models.database.job_result import JobResult -from shared.services.retrieval.graph_service import DocumentGraphService, GraphScope -from shared.services.retrieval.lexical_text import ( - build_content_lexical_text, - build_content_search_text, - build_path_lexical_text, - build_path_search_text, - build_term_search_text, - section_path_from_chunk_path, +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace +from shared.services.retrieval.graph.service import DocumentGraphService, GraphScope +from shared.services.retrieval.publication_content import ( + replace_document_revision_content, +) +from shared.services.retrieval.publication_models import ( + DocumentPublicationScope, + ExistingDocumentScope, + PublishedDocumentState, ) @@ -42,7 +43,7 @@ def get_existing_document_scope( db: Session, *, job_id: str, - ) -> Optional[Dict[str, str]]: + ) -> ExistingDocumentScope | None: job = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() if not job: return None @@ -58,7 +59,10 @@ def get_existing_document_scope( if not document: return None - return {"document_id": document.document_id, "namespace": document.namespace} + return ExistingDocumentScope( + document_id=document.document_id, + namespace=document.namespace, + ) def publish_document_state( self, @@ -66,8 +70,8 @@ def publish_document_state( *, job_id: str, job_result_id: str, - chunks: List[Dict[str, Any]], - ) -> Optional[Dict[str, Any]]: + chunks: list[dict[str, Any]], + ) -> PublishedDocumentState | None: job = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() if not job: logger.warning(f"Job not found for document publication: {job_id}") @@ -86,13 +90,15 @@ def _publish_document_state_for_job( *, job: Job, job_result_id: str, - chunks: List[Dict[str, Any]], - ) -> Optional[Dict[str, Any]]: + chunks: list[dict[str, Any]], + ) -> PublishedDocumentState | None: job_metadata = job.job_metadata or {} - namespace = job_metadata.get("namespace") or "default" + namespace = normalize_retrieval_namespace(job_metadata.get("namespace")) document_id = job_metadata.get("document_id") - source_file_name = job_metadata.get("source_file_name") or job_metadata.get("file_name") + source_file_name = job_metadata.get("source_file_name") or job_metadata.get( + "file_name" + ) deduped_chunks = chunks @@ -102,14 +108,60 @@ def _publish_document_state_for_job( f"⏭️ All chunks are duplicates of existing documents. " f"Skipping document creation for job_id={job.job_id}." ) - return { - "user_id": str(job.user_id), - "namespace": namespace, - "document_id": None, - "skipped_all_duplicate": True, - } - - # ── Document upsert (original logic, but only for deduped chunks) ── + return PublishedDocumentState( + user_id=str(job.user_id), + namespace=namespace, + document_id=None, + skipped_all_duplicate=True, + ) + + document = self._upsert_document_revision( + db, + job=job, + job_result_id=job_result_id, + document_id=str(document_id) if document_id else None, + namespace=namespace, + source_file_name=str(source_file_name) if source_file_name else None, + ) + if document is None: + return None + + self._bind_job_result_document( + db, + job_result_id=job_result_id, + document_id=document.document_id, + ) + namespace = normalize_retrieval_namespace(namespace or document.namespace) + scope = DocumentPublicationScope( + user_id=str(job.user_id), + namespace=namespace, + document_id=document.document_id, + job_result_id=job_result_id, + source_file_name=str(source_file_name) if source_file_name else None, + ) + replace_document_revision_content( + db, + scope=scope, + chunks=deduped_chunks, + ) + + db.flush() + return PublishedDocumentState( + user_id=str(job.user_id), + namespace=namespace, + document_id=document.document_id, + ) + + def _upsert_document_revision( + self, + db: Session, + *, + job: Job, + job_result_id: str, + document_id: str | None, + namespace: str, + source_file_name: str | None, + ) -> Document | None: document = None if document_id: document = db.execute( @@ -132,14 +184,14 @@ def _publish_document_state_for_job( ) db.add(document) else: - namespace = namespace or document.namespace if self._is_stale_document_completion( db, document=document, job=job, ): logger.warning( - f"Skipping stale document publication: job_id={job.job_id}, document_id={document.document_id}" + "Skipping stale document publication: " + f"job_id={job.job_id}, document_id={document.document_id}" ) return None document.status = "active" @@ -149,103 +201,19 @@ def _publish_document_state_for_job( document.updated_at = utc_now_naive() db.flush() - document_id = document.document_id + return document + + def _bind_job_result_document( + self, + db: Session, + *, + job_result_id: str, + document_id: str, + ) -> None: result = db.execute(select(JobResult).where(JobResult.id == job_result_id)) job_result = result.scalar_one_or_none() if job_result: job_result.document_id = document_id - namespace = namespace or document.namespace - - db.execute( - delete(DocumentChunk) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - ) - db.execute( - delete(DocumentSection) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - ) - - # ── Insert only deduped (non-duplicate) chunks ────────────────── - sections_by_path: Dict[str, DocumentSection] = {} - for index, chunk in enumerate(deduped_chunks): - chunk_metadata = chunk.get("metadata") or {} - source_path = chunk_metadata.get("path") or chunk.get("path") - section_path = section_path_from_chunk_path(source_path) - section = sections_by_path.get(section_path) - if section is None: - path_parts = [p for p in section_path.split(" / ") if p] - # Ensure all ancestor sections exist (top-down) - for depth in range(1, len(path_parts) + 1): - ancestor_path = " / ".join(path_parts[:depth]) - if ancestor_path in sections_by_path: - continue - ancestor_parent_id = None - if depth > 1: - parent_path = " / ".join(path_parts[:depth - 1]) - parent = sections_by_path.get(parent_path) - if parent is not None: - ancestor_parent_id = parent.section_id - ancestor_section = DocumentSection( - user_id=str(job.user_id), - namespace=namespace, - document_id=document_id, - job_result_id=job_result_id, - parent_section_id=ancestor_parent_id, - section_path=ancestor_path, - section_title=path_parts[depth - 1], - section_level=depth, - section_metadata={}, - sort_order=len(sections_by_path), - ) - db.add(ancestor_section) - db.flush() - sections_by_path[ancestor_path] = ancestor_section - section = sections_by_path[section_path] - - chunk_id = chunk.get("chunk_id") or f"chunk_{uuid4().hex[:12]}" - section_summary = section.summary if section else None - section_path_str = section.section_path if section else "Root" - section_title_str = section.section_title if section else None - path_text = f"{source_file_name or ''} {section_path_str}".strip() - - db.add( - DocumentChunk( - id=f"dchk_{uuid4().hex[:12]}", - chunk_id=chunk_id, - user_id=str(job.user_id), - namespace=namespace, - document_id=document_id, - job_result_id=job_result_id, - section_id=section.section_id, - chunk_type=chunk.get("type") or chunk.get("chunk_type") or "text", - content=chunk.get("content") or chunk.get("text"), - content_lexical_text=build_content_lexical_text(chunk), - path_lexical_text=build_path_lexical_text(source_path), - content_search_text=build_content_search_text( - chunk, section_summary=section_summary - ), - path_search_text=build_path_search_text( - source_file_name=source_file_name, - section_path=section_path_str, - section_title=section_title_str, - section_summary=section_summary, - ), - term_search_text=build_term_search_text(chunk, path_text=path_text), - source_chunk_path=source_path, - file_path=chunk_metadata.get("file_path") or chunk.get("file_path"), - chunk_metadata=chunk_metadata, - sort_order=chunk.get("order", index), - ) - ) - - db.flush() - return { - "user_id": str(job.user_id), - "namespace": namespace, - "document_id": document_id, - } def publish_document_graph( self, @@ -269,7 +237,7 @@ def _publish_document_graph_for_job( ) -> None: metadata = job.job_metadata or {} - namespace = metadata.get("namespace") or "default" + namespace = normalize_retrieval_namespace(metadata.get("namespace")) document_id = metadata.get("document_id") if not document_id: document = db.execute( diff --git a/packages/shared-python/shared/services/retrieval/search/__init__.py b/packages/shared-python/shared/services/retrieval/search/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/search/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/channels.py b/packages/shared-python/shared/services/retrieval/search/channels.py similarity index 82% rename from packages/shared-python/shared/services/retrieval/channels.py rename to packages/shared-python/shared/services/retrieval/search/channels.py index a597307ca..bd881769e 100644 --- a/packages/shared-python/shared/services/retrieval/channels.py +++ b/packages/shared-python/shared/services/retrieval/search/channels.py @@ -8,12 +8,14 @@ from typing import Any -from loguru import logger from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.graph_service import is_excluded_section -from shared.utils.text_utils import tokenize_for_retrieval +from shared.services.retrieval.search.lexical_ranker import ( + rank_rows_by_bm25, + tokenize_query_for_ranker, +) +from shared.services.retrieval.search.section_filters import is_excluded_section _SCOPED_CORPUS_CTE = """ @@ -193,56 +195,6 @@ async def content_channel( ) -def _tokenize_query(query: str) -> list[str]: - return tokenize_for_retrieval(query, dedupe=True) - - -def _bm25_rerank( - rows: list[dict[str, Any]], - query_tokens: list[str], - *, - search_field: str, -) -> list[dict[str, Any]]: - """Rank matching rows with BM25 over pre-tokenized search text.""" - try: - from rank_bm25 import BM25Okapi - except ImportError: - logger.warning("rank_bm25 not installed, skipping BM25 re-rank") - ranked_rows: list[dict[str, Any]] = [] - query_token_set = set(query_tokens) - for row in rows: - tokens = [token for token in str(row.get(search_field) or "").split() if token] - overlap = len(query_token_set.intersection(tokens)) - if overlap <= 0: - continue - row["score"] = float(overlap) - ranked_rows.append(row) - ranked_rows.sort(key=lambda r: r["score"], reverse=True) - return ranked_rows - - corpus: list[list[str]] = [] - ranked_rows: list[dict[str, Any]] = [] - query_token_set = set(query_tokens) - for row in rows: - tokens = [token for token in str(row.get(search_field) or "").split() if token] - if not tokens or not query_token_set.intersection(tokens): - continue - corpus.append(tokens) - ranked_rows.append(row) - - if not corpus or not query_tokens: - return [] - - bm25 = BM25Okapi(corpus) - scores = bm25.get_scores(query_tokens) - - for i, row in enumerate(ranked_rows): - row["score"] = float(scores[i]) - - ranked_rows.sort(key=lambda r: r["score"], reverse=True) - return ranked_rows - - async def _bm25_channel( db: AsyncSession, *, @@ -260,7 +212,7 @@ async def _bm25_channel( if search_field not in {"content_search_text", "path_search_text"}: raise ValueError(f"Unsupported search_field: {search_field}") - query_tokens = _tokenize_query(query) + query_tokens = tokenize_query_for_ranker(query) if not query_tokens: return [] @@ -287,7 +239,7 @@ async def _bm25_channel( rows = [_row_to_dict(r) for r in result.all()] rows = _filter_excluded_sections(rows, exclude_sections) - ranked_rows = _bm25_rerank(rows, query_tokens, search_field=search_field) + ranked_rows = rank_rows_by_bm25(rows, query_tokens, search_field=search_field) return ranked_rows[:top_k] @@ -312,7 +264,7 @@ async def term_channel( Note: top_k is already effective_recall_k from app_service. """ query_lower = query.lower().strip() - query_tokens = _tokenize_query(query) + query_tokens = tokenize_query_for_ranker(query) if not query_lower or not query_tokens: return [] diff --git a/packages/shared-python/shared/services/retrieval/search/lexical_ranker.py b/packages/shared-python/shared/services/retrieval/search/lexical_ranker.py new file mode 100644 index 000000000..8b50cd404 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/search/lexical_ranker.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +from typing import Any + +from loguru import logger + +from shared.utils.text_utils import tokenize_for_retrieval + + +def tokenize_query_for_ranker(query: str) -> list[str]: + return tokenize_for_retrieval(query, dedupe=True) + + +def rank_rows_by_bm25( + rows: list[dict[str, Any]], + query_tokens: list[str], + *, + search_field: str, +) -> list[dict[str, Any]]: + """Rank matching rows with BM25 over pre-tokenized search text.""" + try: + from rank_bm25 import BM25Okapi + except ImportError: + return _rank_rows_by_token_overlap( + rows, + query_tokens, + search_field=search_field, + ) + + corpus: list[list[str]] = [] + ranked_rows: list[dict[str, Any]] = [] + query_token_set = set(query_tokens) + for row in rows: + tokens = _get_search_tokens(row, search_field=search_field) + if not tokens or not query_token_set.intersection(tokens): + continue + corpus.append(tokens) + ranked_rows.append(row) + + if not corpus or not query_tokens: + return [] + + bm25 = BM25Okapi(corpus) + scores = bm25.get_scores(query_tokens) + + for index, row in enumerate(ranked_rows): + row["score"] = float(scores[index]) + + ranked_rows.sort(key=lambda row: row["score"], reverse=True) + return ranked_rows + + +def _rank_rows_by_token_overlap( + rows: list[dict[str, Any]], + query_tokens: list[str], + *, + search_field: str, +) -> list[dict[str, Any]]: + logger.warning("rank_bm25 not installed, skipping BM25 re-rank") + ranked_rows: list[dict[str, Any]] = [] + query_token_set = set(query_tokens) + for row in rows: + tokens = _get_search_tokens(row, search_field=search_field) + overlap = len(query_token_set.intersection(tokens)) + if overlap <= 0: + continue + row["score"] = float(overlap) + ranked_rows.append(row) + ranked_rows.sort(key=lambda row: row["score"], reverse=True) + return ranked_rows + + +def _get_search_tokens(row: dict[str, Any], *, search_field: str) -> list[str]: + return [token for token in str(row.get(search_field) or "").split() if token] diff --git a/packages/shared-python/shared/services/retrieval/lexical_text.py b/packages/shared-python/shared/services/retrieval/search/lexical_text.py similarity index 85% rename from packages/shared-python/shared/services/retrieval/lexical_text.py rename to packages/shared-python/shared/services/retrieval/search/lexical_text.py index 42d15366f..b6cbc9c8c 100644 --- a/packages/shared-python/shared/services/retrieval/lexical_text.py +++ b/packages/shared-python/shared/services/retrieval/search/lexical_text.py @@ -6,6 +6,7 @@ from typing import Any, Optional +from shared.services.chunks.document_path import split_document_path from shared.utils.text_utils import tokenize_contents_for_retrieval @@ -57,23 +58,36 @@ def build_content_lexical_text(chunk: dict[str, Any]) -> Optional[str]: return "\n".join(lexical_parts) if lexical_parts else content -def section_path_from_chunk_path(source_path: Optional[str]) -> str: +def section_path_from_chunk_path( + source_path: Optional[str], + *, + source_file_name: str | None = None, +) -> str: """Extract section hierarchy from chunk path. - Expected format: "/.ext/
//..." + Expected format: ".ext/
//..." Returns " / "-joined section parts, or "Root" if no section hierarchy. """ if not source_path: return "Root" - parts = split_section_path(source_path) - section_parts = parts[2:] # skip kb_root + filename + _, section_parts = split_document_path( + source_path, + source_file_name=source_file_name, + ) if not section_parts: return "Root" - return normalize_section_path(" / ".join(section_parts)) + return " / ".join(section_parts) -def build_path_lexical_text(source_path: Optional[str]) -> Optional[str]: - section_path = section_path_from_chunk_path(source_path) +def build_path_lexical_text( + source_path: Optional[str], + *, + source_file_name: str | None = None, +) -> Optional[str]: + section_path = section_path_from_chunk_path( + source_path, + source_file_name=source_file_name, + ) if not section_path: return None normalized_path = section_path.replace(" / ", " ") diff --git a/packages/shared-python/shared/services/retrieval/search/ranking.py b/packages/shared-python/shared/services/retrieval/search/ranking.py new file mode 100644 index 000000000..392aa1f13 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/search/ranking.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +import math +from typing import Any + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import RetrievalHitStat +from shared.services.retrieval.stats.service import compute_importance_score +from shared.services.retrieval.search.scoring import get_row_path + + +def get_candidate_key(row: dict[str, Any]) -> str: + path = get_row_path(row) + if path: + return f'path:{path}' + chunk_id = str(row.get('chunk_id') or '').strip() + return f'chunk:{chunk_id}' if chunk_id else '' + + +async def load_chunk_importance_scores( + db: AsyncSession, + *, + user_id: str, + namespace: str, + rows: list[dict[str, Any]], +) -> dict[str, float]: + chunk_ids = sorted({ + str(row.get('chunk_id') or '').strip() + for row in rows + if row.get('chunk_id') + }) + if not chunk_ids: + return {} + stmt = ( + select( + RetrievalHitStat.chunk_id, + RetrievalHitStat.hit_count, + RetrievalHitStat.last_hit_at, + RetrievalHitStat.created_at, + ) + .where(RetrievalHitStat.user_id == user_id) + .where(RetrievalHitStat.namespace == namespace) + .where(RetrievalHitStat.hit_kind == 'chunk') + .where(RetrievalHitStat.chunk_id.in_(chunk_ids)) + ) + result = await db.execute(stmt) + importance_scores: dict[str, float] = {} + for chunk_id, hit_count, last_hit_at, created_at in result.all(): + if not chunk_id: + continue + importance_scores[str(chunk_id)] = compute_importance_score(hit_count, last_hit_at, created_at) + return importance_scores + + +def apply_importance_multiplier( + rows: list[dict[str, Any]], + *, + raw_field: str = 'importance_raw_score', + low: float = 0.1, + high: float = 2.0, +) -> None: + if not rows: + return + + values = sorted(float(row.get(raw_field, 0.0) or 0.0) for row in rows) + item_count = len(values) + median = values[item_count // 2] if item_count % 2 else (values[item_count // 2 - 1] + values[item_count // 2]) / 2 + q1 = values[item_count // 4] if item_count >= 4 else values[0] + q3 = values[3 * item_count // 4] if item_count >= 4 else values[-1] + iqr = q3 - q1 + + for row in rows: + raw_score = float(row.get(raw_field, 0.0) or 0.0) + if iqr <= 1e-9: + multiplier = 1.0 + else: + z_score = (raw_score - median) / iqr + sigmoid_score = 1.0 / (1.0 + math.exp(-z_score)) + multiplier = low + (high - low) * sigmoid_score + row['importance_multiplier'] = round(multiplier, 4) + row['agent_score'] = round( + float(row.get('agent_score', 0.0) or 0.0) * multiplier, + 6, + ) + row['discovery_score'] = round( + float(row.get('discovery_score', 0.0) or 0.0) * multiplier, + 6, + ) + + +def rank_candidates_by_path( + discovery_rows: list[dict[str, Any]], + routed_rows: list[dict[str, Any]], + top_k: int, + *, + importance_scores: dict[str, float] | None = None, +) -> list[dict[str, Any]]: + merged: dict[str, dict[str, Any]] = {} + insertion_order: dict[str, int] = {} + counter = 0 + + for row in discovery_rows: + key = get_candidate_key(row) + if not key: + continue + candidate = dict(row) + candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0) + candidate['agent_score'] = 0.0 + candidate.setdefault('hydrate_mode', 'chunks') + merged[key] = candidate + insertion_order[key] = counter + counter += 1 + + for row in routed_rows: + key = get_candidate_key(row) + if not key: + continue + routed_agent_score = float(row.get('agent_score', 0.0) or 0.0) + if key not in merged: + candidate = dict(row) + candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0) + candidate['agent_score'] = routed_agent_score + merged[key] = candidate + insertion_order[key] = counter + counter += 1 + continue + candidate = merged[key] + candidate['agent_score'] = max(float(candidate.get('agent_score', 0.0) or 0.0), routed_agent_score) + if not candidate.get('source_chunk_path') and row.get('source_chunk_path'): + candidate['source_chunk_path'] = row.get('source_chunk_path') + if not candidate.get('section_path') and row.get('section_path'): + candidate['section_path'] = row.get('section_path') + + for row in merged.values(): + row['importance_raw_score'] = float( + (importance_scores or {}).get(str(row.get('chunk_id') or ''), 0.0) or 0.0 + ) + apply_importance_multiplier(list(merged.values())) + + has_agent_results = len(routed_rows) > 0 + primary_rows: list[dict[str, Any]] = [] + fallback_rows: list[dict[str, Any]] = [] + + for key, row in merged.items(): + agent_score = float(row.get('agent_score', 0.0) or 0.0) + discovery_score = float(row.get('discovery_score', 0.0) or 0.0) + row['evidence_score'] = round(agent_score if has_agent_results else max(discovery_score, agent_score), 6) + row['score'] = row['evidence_score'] + row['_candidate_order'] = insertion_order[key] + + if has_agent_results and agent_score <= 0.0: + fallback_rows.append(row) + else: + primary_rows.append(row) + + def get_sort_key(row: dict[str, Any]) -> tuple[float, float, int]: + return ( + float(row.get('agent_score', 0.0) or 0.0), + float(row.get('discovery_score', 0.0) or 0.0), + -int(row.get('_candidate_order', 0) or 0), + ) + + primary_rows.sort(key=get_sort_key, reverse=True) + ranked_rows = primary_rows[:top_k] + + if len(ranked_rows) < top_k and fallback_rows: + fallback_rows.sort(key=get_sort_key, reverse=True) + ranked_rows.extend(fallback_rows[:top_k - len(ranked_rows)]) + + for row in ranked_rows: + row.pop('_candidate_order', None) + return ranked_rows + + +async def rank_retrieval_candidates( + db: AsyncSession, + *, + user_id: str, + namespace: str, + discovery_rows: list[dict[str, Any]], + routed_rows: list[dict[str, Any]], + top_k: int, +) -> list[dict[str, Any]]: + try: + importance_scores = await load_chunk_importance_scores( + db, + user_id=user_id, + namespace=namespace, + rows=[*discovery_rows, *routed_rows], + ) + except Exception as exc: + logger.warning(f'Failed to load chunk importance scores, continuing without importance: {exc}') + importance_scores = {} + return rank_candidates_by_path( + discovery_rows, + routed_rows, + top_k, + importance_scores=importance_scores, + ) diff --git a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py new file mode 100644 index 000000000..3f7e4b3d3 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +from typing import Any + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.job_result import JobResult +from shared.services.retrieval.search.section_filters import is_excluded_section + + +async def count_scoped_chunks( + db: AsyncSession, + *, + user_id: str, + namespace: str, + exclude_document_ids: list[str], + allowed_chunk_types: set[str] | None, +) -> int: + stmt = ( + select(func.count(DocumentChunk.id)) + .join( + Document, + (Document.document_id == DocumentChunk.document_id) + & (Document.current_job_result_id == DocumentChunk.job_result_id), + ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + ) + if exclude_document_ids: + stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) + if allowed_chunk_types is not None: + stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types))) + result = await db.execute(stmt) + return result.scalar() or 0 + + +async def load_all_scoped_chunks( + db: AsyncSession, + *, + user_id: str, + namespace: str, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + allowed_chunk_types: set[str] | None, + signal_paths: list[str], + filter_mode: str, +) -> list[dict[str, Any]]: + stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .join(JobResult, JobResult.id == DocumentChunk.job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .order_by(DocumentChunk.sort_order) + ) + if exclude_document_ids: + stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) + if allowed_chunk_types is not None: + stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types))) + + result = await db.execute(stmt) + rows: list[dict[str, Any]] = [] + for document, chunk, section, job_result in result.all(): + section_path = section.section_path if section else None + if is_excluded_section( + document_id=document.document_id, + section_path=section_path, + exclude_sections=exclude_sections, + ): + continue + if signal_paths and section_path: + path_lower = section_path.lower() + matches_any = any(keyword.lower() in path_lower for keyword in signal_paths) + if filter_mode == 'keep' and not matches_any: + continue + if filter_mode == 'delete' and matches_any: + continue + rows.append({ + 'document_id': document.document_id, + 'chunk_id': chunk.chunk_id, + 'section_id': chunk.section_id, + 'section_path': section_path, + 'source_file_name': document.source_file_name, + 'chunk_type': chunk.chunk_type, + 'content': chunk.content, + 'score': 1.0, + 'file_path': chunk.file_path, + 'chunk_metadata': chunk.chunk_metadata or {}, + 'job_result_id': chunk.job_result_id, + 'job_id': job_result.job_id if job_result else None, + 'sort_order': chunk.sort_order, + }) + return rows diff --git a/packages/shared-python/shared/services/retrieval/search/scoring.py b/packages/shared-python/shared/services/retrieval/search/scoring.py new file mode 100644 index 000000000..848a3adac --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/search/scoring.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +from typing import Any + +from shared.services.retrieval.settings import RRF_K + + +def get_row_path(row: dict[str, Any]) -> str: + """Extract the canonical path from a row for deduplication.""" + return str(row.get('section_path') or row.get('source_chunk_path') or '') + + +def merge_same_section_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: + if not rows: + return rows + groups: dict[str, list[dict[str, Any]]] = {} + order: list[str] = [] + for row in rows: + section_path = row.get('section_path') + if section_path: + key = f"{row.get('document_id', '')}::{section_path}" + else: + key = row.get('chunk_id', '') + if key not in groups: + groups[key] = [] + order.append(key) + groups[key].append(row) + + merged: list[dict[str, Any]] = [] + for key in order: + group = groups[key] + if len(group) == 1: + merged.append(group[0]) + continue + base = dict(group[0]) + base['content'] = '\n'.join(str(row.get('content', '')) for row in group) + base['score'] = max(row.get('score', 0.0) for row in group) + merged.append(base) + return merged + + +def merge_channels_rrf( + channels: list[list[dict[str, Any]]], + weights: list[float], + top_k: int, + k: int = RRF_K, +) -> list[dict[str, Any]]: + """Reciprocal Rank Fusion across multiple retrieval channels.""" + score_dict: dict[str, float] = {} + row_by_chunk_id: dict[str, dict[str, Any]] = {} + + for channel_idx, channel_rows in enumerate(channels): + weight = weights[channel_idx] if channel_idx < len(weights) else 1.0 + for rank, row in enumerate(channel_rows): + chunk_id = str(row.get('chunk_id') or '') + if not chunk_id: + continue + rrf_score = weight / (k + rank + 1) + score_dict[chunk_id] = score_dict.get(chunk_id, 0.0) + rrf_score + if chunk_id not in row_by_chunk_id: + row_by_chunk_id[chunk_id] = row + + ranked = sorted(score_dict.items(), key=lambda x: x[1], reverse=True) + results: list[dict[str, Any]] = [] + for chunk_id, fused_score in ranked[:top_k]: + row = row_by_chunk_id[chunk_id] + results.append(dict(row, score=round(fused_score, 6))) + return results + + +def normalize_row_scores( + rows: list[dict[str, Any]], + *, + source_field: str, + target_field: str, + default: float, +) -> None: + if not rows: + return + values = [float(row.get(source_field, 0.0) or 0.0) for row in rows] + min_score = min(values) + max_score = max(values) + if max_score <= 0.0 and min_score <= 0.0: + for row in rows: + row[target_field] = 0.0 + return + if max_score == min_score: + for row in rows: + row[target_field] = default + return + denominator = max_score - min_score + for row in rows: + raw_score = float(row.get(source_field, 0.0) or 0.0) + row[target_field] = round((raw_score - min_score) / denominator, 6) diff --git a/packages/shared-python/shared/services/retrieval/search/section_filters.py b/packages/shared-python/shared/services/retrieval/search/section_filters.py new file mode 100644 index 000000000..f74b7f3c3 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/search/section_filters.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +from collections.abc import Iterable + + +def is_excluded_section( + *, + document_id: str | None, + section_path: str | None, + exclude_sections: Iterable[dict[str, str]], +) -> bool: + document_id = str(document_id or '').strip() + section_path = str(section_path or '').strip() + if not document_id or not section_path: + return False + for item in exclude_sections: + if not isinstance(item, dict): + continue + exc_doc = str(item.get('document_id') or '').strip() + exc_path = str(item.get('section_path') or '').strip() + if document_id == exc_doc and ( + section_path == exc_path or section_path.startswith(exc_path + ' / ') + ): + return True + return False diff --git a/packages/shared-python/shared/services/retrieval/settings.py b/packages/shared-python/shared/services/retrieval/settings.py new file mode 100644 index 000000000..49f2461b5 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/settings.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +CHANNEL_WEIGHT_PATH = 1.0 +CHANNEL_WEIGHT_CONTENT = 2.0 +CHANNEL_WEIGHT_TERM = 1.5 +INTERNAL_RECALL_K_MULTIPLIER = 2 +RRF_K = 60 + +DATA_TYPE_ALLOWED_CHUNK_TYPES: dict[int, set[str] | None] = { + 1: None, + 2: {'text'}, + 3: {'image'}, + 4: {'table'}, + 5: {'text', 'image'}, + 6: {'text', 'table'}, +} + + +def resolve_allowed_chunk_types(data_type: int) -> set[str] | None: + return DATA_TYPE_ALLOWED_CHUNK_TYPES.get(data_type) diff --git a/packages/shared-python/shared/services/retrieval/stats/__init__.py b/packages/shared-python/shared/services/retrieval/stats/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/stats/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/stats/recorder.py b/packages/shared-python/shared/services/retrieval/stats/recorder.py new file mode 100644 index 000000000..967086a4b --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/stats/recorder.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +import asyncio +from typing import Any + +from loguru import logger + +from shared.services.retrieval.stats.service import record_retrieval_hits + + +_pending_retrieval_hit_stat_tasks: set[asyncio.Task[None]] = set() + + +def _finalize_retrieval_hit_stats_task(task: asyncio.Task[None]) -> None: + _pending_retrieval_hit_stat_tasks.discard(task) + + try: + task.result() + except asyncio.CancelledError: + logger.debug('Retrieval hit stats task was cancelled') + except Exception as exc: + logger.warning(f'Failed to record retrieval hit stats (ignored): {exc}') + + +def schedule_retrieval_hit_stats_update(*, user_id: str, namespace: str, results: list[dict[str, Any]]) -> None: + try: + task = asyncio.create_task( + _record_retrieval_hit_stats_best_effort( + user_id=user_id, + namespace=namespace, + results=results, + ), + name=f'retrieval_hit_stats:{user_id}:{namespace}', + ) + _pending_retrieval_hit_stat_tasks.add(task) + task.add_done_callback(_finalize_retrieval_hit_stats_task) + except Exception as exc: + logger.warning(f'Failed to schedule retrieval hit stats update (ignored): {exc}') + + +async def drain_retrieval_hit_stats_updates(timeout_seconds: float = 2.0) -> None: + if not _pending_retrieval_hit_stat_tasks: + return + + pending_tasks = tuple(_pending_retrieval_hit_stat_tasks) + + try: + await asyncio.wait_for( + asyncio.gather(*pending_tasks, return_exceptions=True), + timeout=timeout_seconds, + ) + except asyncio.TimeoutError: + for task in pending_tasks: + if not task.done(): + task.cancel() + + await asyncio.gather(*pending_tasks, return_exceptions=True) + + +async def _record_retrieval_hit_stats_best_effort(*, user_id: str, namespace: str, results: list[dict[str, Any]]) -> None: + try: + from shared.core.database import get_db_context + + async with get_db_context() as db: + await record_retrieval_hits(db, user_id=user_id, namespace=namespace, results=results) + await db.commit() + except Exception as exc: + logger.warning(f'Failed to record retrieval hit stats (ignored): {exc}') diff --git a/packages/shared-python/shared/services/retrieval/hit_stats_service.py b/packages/shared-python/shared/services/retrieval/stats/service.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/hit_stats_service.py rename to packages/shared-python/shared/services/retrieval/stats/service.py diff --git a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py index abce68b36..234927d04 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py @@ -2,37 +2,63 @@ from __future__ import annotations import asyncio -import os import time -from typing import Any +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager from uuid import uuid4 from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession -from shared.core.database import get_db_context -from shared.services.retrieval.agentic.budget import BudgetLedger -from shared.services.retrieval.agentic.orchestrator import RetrievalAgent, _load_budget_inventory -from shared.services.retrieval.agentic.types import AgenticResult -from shared.services.retrieval.cache_service import ( - get_cached_workflow_plan, - set_cached_workflow_plan, -) +from shared.services.retrieval.agentic.core.budget import BudgetLedger +from shared.services.retrieval.agentic.orchestrator import _load_budget_inventory from shared.services.retrieval.llm_adapter import ( create_retrieval_llm_fn, create_retrieval_planner_fn, ) -from shared.services.retrieval.workflow.planner import QueryPlanner -from shared.services.retrieval.workflow.synthesizer import compose_final_answer, synthesize_step -from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, StepResult, WorkflowResult +from shared.services.retrieval.workflow.plan_service import WorkflowPlanService +from shared.services.retrieval.workflow.reference_projection import WorkflowReferenceProjection +from shared.services.retrieval.workflow.run_request import WorkflowRunRequest +from shared.services.retrieval.workflow.runtime_config import WorkflowRuntimeConfig +from shared.services.retrieval.workflow.step_runner import WorkflowStepRunner +from shared.services.retrieval.workflow.synthesizer import compose_final_answer +from shared.services.retrieval.workflow.types import StepResult, WorkflowResult from shared.services.retrieval.workflow.wallet import BudgetWallet +DbSessionFactory = Callable[[], AbstractAsyncContextManager[AsyncSession]] +WorkflowStepRunnerFactory = Callable[[DbSessionFactory, str], WorkflowStepRunner] + + +def _create_workflow_step_runner( + db_factory: DbSessionFactory, + parent_run_id: str, +) -> WorkflowStepRunner: + return WorkflowStepRunner(db_factory=db_factory, parent_run_id=parent_run_id) + class WorkflowOrchestrator: """Plan and execute a query workflow DAG.""" - def __init__(self) -> None: + def __init__( + self, + db_factory: DbSessionFactory | None = None, + plan_service: WorkflowPlanService | None = None, + step_runner_factory: WorkflowStepRunnerFactory | None = None, + ) -> None: self.parent_run_id = f'wret_{uuid4().hex[:12]}' + self._db_factory = db_factory + self._plan_service = plan_service or WorkflowPlanService() + self._step_runner_factory = ( + step_runner_factory or _create_workflow_step_runner + ) + + def _get_db_factory(self) -> DbSessionFactory: + if self._db_factory is not None: + return self._db_factory + + from shared.core.database import get_db_context + + return get_db_context async def run( self, @@ -49,72 +75,90 @@ async def run( filter_mode: str = 'delete', channels: list[str] | None = None, channel_weights: dict[str, float] | None = None, + internal_recall_k: int | None = None, + rerank: bool = False, + threshold: float = 0.0, + llm_fn=None, + ) -> WorkflowResult: + request = WorkflowRunRequest( + user_id=user_id, + namespace=namespace, + query=query, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + data_type=data_type, + signal_paths=signal_paths, + filter_mode=filter_mode, + channels=channels, + channel_weights=channel_weights, + internal_recall_k=internal_recall_k, + rerank=rerank, + threshold=threshold, + ) + return await self.run_request(db, request=request, llm_fn=llm_fn) + + async def run_request( + self, + db: AsyncSession, + *, + request: WorkflowRunRequest, llm_fn=None, ) -> WorkflowResult: t0 = time.monotonic() + config = WorkflowRuntimeConfig.from_env() llm_fn = llm_fn or create_retrieval_llm_fn() planner_llm = create_retrieval_planner_fn(thinking=True) - planner_budget = _env_int('RETRIEVAL_PLANNER_THINKING_BUDGET', 4000) - wallet_total = _env_int('RETRIEVAL_WALLET_TOTAL_BUDGET', 200000) - per_retrieve = _env_int('RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET', 40000) - per_synthesize = _env_int('RETRIEVAL_WALLET_PER_SYNTHESIZE_STEP_BUDGET', 6000) - max_steps = _env_int('RETRIEVAL_DECOMPOSITION_MAX_STEPS', 5) planner_ledger = BudgetLedger( - total=planner_budget, + total=config.planner_budget, planning_ratio=0.0, - bootstrap=planner_budget, + bootstrap=config.planner_budget, per_doc_min_share=0, ) total_chunks, total_docs, _chunks_count_by_doc = await _load_budget_inventory( db, - user_id=user_id, - namespace=namespace, - exclude_document_ids=exclude_document_ids, + user_id=request.user_id, + namespace=request.namespace, + exclude_document_ids=request.exclude_document_ids, ) planner_ledger.total_chunks = total_chunks planner_ledger.total_docs = total_docs - plan = await self._load_or_plan( - user_id=user_id, - namespace=namespace, - query=query, + plan = await self._plan_service.load_or_create( + user_id=request.user_id, + namespace=request.namespace, + query=request.query, planner_llm=planner_llm, planner_ledger=planner_ledger, - max_steps=max_steps, - wallet_total=wallet_total, - per_retrieve=per_retrieve, - kb_total_docs=total_docs, - kb_total_chunks=total_chunks, + max_steps=config.max_steps, + wallet_total=config.wallet_total_budget, + per_retrieve=config.per_retrieve_step_budget, + corpus_total_docs=total_docs, + corpus_total_chunks=total_chunks, ) wallet = BudgetWallet( - total=wallet_total, - per_retrieve_step_default=per_retrieve, - per_synthesize_step_default=per_synthesize, + total=config.wallet_total_budget, + per_retrieve_step_default=config.per_retrieve_step_budget, + per_synthesize_step_default=config.per_synthesize_step_budget, ) ledgers = await wallet.allocate(plan) results_by_id: dict[str, StepResult] = {} - sem = asyncio.Semaphore(_env_int('RETRIEVAL_WORKFLOW_PARALLEL_MAX', 3)) + sem = asyncio.Semaphore(config.parallel_max) + step_runner = self._step_runner_factory( + self._get_db_factory(), + self.parent_run_id, + ) for batch in plan.topological_batches(): await asyncio.gather( *[ - self._run_step( - db, + step_runner.run_step( step=step, ledger=ledgers[step.id], results_by_id=results_by_id, semaphore=sem, - user_id=user_id, - namespace=namespace, - top_k=step.top_k or top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - data_type=step.data_type or data_type, - signal_paths=signal_paths, - filter_mode=filter_mode, - channels=channels, - channel_weights=channel_weights, + request=request.for_step(step), llm_fn=llm_fn, ) for step in batch @@ -125,10 +169,11 @@ async def run( answer_text = compose_final_answer(plan, results_by_id) ordered_results = [results_by_id[step.id] for step in plan.steps if step.id in results_by_id] - referenced_chunks = _dedupe_references( + reference_projection = WorkflowReferenceProjection() + referenced_chunks = reference_projection.dedupe( ref for step_result in ordered_results for ref in step_result.referenced_chunks ) - api_results = _references_to_results(referenced_chunks) + api_results = reference_projection.to_api_results(referenced_chunks) elapsed_ms = int((time.monotonic() - t0) * 1000) logger.info( 'workflow retrieval DONE: steps={} refs={} answer_chars={} elapsed={}ms', @@ -138,8 +183,8 @@ async def run( elapsed_ms, ) return WorkflowResult( - namespace=namespace, - query=query, + namespace=request.namespace, + query=request.query, router_used='workflow_decomposed' if len(plan.steps) > 1 else 'workflow_single_step', answer_text=answer_text, plan=plan, @@ -151,256 +196,3 @@ async def run( planner_snapshot=planner_ledger.snapshot(), parent_run_id=self.parent_run_id, ) - - async def _load_or_plan( - self, - *, - user_id: str, - namespace: str, - query: str, - planner_llm, - planner_ledger: BudgetLedger, - max_steps: int, - wallet_total: int, - per_retrieve: int, - kb_total_docs: int, - kb_total_chunks: int, - ) -> QueryPlan: - try: - cached = await get_cached_workflow_plan(user_id=user_id, namespace=namespace, query=query) - if cached: - return QueryPlan.from_dict(cached, original_query=query) - except Exception as exc: - logger.warning(f'workflow plan cache read failed (ignored): {exc}') - - planner = QueryPlanner( - llm_fn=planner_llm, - planner_ledger=planner_ledger, - max_steps=max_steps, - total_budget=wallet_total, - per_step_budget=per_retrieve, - ) - plan = await planner.plan( - query=query, - kb_total_docs=kb_total_docs, - kb_total_chunks=kb_total_chunks, - ) - try: - await set_cached_workflow_plan( - user_id=user_id, - namespace=namespace, - query=query, - plan=plan.to_dict(), - ) - except Exception as exc: - logger.warning(f'workflow plan cache write failed (ignored): {exc}') - return plan - - async def _run_step( - self, - db: AsyncSession, - *, - step: PlannedStep, - ledger: BudgetLedger, - results_by_id: dict[str, StepResult], - semaphore: asyncio.Semaphore, - user_id: str, - namespace: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - data_type: int, - signal_paths: list[str] | None, - filter_mode: str, - channels: list[str] | None, - channel_weights: dict[str, float] | None, - llm_fn, - ) -> None: - async with semaphore: - if step.step_kind == 'synthesize': - await self._run_synthesize_step(step, ledger, results_by_id, llm_fn) - return - await self._run_retrieve_step( - db, - step=step, - ledger=ledger, - results_by_id=results_by_id, - user_id=user_id, - namespace=namespace, - top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - data_type=data_type, - signal_paths=signal_paths, - filter_mode=filter_mode, - channels=channels, - channel_weights=channel_weights, - llm_fn=llm_fn, - ) - - async def _run_retrieve_step( - self, - db: AsyncSession, - *, - step: PlannedStep, - ledger: BudgetLedger, - results_by_id: dict[str, StepResult], - user_id: str, - namespace: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - data_type: int, - signal_paths: list[str] | None, - filter_mode: str, - channels: list[str] | None, - channel_weights: dict[str, float] | None, - llm_fn, - ) -> None: - try: - # AsyncSession is not safe for concurrent use. Workflow steps may - # run in the same topological batch, so each retrieve step opens an - # isolated session and leaves the parent session untouched. - async with get_db_context() as step_db: - agentic_result = await RetrievalAgent().run( - step_db, - user_id=user_id, - namespace=namespace, - query=step.sub_query, - top_k=top_k, - llm_fn=llm_fn, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - data_type=data_type, - signal_paths=signal_paths, - filter_mode=filter_mode, - channels=channels, - channel_weights=channel_weights, - ledger=ledger, - parent_run_id=self.parent_run_id, - workflow_step_id=step.id, - ) - results_by_id[step.id] = _step_result_from_agentic(step, agentic_result) - except Exception as exc: - logger.exception(f'workflow retrieve step failed: step_id={step.id}') - results_by_id[step.id] = StepResult( - step_id=step.id, - sub_query=step.sub_query, - step_kind=step.step_kind, - depends_on=step.depends_on, - output_role=step.output_role, - status='error', - error=str(exc), - budget_snapshot=ledger.snapshot(), - ) - - async def _run_synthesize_step( - self, - step: PlannedStep, - ledger: BudgetLedger, - results_by_id: dict[str, StepResult], - llm_fn, - ) -> None: - if llm_fn is None: - results_by_id[step.id] = StepResult( - step_id=step.id, - sub_query=step.sub_query, - step_kind=step.step_kind, - depends_on=step.depends_on, - output_role=step.output_role, - status='skipped', - answer_text='', - error='llm unavailable for synthesis', - budget_snapshot=ledger.snapshot(), - ) - return - prior = {dep: results_by_id[dep] for dep in step.depends_on if dep in results_by_id} - try: - answer = await synthesize_step(step, prior_results=prior, llm_fn=llm_fn, ledger=ledger) - refs = _dedupe_references( - ref for result in prior.values() for ref in result.referenced_chunks - ) - results_by_id[step.id] = StepResult( - step_id=step.id, - sub_query=step.sub_query, - step_kind=step.step_kind, - depends_on=step.depends_on, - output_role=step.output_role, - status='done', - answer_text=answer, - referenced_chunks=refs, - budget_snapshot=ledger.snapshot(), - ) - except Exception as exc: - results_by_id[step.id] = StepResult( - step_id=step.id, - sub_query=step.sub_query, - step_kind=step.step_kind, - depends_on=step.depends_on, - output_role=step.output_role, - status='budget_stop' if 'budget' in str(exc).lower() else 'error', - answer_text='(budget exhausted)' if 'budget' in str(exc).lower() else '', - error=str(exc), - budget_snapshot=ledger.snapshot(), - ) - - -def _step_result_from_agentic(step: PlannedStep, result: AgenticResult) -> StepResult: - if result.answer_text: - status = 'done' - elif result.failure_reason: - status = 'not_found' - elif 'budget' in (result.stop_reason or ''): - status = 'budget_stop' - else: - status = 'done' - return StepResult( - step_id=step.id, - sub_query=step.sub_query, - step_kind=step.step_kind, - depends_on=step.depends_on, - output_role=step.output_role, - status=status, # type: ignore[arg-type] - answer_text=result.answer_text, - evidence_text=result.evidence_text, - referenced_chunks=result.referenced_chunks, - budget_snapshot=result.budget_snapshot, - router_used=result.router_used, - stop_reason=result.stop_reason, - failure_reason=result.failure_reason, - ) - - -def _dedupe_references(refs) -> list[dict[str, Any]]: - seen: set[str] = set() - out: list[dict[str, Any]] = [] - for ref in refs: - chunk_id = str(ref.get('chunk_id') or '') - key = chunk_id or str(ref) - if key in seen: - continue - seen.add(key) - out.append(dict(ref)) - return out - - -def _references_to_results(refs: list[dict[str, Any]]) -> list[dict[str, Any]]: - return [ - { - 'chunk_id': ref.get('chunk_id'), - 'document_id': ref.get('document_id'), - 'chunk_type': ref.get('chunk_type'), - 'source': { - 'document_id': ref.get('document_id'), - 'section_path': ref.get('section_path'), - }, - } - for ref in refs - ] - - -def _env_int(name: str, default: int) -> int: - try: - return int(os.environ.get(name, str(default))) - except (TypeError, ValueError): - return default diff --git a/packages/shared-python/shared/services/retrieval/workflow/plan_service.py b/packages/shared-python/shared/services/retrieval/workflow/plan_service.py new file mode 100644 index 000000000..dec625979 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/workflow/plan_service.py @@ -0,0 +1,59 @@ +"""Plan loading and creation for decomposed retrieval workflows.""" +from __future__ import annotations + +from loguru import logger + +from shared.services.retrieval.agentic.core.budget import BudgetLedger +from shared.services.retrieval.cache_service import ( + get_cached_workflow_plan, + set_cached_workflow_plan, +) +from shared.services.retrieval.llm_adapter import LLMFn +from shared.services.retrieval.workflow.planner import QueryPlanner +from shared.services.retrieval.workflow.types import QueryPlan + + +class WorkflowPlanService: + async def load_or_create( + self, + *, + user_id: str, + namespace: str, + query: str, + planner_llm: LLMFn | None, + planner_ledger: BudgetLedger, + max_steps: int, + wallet_total: int, + per_retrieve: int, + corpus_total_docs: int, + corpus_total_chunks: int, + ) -> QueryPlan: + try: + cached = await get_cached_workflow_plan(user_id=user_id, namespace=namespace, query=query) + if cached: + return QueryPlan.from_dict(cached, original_query=query) + except Exception as exc: + logger.warning(f"workflow plan cache read failed (ignored): {exc}") + + planner = QueryPlanner( + llm_fn=planner_llm, + planner_ledger=planner_ledger, + max_steps=max_steps, + total_budget=wallet_total, + per_step_budget=per_retrieve, + ) + plan = await planner.plan( + query=query, + corpus_total_docs=corpus_total_docs, + corpus_total_chunks=corpus_total_chunks, + ) + try: + await set_cached_workflow_plan( + user_id=user_id, + namespace=namespace, + query=query, + plan=plan.to_dict(), + ) + except Exception as exc: + logger.warning(f"workflow plan cache write failed (ignored): {exc}") + return plan diff --git a/packages/shared-python/shared/services/retrieval/workflow/planner.py b/packages/shared-python/shared/services/retrieval/workflow/planner.py index e34713db6..82b574677 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/planner.py +++ b/packages/shared-python/shared/services/retrieval/workflow/planner.py @@ -8,7 +8,7 @@ from loguru import logger -from shared.services.retrieval.agentic.budget import BudgetExceeded, BudgetLedger +from shared.services.retrieval.agentic.core.budget import BudgetExceeded, BudgetLedger from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage from shared.services.retrieval.workflow.types import FinalStrategy, OutputRole, PlannedStep, QueryPlan, StepKind from shared.utils.token_estimate import estimate_tokens @@ -34,7 +34,7 @@ You are a retrieval workflow planner. Think step by step before answering. User query: {query} -Knowledge base inventory: {kb_total_docs} docs / {kb_total_chunks} chunks. +Document corpus inventory: {corpus_total_docs} docs / {corpus_total_chunks} chunks. Wallet status: total_budget={total_budget} tokens (planner_used={planner_used}). Decide whether the query needs decomposition into multiple sub-queries. @@ -80,16 +80,16 @@ async def plan( self, *, query: str, - kb_total_docs: int = 0, - kb_total_chunks: int = 0, + corpus_total_docs: int = 0, + corpus_total_chunks: int = 0, ) -> QueryPlan: if self._llm_fn is None: return QueryPlan.single_step(query, reason="planner_llm_unavailable") prompt = _PLANNER_PROMPT.format( query=query, - kb_total_docs=kb_total_docs, - kb_total_chunks=kb_total_chunks, + corpus_total_docs=corpus_total_docs, + corpus_total_chunks=corpus_total_chunks, total_budget=self._total_budget, planner_used=self._planner_used(), max_steps=self._max_steps, @@ -205,10 +205,10 @@ def _extract_json_object(text: str) -> dict[str, Any]: text = text.strip() try: parsed = json.loads(text) - if isinstance(parsed, dict): - return parsed except (ValueError, json.JSONDecodeError): - pass + parsed = None + if isinstance(parsed, dict): + return parsed match = re.search(r"\{.*\}", text, re.DOTALL) if not match: raise ValueError("no JSON object found") diff --git a/packages/shared-python/shared/services/retrieval/workflow/reference_projection.py b/packages/shared-python/shared/services/retrieval/workflow/reference_projection.py new file mode 100644 index 000000000..992085c00 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/workflow/reference_projection.py @@ -0,0 +1,40 @@ +"""Reference projection for decomposed retrieval workflows.""" +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + + +class WorkflowReferenceProjection: + def dedupe(self, refs: Iterable[dict[str, Any]]) -> list[dict[str, Any]]: + seen: set[str] = set() + out: list[dict[str, Any]] = [] + for ref in refs: + document_id = str(ref.get("document_id") or "").strip() + chunk_id = str(ref.get("chunk_id") or "").strip() + section_path = str(ref.get("section_path") or "").strip() + file_path = str(ref.get("file_path") or "").strip() + key = ( + f"{document_id}:{chunk_id}:{section_path}:{file_path}" + if document_id and chunk_id + else str(ref) + ) + if key in seen: + continue + seen.add(key) + out.append(dict(ref)) + return out + + def to_api_results(self, refs: list[dict[str, Any]]) -> list[dict[str, Any]]: + return [ + { + "chunk_id": ref.get("chunk_id"), + "document_id": ref.get("document_id"), + "chunk_type": ref.get("chunk_type"), + "source": { + "document_id": ref.get("document_id"), + "section_path": ref.get("section_path"), + }, + } + for ref in refs + ] diff --git a/packages/shared-python/shared/services/retrieval/workflow/run_request.py b/packages/shared-python/shared/services/retrieval/workflow/run_request.py new file mode 100644 index 000000000..05ee21cba --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/workflow/run_request.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from shared.services.retrieval.execution.route_types import RetrievalRouteContext +from shared.services.retrieval.settings import INTERNAL_RECALL_K_MULTIPLIER +from shared.services.retrieval.workflow.types import PlannedStep + + +@dataclass(frozen=True) +class WorkflowRunRequest: + user_id: str + namespace: str + query: str + top_k: int + exclude_document_ids: list[str] + exclude_sections: list[dict[str, str]] + data_type: int = 1 + signal_paths: list[str] | None = None + filter_mode: str = "delete" + channels: list[str] | None = None + channel_weights: dict[str, float] | None = None + internal_recall_k: int | None = None + rerank: bool = False + threshold: float = 0.0 + + @classmethod + def from_route_context( + cls, + context: RetrievalRouteContext, + ) -> WorkflowRunRequest: + return cls( + user_id=context.user_id, + namespace=context.namespace, + query=context.query, + top_k=context.top_k, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + data_type=context.data_type, + signal_paths=context.signal_paths, + filter_mode=context.filter_mode, + channels=context.channels, + channel_weights=context.channel_weights, + internal_recall_k=context.internal_recall_k, + rerank=context.rerank, + threshold=context.threshold, + ) + + def for_step(self, step: PlannedStep) -> WorkflowStepRequest: + step_top_k = step.top_k or self.top_k + return WorkflowStepRequest( + user_id=self.user_id, + namespace=self.namespace, + query=step.sub_query, + top_k=step_top_k, + exclude_document_ids=self.exclude_document_ids, + exclude_sections=self.exclude_sections, + data_type=step.data_type or self.data_type, + signal_paths=self.signal_paths, + filter_mode=self.filter_mode, + channels=self.channels, + channel_weights=self.channel_weights, + internal_recall_k=self.resolve_effective_recall_k(step_top_k), + rerank=self.rerank, + threshold=self.threshold, + ) + + def resolve_effective_recall_k(self, top_k: int) -> int: + if self.internal_recall_k is not None: + return self.internal_recall_k + return top_k * INTERNAL_RECALL_K_MULTIPLIER + + +@dataclass(frozen=True) +class WorkflowStepRequest: + user_id: str + namespace: str + query: str + top_k: int + exclude_document_ids: list[str] + exclude_sections: list[dict[str, str]] + data_type: int + signal_paths: list[str] | None + filter_mode: str + channels: list[str] | None + channel_weights: dict[str, float] | None + internal_recall_k: int | None + rerank: bool + threshold: float diff --git a/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py b/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py new file mode 100644 index 000000000..bcb6ab262 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py @@ -0,0 +1,33 @@ +"""Runtime configuration for decomposed retrieval workflows.""" +from __future__ import annotations + +import os +from dataclasses import dataclass + + +@dataclass(frozen=True) +class WorkflowRuntimeConfig: + planner_budget: int = 4000 + wallet_total_budget: int = 200000 + per_retrieve_step_budget: int = 40000 + per_synthesize_step_budget: int = 6000 + max_steps: int = 5 + parallel_max: int = 3 + + @classmethod + def from_env(cls) -> "WorkflowRuntimeConfig": + return cls( + planner_budget=_env_int("RETRIEVAL_PLANNER_THINKING_BUDGET", 4000), + wallet_total_budget=_env_int("RETRIEVAL_WALLET_TOTAL_BUDGET", 200000), + per_retrieve_step_budget=_env_int("RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET", 40000), + per_synthesize_step_budget=_env_int("RETRIEVAL_WALLET_PER_SYNTHESIZE_STEP_BUDGET", 6000), + max_steps=_env_int("RETRIEVAL_DECOMPOSITION_MAX_STEPS", 5), + parallel_max=_env_int("RETRIEVAL_WORKFLOW_PARALLEL_MAX", 3), + ) + + +def _env_int(name: str, default: int) -> int: + try: + return int(os.environ.get(name, str(default))) + except (TypeError, ValueError): + return default diff --git a/packages/shared-python/shared/services/retrieval/workflow/step_runner.py b/packages/shared-python/shared/services/retrieval/workflow/step_runner.py new file mode 100644 index 000000000..18a2e7059 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/workflow/step_runner.py @@ -0,0 +1,175 @@ +"""Step execution for decomposed retrieval workflows.""" +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager + +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.agentic.core.budget import BudgetLedger +from shared.services.retrieval.agentic.orchestrator import RetrievalAgent +from shared.services.retrieval.agentic.core.types import AgenticResult +from shared.services.retrieval.llm_adapter import LLMFn +from shared.services.retrieval.workflow.reference_projection import WorkflowReferenceProjection +from shared.services.retrieval.workflow.run_request import WorkflowStepRequest +from shared.services.retrieval.workflow.synthesizer import synthesize_step +from shared.services.retrieval.workflow.types import PlannedStep, StepResult, StepStatus + +DbSessionFactory = Callable[[], AbstractAsyncContextManager[AsyncSession]] +RetrievalAgentFactory = Callable[[], RetrievalAgent] + + +class WorkflowStepRunner: + def __init__( + self, + *, + db_factory: DbSessionFactory, + parent_run_id: str, + agent_factory: RetrievalAgentFactory | None = None, + ) -> None: + self._db_factory = db_factory + self._parent_run_id = parent_run_id + self._agent_factory = agent_factory or RetrievalAgent + self._references = WorkflowReferenceProjection() + + async def run_step( + self, + *, + step: PlannedStep, + ledger: BudgetLedger, + results_by_id: dict[str, StepResult], + semaphore: asyncio.Semaphore, + request: WorkflowStepRequest, + llm_fn: LLMFn | None, + ) -> None: + async with semaphore: + if step.step_kind == "synthesize": + await self._run_synthesize_step(step, ledger, results_by_id, llm_fn) + return + await self._run_retrieve_step( + step=step, + ledger=ledger, + results_by_id=results_by_id, + request=request, + llm_fn=llm_fn, + ) + + async def _run_retrieve_step( + self, + *, + step: PlannedStep, + ledger: BudgetLedger, + results_by_id: dict[str, StepResult], + request: WorkflowStepRequest, + llm_fn: LLMFn | None, + ) -> None: + try: + async with self._db_factory() as step_db: + agentic_result = await self._agent_factory().run( + step_db, + user_id=request.user_id, + namespace=request.namespace, + query=request.query, + top_k=request.top_k, + llm_fn=llm_fn, + exclude_document_ids=request.exclude_document_ids, + exclude_sections=request.exclude_sections, + data_type=request.data_type, + signal_paths=request.signal_paths, + filter_mode=request.filter_mode, + channels=request.channels, + channel_weights=request.channel_weights, + internal_recall_k=request.internal_recall_k, + ledger=ledger, + parent_run_id=self._parent_run_id, + workflow_step_id=step.id, + ) + results_by_id[step.id] = _step_result_from_agentic(step, agentic_result) + except Exception as exc: + logger.exception(f"workflow retrieve step failed: step_id={step.id}") + results_by_id[step.id] = StepResult( + step_id=step.id, + sub_query=step.sub_query, + step_kind=step.step_kind, + depends_on=step.depends_on, + output_role=step.output_role, + status="error", + error=str(exc), + budget_snapshot=ledger.snapshot(), + ) + + async def _run_synthesize_step( + self, + step: PlannedStep, + ledger: BudgetLedger, + results_by_id: dict[str, StepResult], + llm_fn: LLMFn | None, + ) -> None: + if llm_fn is None: + results_by_id[step.id] = StepResult( + step_id=step.id, + sub_query=step.sub_query, + step_kind=step.step_kind, + depends_on=step.depends_on, + output_role=step.output_role, + status="skipped", + answer_text="", + error="llm unavailable for synthesis", + budget_snapshot=ledger.snapshot(), + ) + return + prior = {dep: results_by_id[dep] for dep in step.depends_on if dep in results_by_id} + try: + answer = await synthesize_step(step, prior_results=prior, llm_fn=llm_fn, ledger=ledger) + refs = self._references.dedupe(ref for result in prior.values() for ref in result.referenced_chunks) + results_by_id[step.id] = StepResult( + step_id=step.id, + sub_query=step.sub_query, + step_kind=step.step_kind, + depends_on=step.depends_on, + output_role=step.output_role, + status="done", + answer_text=answer, + referenced_chunks=refs, + budget_snapshot=ledger.snapshot(), + ) + except Exception as exc: + results_by_id[step.id] = StepResult( + step_id=step.id, + sub_query=step.sub_query, + step_kind=step.step_kind, + depends_on=step.depends_on, + output_role=step.output_role, + status="budget_stop" if "budget" in str(exc).lower() else "error", + answer_text="(budget exhausted)" if "budget" in str(exc).lower() else "", + error=str(exc), + budget_snapshot=ledger.snapshot(), + ) + + +def _step_result_from_agentic(step: PlannedStep, result: AgenticResult) -> StepResult: + if result.answer_text: + status: StepStatus = "done" + elif result.failure_reason: + status = "not_found" + elif "budget" in (result.stop_reason or ""): + status = "budget_stop" + else: + status = "done" + return StepResult( + step_id=step.id, + sub_query=step.sub_query, + step_kind=step.step_kind, + depends_on=step.depends_on, + output_role=step.output_role, + status=status, + answer_text=result.answer_text, + evidence_text=result.evidence_text, + referenced_chunks=result.referenced_chunks, + budget_snapshot=result.budget_snapshot, + router_used=result.router_used, + stop_reason=result.stop_reason, + failure_reason=result.failure_reason, + ) diff --git a/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py b/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py index 89d4ec031..a81bb517d 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py +++ b/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py @@ -3,7 +3,7 @@ import re -from shared.services.retrieval.agentic.budget import BudgetLedger +from shared.services.retrieval.agentic.core.budget import BudgetLedger from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, StepResult from shared.utils.token_estimate import estimate_tokens diff --git a/packages/shared-python/shared/services/retrieval/workflow/wallet.py b/packages/shared-python/shared/services/retrieval/workflow/wallet.py index 0a540a120..624f03f7c 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/wallet.py +++ b/packages/shared-python/shared/services/retrieval/workflow/wallet.py @@ -5,7 +5,7 @@ import os from dataclasses import dataclass, field -from shared.services.retrieval.agentic.budget import BudgetLedger +from shared.services.retrieval.agentic.core.budget import BudgetLedger from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan diff --git a/packages/shared-python/shared/services/storage/file_encryptor_service.py b/packages/shared-python/shared/services/storage/file_encryptor_service.py deleted file mode 100755 index e542a705a..000000000 --- a/packages/shared-python/shared/services/storage/file_encryptor_service.py +++ /dev/null @@ -1,44 +0,0 @@ -import os -import pickle -from typing import Any - -from cryptography.fernet import Fernet - - -class FernetPickleEncryptor: - encrypt = False - - def __init__(self, key: bytes = b"nc1BPZSkNb7Oc82_Wo3QoZTmJCEnQtpKZ2n-Z5F4CwY="): - self.cipher = Fernet(key) - - def save_to_file(self, data: Any, file_path: str) -> None: - serialized_data = pickle.dumps(data) # Serialize the input payload. - encrypted_data = self.cipher.encrypt( - serialized_data - ) # Encrypt the serialized bytes. - with open(file_path, "wb") as f: - f.write(encrypted_data) - - def load_from_file(self, file_path: str) -> Any: - if not os.path.exists(file_path): - raise FileNotFoundError(f"File {file_path} does not exist.") - with open(file_path, "rb") as f: - encrypted_data = f.read() - decrypted_data = self.cipher.decrypt(encrypted_data) # Decrypt the file bytes. - loaded_data = pickle.loads(decrypted_data) # Deserialize the decrypted payload. - return loaded_data - - -encryptor = FernetPickleEncryptor() - -if __name__ == "__main__": - # 2. Encrypt the payload. - data = {"key": "value"} - # 3. Save the encrypted payload to a file. - file_path = "data.pkl" - encryptor.save_to_file(data, file_path) - print(f"Encrypted data saved to file: {file_path}") - - # 4. Load the encrypted payload back from disk. - decrypted_data = encryptor.load_from_file(file_path) - print("Decrypted data loaded from file:", decrypted_data) diff --git a/packages/shared-python/shared/services/storage/file_upload_service.py b/packages/shared-python/shared/services/storage/file_upload_service.py index c60f3c557..8b24793a4 100644 --- a/packages/shared-python/shared/services/storage/file_upload_service.py +++ b/packages/shared-python/shared/services/storage/file_upload_service.py @@ -1,147 +1,30 @@ -"""Storage upload service.""" +"""Async adapter for shared Job file storage.""" import asyncio -import json -import os -from typing import Any, Dict, Optional +from typing import Any, Optional from loguru import logger -from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import ( - KnowhereException, - StorageServiceException, -) -from shared.utils.pinned_outbound_http import ( - download_pinned_outbound_file_async, -) -from shared.utils.url_security import validate_http_url_and_resolve_ip_async +from shared.core.exceptions.domain_exceptions import StorageServiceException +from shared.services.storage.job_file_storage import JobFileStorage class FileUploadService: - """File upload service supporting S3, OSS, and MinIO.""" + """Async adapter over the shared Job file storage module.""" - def __init__(self): - self.adapter = settings.get_storage_adapter() - self.uploads_bucket = settings.S3_BUCKET_NAME - self.results_bucket = getattr( - settings, "S3_RESULTS_BUCKET", settings.S3_BUCKET_NAME - ) - - async def handle_direct_upload(self, file_path: str, job_id: str) -> str: - """ - Handle a direct file upload. - - Args: - file_path: Local file path. - job_id: Job ID. - - Returns: - str: Storage key. - """ - try: - # Build the storage key. - file_extension = os.path.splitext(file_path)[1] - s3_key = f"uploads/{job_id}{file_extension}" - - # Upload the file. - await self._upload_to_s3(file_path, s3_key, self.uploads_bucket) - - logger.info(f"Direct file upload succeeded: {file_path} -> {s3_key}") - return s3_key - - except KnowhereException: - raise - except Exception as e: - logger.error(f"Direct file upload failed: {e}") - raise StorageServiceException( - internal_message=f"Direct file upload failed: {str(e)}", - operation="direct_upload", - original_exception=e, - ) - - async def handle_url_upload(self, file_url: str, job_id: str) -> str: - """ - Handle a URL-based upload flow. - - Args: - file_url: File URL. - job_id: Job ID. - - Returns: - str: Storage key. - """ - try: - # Download the file into a temporary location first. - temp_file_path = await self._download_file_from_url(file_url) - - try: - # Build the storage key. - file_extension = os.path.splitext(file_url.split("?")[0])[1] - s3_key = f"uploads/{job_id}{file_extension}" - - # Upload the downloaded file. - await self._upload_to_s3(temp_file_path, s3_key, self.uploads_bucket) - - logger.info( - f"URL file download and upload succeeded: {file_url} -> {s3_key}" - ) - return s3_key - - finally: - # Clean up the temporary file. - if os.path.exists(temp_file_path): - os.remove(temp_file_path) - - except KnowhereException: - raise - except Exception as e: - logger.error(f"URL file handling failed: {e}") - raise StorageServiceException( - internal_message=f"URL file handling failed: {str(e)}", - operation="url_upload", - original_exception=e, - ) + def __init__(self, *, storage: JobFileStorage | None = None) -> None: + self._storage = storage or JobFileStorage() async def generate_upload_url( self, job_id: str, file_extension: str = "" - ) -> Dict[str, Any]: - """ - Generate a presigned upload URL. - - Args: - job_id: Job ID. - file_extension: File extension. - - Returns: - Dict: Upload URL payload including the storage key. - """ + ) -> dict[str, Any]: try: - s3_key = f"uploads/{job_id}{file_extension}" - - # Infer a Content-Type from the file extension. - content_type = self.get_content_type(file_extension) - - # Use the job waiting expiry as the upload URL TTL. - upload_url = self.adapter.generate_presigned_url( - s3_key, - expiration=settings.JOB_WAITING_EXPIRE_SECONDS, - bucket=self.uploads_bucket, - method="PUT", - headers={"Content-Type": content_type}, + return await asyncio.to_thread( + self._storage.generate_upload_url, + job_id=job_id, + file_extension=file_extension, ) - logger.info(f"Generated presigned upload URL: {upload_url}") - - return { - "upload_url": upload_url, - "s3_key": s3_key, - "expires_in": settings.JOB_WAITING_EXPIRE_SECONDS, - "upload_headers": {"Content-Type": content_type}, - } - - except KnowhereException: - raise except Exception as e: logger.error(f"Failed to generate upload URL: {e}") raise StorageServiceException( @@ -152,29 +35,15 @@ async def generate_upload_url( async def generate_download_url( self, s3_key: str, bucket: Optional[str] = None, expires_in: int = 3600 - ) -> Dict[str, Any]: - """ - Generate a presigned download URL. - - Args: - s3_key: Storage key. - bucket: Optional bucket name. - - Returns: - str: Download URL. - """ + ) -> dict[str, Any]: try: - bucket_name = bucket or self.results_bucket - - # Generate a one-hour presigned URL by default. - download_url = self.adapter.generate_presigned_url( - s3_key, expiration=expires_in, bucket=bucket_name, method="GET" + return await asyncio.to_thread( + self._storage.generate_download_url, + s3_key, + bucket=bucket or self._storage.results_bucket, + expires_in=expires_in, ) - return {"download_url": download_url, "expires_in": expires_in} - - except KnowhereException: - raise except Exception as e: logger.error(f"Failed to generate download URL: {e}") raise StorageServiceException( @@ -183,411 +52,19 @@ async def generate_download_url( original_exception=e, ) - async def get_file_info( + async def verify_s3_file_exists( self, s3_key: str, bucket: Optional[str] = None - ) -> Optional[Dict[str, Any]]: - """ - Get file information. - - Args: - s3_key: Storage key. - bucket: Optional bucket name. - - Returns: - Dict: File metadata. - """ - try: - bucket_name = bucket or self.results_bucket - - # Check existence and load the object size. - if not self.adapter.exists(s3_key, bucket_name): - return None - - size = self.adapter.get_object_size(s3_key, bucket_name) - return { - "size": size, - "content_type": None, # The adapter interface does not expose content_type yet. - "last_modified": None, - "etag": None, - } - - except Exception as e: - # Treat not-found responses as a missing object. - if "404" in str(e) or "not found" in str(e).lower(): - return None - logger.error(f"Failed to get file info: {e}") - raise StorageServiceException( - internal_message=f"Failed to get file info: {str(e)}", - operation="get_file_info", - original_exception=e, - ) - - async def upload_result_file( - self, local_file_path: str, job_id: str, file_extension: str = "" - ) -> str: - """ - Upload a result file. - - Args: - local_file_path: Local file path. - job_id: Job ID. - file_extension: File extension. - - Returns: - str: Storage key. - """ - try: - s3_key = f"results/{job_id}{file_extension}" - await self._upload_to_s3(local_file_path, s3_key, self.results_bucket) - - logger.info(f"Result file upload succeeded: {local_file_path} -> {s3_key}") - return s3_key - - except KnowhereException: - raise - except Exception as e: - logger.error(f"Result file upload failed: {e}") - raise StorageServiceException( - internal_message=f"Result file upload failed: {str(e)}", - operation="upload_result_file", - original_exception=e, - ) - - async def upload_json_result( - self, - job_id: str, - result_data: Dict[str, Any], - *, - content_type: str = "application/json", - ) -> str: - """Upload a JSON result file; deprecated but kept for compatibility.""" + ) -> dict[str, Any]: try: - s3_key = f"results/{job_id}.json" - from io import BytesIO - - body = json.dumps(result_data, ensure_ascii=False).encode("utf-8") - self.adapter.upload_fileobj( - BytesIO(body), + return await asyncio.to_thread( + self._storage.verify_exists, s3_key, - bucket=self.results_bucket, - content_type=content_type, + bucket=bucket or self._storage.uploads_bucket, ) - logger.info(f"Result JSON upload succeeded: job_id={job_id}, key={s3_key}") - return s3_key - except KnowhereException: - raise except Exception as e: - logger.error(f"Failed to upload result JSON: {e}") - raise StorageServiceException( - internal_message=f"Failed to upload result JSON: {str(e)}", - operation="upload_json_result", - original_exception=e, - ) - - async def upload_zip_result( - self, - job_id: str, - zip_file_path: str, - ) -> str: - """Upload a ZIP result file.""" - try: - s3_key = f"results/{job_id}.zip" - await self._upload_to_s3(zip_file_path, s3_key, self.results_bucket) - logger.info(f"Result ZIP upload succeeded: job_id={job_id}, key={s3_key}") - - # Clean up the temporary ZIP after upload. - try: - if os.path.exists(zip_file_path): - os.remove(zip_file_path) - except Exception as e: - logger.warning(f"Failed to clean up temporary ZIP file: {e}") - - return s3_key - except KnowhereException: - raise - except Exception as e: - logger.error(f"Failed to upload result ZIP: {e}") - raise StorageServiceException( - internal_message=f"Failed to upload result ZIP: {str(e)}", - operation="upload_zip_result", - original_exception=e, - ) - - def _ensure_bucket_exists(self, bucket_name: str) -> bool: - """ - Ensure the bucket is accessible. - - Args: - bucket_name: Bucket name. - - Returns: - bool: Whether the bucket check succeeded. - """ - try: - # In adapter mode, probe accessibility by listing objects. - adapter = settings.get_storage_adapter() - list(adapter.list_objects(prefix="", bucket=bucket_name)) - logger.debug(f"Bucket {bucket_name} is accessible") - return True - except Exception as e: - # The bucket is missing or inaccessible. - # For OSS, buckets should already exist; only accessibility is checked here. - logger.warning( - f"Bucket {bucket_name} may not exist or may be inaccessible: {e}" - ) - # In production, buckets should already be provisioned, so continue. - # Return False here instead if strict enforcement is ever needed. - return True - - async def _ensure_bucket_exists_async(self, bucket_name: str) -> bool: - """ - Asynchronously ensure the bucket is accessible. - - Args: - bucket_name: Bucket name. - - Returns: - bool: Whether the bucket check succeeded. - """ - - def _check_and_create(): - try: - # In adapter mode, probe accessibility by listing objects. - adapter = settings.get_storage_adapter() - list(adapter.list_objects(prefix="", bucket=bucket_name)) - logger.debug(f"Bucket {bucket_name} is accessible") - return True - except Exception as e: - # The bucket is missing or inaccessible. - logger.warning( - f"Bucket {bucket_name} may not exist or may be inaccessible: {e}" - ) - # In production, buckets should already be provisioned, so continue. - return True - - # Run the synchronous probe in a thread pool. - loop = asyncio.get_event_loop() - return await loop.run_in_executor(None, _check_and_create) - - async def _upload_to_s3(self, local_file_path: str, s3_key: str, bucket: str): - """Upload a file to storage.""" - # Ensure the bucket is accessible before uploading. - if not await self._ensure_bucket_exists_async(bucket): - raise StorageServiceException( - internal_message=f"Could not ensure bucket {bucket} exists", - operation="ensure_bucket", - ) - - def _upload(): - self.adapter.upload_file(local_file_path, s3_key, bucket) - - # Run the blocking upload in a thread pool. - loop = asyncio.get_event_loop() - await loop.run_in_executor(None, _upload) - - async def download_from_s3(self, s3_key: str, bucket: Optional[str] = None) -> str: - """Download a file from storage into a local temporary directory.""" - import uuid - - if bucket is None: - bucket = settings.S3_BUCKET_NAME - - # Create the temporary destination directory. - temp_dir = getattr(settings, "TMP_PATH", "/tmp") - os.makedirs(temp_dir, exist_ok=True) - - # Generate a temporary filename while preserving the original extension. - file_extension = os.path.splitext(s3_key)[1] - temp_filename = f"temp_{uuid.uuid4().hex}{file_extension}" - temp_file_path = os.path.join(temp_dir, temp_filename) - - try: - # Use the adapter to download the file. - def _download(): - self.adapter.download_file(s3_key, temp_file_path, bucket) - - # Run the blocking download in the event loop executor. - loop = asyncio.get_event_loop() - await loop.run_in_executor(None, _download) - - return temp_file_path - - except KnowhereException: - if os.path.exists(temp_file_path): - os.remove(temp_file_path) - raise - except Exception as e: - # Clean up the temporary file on failure. - if os.path.exists(temp_file_path): - os.remove(temp_file_path) - raise StorageServiceException( - internal_message=f"Failed to download file from S3: {str(e)}", - operation="download_from_s3", - original_exception=e, - ) - - async def _download_file_from_url(self, file_url: str) -> str: - """Download a file from a URL into a temporary directory.""" - temp_file_path = "" - try: - validation = await validate_http_url_and_resolve_ip_async(file_url) - if not validation.is_valid or not validation.validated_ip: - raise StorageServiceException( - internal_message=f"Invalid URL: {validation.error_message}", - operation="download_from_url", - ) - - temp_dir = getattr(settings, "TMP_PATH", "/tmp") - os.makedirs(temp_dir, exist_ok=True) - download_result = await download_pinned_outbound_file_async( - url=validation.url, - pinned_ip=validation.validated_ip, - timeout_seconds=300, - user_agent="Knowhere-FileDownloader/1.0", - temp_dir=temp_dir, - ) - temp_file_path = download_result.temp_file_path - return temp_file_path - - except KnowhereException: - raise - except Exception as e: - if os.path.exists(temp_file_path): - os.remove(temp_file_path) - raise StorageServiceException( - internal_message=f"Failed to download file: {str(e)}", - operation="download_from_url", - original_exception=e, - ) - - async def verify_s3_file_exists( - self, s3_key: str, bucket: Optional[str] = None - ) -> Dict[str, Any]: - """ - Verify whether a file exists in storage. - - Args: - s3_key: Storage key. - bucket: Optional bucket name. - - Returns: - Dict: File info payload, or `{"exists": False}` when missing. - """ - try: - bucket_name = bucket or self.uploads_bucket - - # Use the adapter to check object existence. - exists = self.adapter.exists(s3_key, bucket_name) - if not exists: - return {"exists": False} - - size = self.adapter.get_object_size(s3_key, bucket_name) - return { - "exists": True, - "size": size, - "content_type": None, - "last_modified": None, - "etag": None, - } - - except Exception as e: - # Treat not-found responses as a missing object. - if "404" in str(e) or "not found" in str(e).lower(): - return {"exists": False} logger.error(f"Failed to verify file existence: {e}") raise StorageServiceException( internal_message=f"Failed to verify file existence: {str(e)}", operation="verify_s3_file_exists", original_exception=e, ) - - def get_content_type(self, file_extension: str) -> str: - """ - Return a Content-Type for a file extension. - - Args: - file_extension: File extension, such as `.pdf` or `.docx`. - - Returns: - str: Content-Type - """ - content_types = { - ".pdf": "application/pdf", - ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", - ".doc": "application/msword", - ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", - ".xls": "application/vnd.ms-excel", - ".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation", - ".ppt": "application/vnd.ms-powerpoint", - ".csv": "text/csv", - ".txt": "text/plain", - ".md": "text/markdown", - ".json": "application/json", - ".xml": "application/xml", - ".html": "text/html", - ".htm": "text/html", - ".png": "image/png", - ".jpg": "image/jpeg", - ".jpeg": "image/jpeg", - ".gif": "image/gif", - ".bmp": "image/bmp", - ".tiff": "image/tiff", - ".svg": "image/svg+xml", - ".zip": "application/zip", - ".rar": "application/x-rar-compressed", - ".7z": "application/x-7z-compressed", - ".tar": "application/x-tar", - ".gz": "application/gzip", - } - return content_types.get(file_extension.lower(), "application/octet-stream") - - async def get_file_url( - self, s3_key: str, bucket: Optional[str] = None, expires_in: int = 3600 - ) -> str: - """ - Get a file URL from a storage key. - - Args: - s3_key: Storage key. - bucket: Optional bucket name. - expires_in: URL TTL in seconds, defaulting to one hour. - - Returns: - str: File URL. - """ - try: - bucket_name = bucket or self.uploads_bucket - - # Generate a presigned GET URL. - file_url = self.adapter.generate_presigned_url( - s3_key, expiration=expires_in, bucket=bucket_name, method="GET" - ) - - logger.info(f"Generated file URL successfully: {s3_key} -> {file_url}") - return file_url - - except KnowhereException: - raise - except Exception as e: - logger.error(f"Failed to get file URL: {e}") - raise StorageServiceException( - internal_message=f"Failed to get file URL: {str(e)}", - operation="get_file_url", - original_exception=e, - ) - - def generate_s3_key( - self, job_id: str, file_extension: str = "", prefix: str = "uploads" - ) -> str: - """ - Generate a storage key. - - Args: - job_id: Job ID. - file_extension: File extension. - prefix: Key prefix such as `uploads` or `results`. - - Returns: - str: Storage key. - """ - return f"{prefix}/{job_id}{file_extension}" diff --git a/packages/shared-python/shared/services/storage/job_file_storage.py b/packages/shared-python/shared/services/storage/job_file_storage.py new file mode 100644 index 000000000..77d49e2b4 --- /dev/null +++ b/packages/shared-python/shared/services/storage/job_file_storage.py @@ -0,0 +1,316 @@ +from __future__ import annotations + +import os +import tempfile +from typing import Any, BinaryIO + +from shared.core.config import settings +from shared.core.config.storage import get_cached_storage_adapter +from shared.core.exceptions.domain_exceptions import StorageServiceException +from shared.services.storage.storage_adapter import StorageAdapter +from shared.services.http.pinned_outbound import download_pinned_outbound_file +from shared.services.http.url_security import validate_http_url_and_resolve_ip + + +class JobFileStorage: + """Own storage rules for Job source files and Job Result bundles.""" + + def __init__( + self, + *, + storage_adapter: StorageAdapter | None = None, + uploads_bucket: str | None = None, + results_bucket: str | None = None, + ) -> None: + self._storage_adapter = storage_adapter + self.uploads_bucket = uploads_bucket or settings.S3_BUCKET_NAME + self.results_bucket = results_bucket or getattr( + settings, + "S3_RESULTS_BUCKET", + settings.S3_BUCKET_NAME, + ) + + @property + def storage_adapter(self) -> StorageAdapter: + if self._storage_adapter is None: + self._storage_adapter = get_cached_storage_adapter() + return self._storage_adapter + + def build_upload_key(self, *, job_id: str, file_extension: str = "") -> str: + return f"uploads/{job_id}{file_extension}" + + def build_result_key(self, *, job_id: str, file_extension: str = "") -> str: + return f"results/{job_id}{file_extension}" + + def build_result_zip_key(self, *, job_id: str) -> str: + return self.build_result_key(job_id=job_id, file_extension=".zip") + + def build_result_raw_prefix(self, *, job_id: str) -> str: + return f"results/{job_id}/" + + def generate_upload_url( + self, + *, + job_id: str, + file_extension: str = "", + ) -> dict[str, Any]: + storage_key = self.build_upload_key( + job_id=job_id, + file_extension=file_extension, + ) + content_type = self.get_content_type(file_extension) + upload_url = self.storage_adapter.generate_presigned_url( + storage_key, + expiration=settings.JOB_WAITING_EXPIRE_SECONDS, + bucket=self.uploads_bucket, + method="PUT", + headers={"Content-Type": content_type}, + ) + return { + "upload_url": upload_url, + "s3_key": storage_key, + "expires_in": settings.JOB_WAITING_EXPIRE_SECONDS, + "upload_headers": {"Content-Type": content_type}, + } + + def generate_download_url( + self, + storage_key: str, + *, + bucket: str, + expires_in: int = 3600, + ) -> dict[str, Any]: + download_url = self.storage_adapter.generate_presigned_url( + storage_key, + expiration=expires_in, + bucket=bucket, + method="GET", + ) + return {"download_url": download_url, "expires_in": expires_in} + + def generate_upload_download_url( + self, + storage_key: str, + *, + expires_in: int = 3600, + ) -> dict[str, Any]: + return self.generate_download_url( + storage_key, + bucket=self.uploads_bucket, + expires_in=expires_in, + ) + + def verify_exists( + self, + storage_key: str, + *, + bucket: str, + ) -> dict[str, Any]: + try: + if not self.storage_adapter.exists(storage_key, bucket): + return {"exists": False} + + size = self.storage_adapter.get_object_size(storage_key, bucket) + return { + "exists": True, + "size": size, + "content_type": None, + "last_modified": None, + "etag": None, + } + except Exception as exc: + if "404" in str(exc) or "not found" in str(exc).lower(): + return {"exists": False} + raise StorageServiceException( + internal_message=f"Storage file verification failed: {exc}", + operation="verify_exists", + original_exception=exc, + ) from exc + + def verify_upload_exists(self, storage_key: str) -> dict[str, Any]: + return self.verify_exists(storage_key, bucket=self.uploads_bucket) + + def upload_local_file( + self, + local_file_path: str, + storage_key: str, + *, + bucket: str, + ) -> dict[str, Any]: + try: + return self.storage_adapter.upload_file(local_file_path, storage_key, bucket) + except Exception as exc: + raise StorageServiceException( + internal_message=f"Storage upload failed: {exc}", + operation="upload_local_file", + original_exception=exc, + ) from exc + + def upload_source_file( + self, + local_file_path: str, + storage_key: str, + ) -> dict[str, Any]: + return self.upload_local_file( + local_file_path, + storage_key, + bucket=self.uploads_bucket, + ) + + def upload_fileobj( + self, + file_obj: BinaryIO, + storage_key: str, + *, + bucket: str, + content_type: str | None = None, + ) -> dict[str, Any]: + try: + return self.storage_adapter.upload_fileobj( + file_obj, + storage_key, + bucket=bucket, + content_type=content_type, + ) + except Exception as exc: + raise StorageServiceException( + internal_message=f"Storage upload file object failed: {exc}", + operation="upload_fileobj", + original_exception=exc, + ) from exc + + def download_to_path( + self, + storage_key: str, + local_path: str, + *, + bucket: str, + ) -> str: + try: + return self.storage_adapter.download_file(storage_key, local_path, bucket) + except Exception as exc: + raise StorageServiceException( + internal_message=f"Storage download failed: {exc}", + operation="download_to_path", + original_exception=exc, + ) from exc + + def download_to_temp( + self, + storage_key: str, + *, + suffix: str, + temp_dir: str, + bucket: str, + ) -> str: + local_temp_path: str | None = None + + try: + os.makedirs(temp_dir, exist_ok=True) + with tempfile.NamedTemporaryFile( + delete=False, + suffix=suffix, + dir=temp_dir, + ) as temp_file: + local_temp_path = temp_file.name + + self.download_to_path( + storage_key, + local_temp_path, + bucket=bucket, + ) + return local_temp_path + except Exception as exc: + if local_temp_path and os.path.exists(local_temp_path): + os.remove(local_temp_path) + raise StorageServiceException( + internal_message=( + "Failed to download object-storage file to temp path: " + f"storage_key={storage_key}, temp_dir={temp_dir}, error={exc}" + ), + operation="download_to_temp", + original_exception=exc, + ) from exc + + def download_upload_to_temp( + self, + storage_key: str, + *, + suffix: str, + temp_dir: str, + ) -> str: + return self.download_to_temp( + storage_key, + suffix=suffix, + temp_dir=temp_dir, + bucket=self.uploads_bucket, + ) + + def download_file_from_url( + self, + file_url: str, + *, + temp_dir: str | None = None, + ) -> str: + temp_file_path = "" + try: + validation = validate_http_url_and_resolve_ip(file_url) + if not validation.is_valid or not validation.validated_ip: + raise StorageServiceException( + internal_message=f"Invalid URL: {validation.error_message}", + operation="download_from_url", + ) + + effective_temp_dir = temp_dir or getattr(settings, "TMP_PATH", "/tmp") + os.makedirs(effective_temp_dir, exist_ok=True) + download_result = download_pinned_outbound_file( + url=validation.url, + pinned_ip=validation.validated_ip, + timeout_seconds=300, + user_agent="Knowhere-FileDownloader/1.0", + temp_dir=effective_temp_dir, + ) + temp_file_path = download_result.temp_file_path + return temp_file_path + except StorageServiceException: + raise + except Exception as exc: + if temp_file_path and os.path.exists(temp_file_path): + os.remove(temp_file_path) + raise StorageServiceException( + internal_message=f"Failed to download file: {exc}", + operation="download_from_url", + original_exception=exc, + ) from exc + + @staticmethod + def get_content_type(file_extension: str) -> str: + content_types = { + ".pdf": "application/pdf", + ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + ".doc": "application/msword", + ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + ".xls": "application/vnd.ms-excel", + ".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation", + ".ppt": "application/vnd.ms-powerpoint", + ".csv": "text/csv", + ".txt": "text/plain", + ".md": "text/markdown", + ".json": "application/json", + ".xml": "application/xml", + ".html": "text/html", + ".htm": "text/html", + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".gif": "image/gif", + ".bmp": "image/bmp", + ".tiff": "image/tiff", + ".svg": "image/svg+xml", + ".zip": "application/zip", + ".rar": "application/x-rar-compressed", + ".7z": "application/x-7z-compressed", + ".tar": "application/x-tar", + ".gz": "application/gzip", + } + return content_types.get(file_extension.lower(), "application/octet-stream") diff --git a/packages/shared-python/shared/services/storage/result_storage.py b/packages/shared-python/shared/services/storage/result_storage.py index 52d26783f..e16b8dab1 100644 --- a/packages/shared-python/shared/services/storage/result_storage.py +++ b/packages/shared-python/shared/services/storage/result_storage.py @@ -1,10 +1,16 @@ from __future__ import annotations import os +from collections.abc import Iterator from dataclasses import dataclass from pathlib import Path from typing import Protocol +from loguru import logger + +from shared.services.storage.job_file_storage import JobFileStorage +from shared.services.storage.storage_adapter import StorageAdapter + _EXCLUDED_FILE_NAMES = {".DS_Store", "Thumbs.db"} _EXCLUDED_DIR_NAMES = {"tmp", "temp", "__pycache__"} _CLIENT_ARTIFACT_DIRS = {"images", "tables"} @@ -20,41 +26,36 @@ class UploadedResultBundle: class ResultStorage(Protocol): def upload( self, *, job_id: str, result_dir: str, zip_file_path: str - ) -> UploadedResultBundle: ... + ) -> UploadedResultBundle: + raise NotImplementedError def generate_artifact_url( self, *, job_id: str, artifact_ref: str, expires_in: int = 3600 - ) -> str | None: ... + ) -> str | None: + raise NotImplementedError - def normalize_artifact_ref(self, artifact_ref: str | None) -> str | None: ... + def normalize_artifact_ref(self, artifact_ref: str | None) -> str | None: + raise NotImplementedError -class ResultS3: +class JobResultStorage: def __init__( - self, *, results_bucket: str | None = None, storage_adapter=None + self, + *, + results_bucket: str | None = None, + storage_adapter: StorageAdapter | None = None, ) -> None: - if results_bucket is None: - from shared.core.config import settings - - results_bucket = getattr( - settings, "S3_RESULTS_BUCKET", settings.S3_BUCKET_NAME - ) - self.results_bucket = results_bucket - self._storage_adapter = storage_adapter - - @property - def storage_adapter(self): - if self._storage_adapter is None: - from shared.core.config.storage import get_cached_storage_adapter - - self._storage_adapter = get_cached_storage_adapter() - return self._storage_adapter + self._job_file_storage = JobFileStorage( + storage_adapter=storage_adapter, + results_bucket=results_bucket, + ) + self.results_bucket = self._job_file_storage.results_bucket def build_zip_key(self, *, job_id: str) -> str: - return f"results/{job_id}.zip" + return self._job_file_storage.build_result_zip_key(job_id=job_id) def build_raw_prefix(self, *, job_id: str) -> str: - return f"results/{job_id}/" + return self._job_file_storage.build_result_raw_prefix(job_id=job_id) def build_raw_key(self, *, job_id: str, relative_path: str) -> str: normalized = self._normalize_raw_relative_path(relative_path) @@ -82,15 +83,21 @@ def upload( if not zip_path.is_file(): raise ValueError(f"Result ZIP file does not exist: {zip_file_path}") zip_key = self.build_zip_key(job_id=job_id) - self.storage_adapter.upload_file(str(zip_path), zip_key, self.results_bucket) + self._job_file_storage.upload_local_file( + str(zip_path), + zip_key, + bucket=self.results_bucket, + ) self._cleanup_file(zip_path) raw_files: dict[str, str] = {} for file_path in self._iter_raw_files(result_path): relative_path = file_path.relative_to(result_path).as_posix() raw_key = self.build_raw_key(job_id=job_id, relative_path=relative_path) - self.storage_adapter.upload_file( - str(file_path), raw_key, self.results_bucket + self._job_file_storage.upload_local_file( + str(file_path), + raw_key, + bucket=self.results_bucket, ) raw_files[relative_path] = raw_key @@ -101,12 +108,11 @@ def upload( ) def generate_url(self, *, storage_key: str, expires_in: int = 3600) -> str | None: - return self.storage_adapter.generate_presigned_url( + return self._job_file_storage.generate_download_url( storage_key, - expiration=expires_in, bucket=self.results_bucket, - method="GET", - ) + expires_in=expires_in, + )["download_url"] def generate_artifact_url( self, *, job_id: str, artifact_ref: str, expires_in: int = 3600 @@ -119,7 +125,7 @@ def generate_artifact_url( expires_in=expires_in, ) - def _iter_raw_files(self, result_dir: Path): + def _iter_raw_files(self, result_dir: Path) -> Iterator[Path]: for root, dir_names, file_names in os.walk(result_dir): dir_names[:] = [ dir_name @@ -155,9 +161,9 @@ def _is_excluded_dir(self, dir_name: str) -> bool: def _cleanup_file(self, file_path: Path) -> None: try: file_path.unlink(missing_ok=True) - except Exception: - pass + except Exception as exc: + logger.debug(f"Failed to clean up result file {file_path}: {exc}") def get_result_storage() -> ResultStorage: - return ResultS3() + return JobResultStorage() diff --git a/packages/shared-python/shared/services/storage/zip_chunk_schema.py b/packages/shared-python/shared/services/storage/zip_chunk_schema.py new file mode 100644 index 000000000..2dbf973d1 --- /dev/null +++ b/packages/shared-python/shared/services/storage/zip_chunk_schema.py @@ -0,0 +1,182 @@ +"""Chunk projection for Knowhere ZIP result packages.""" + +from __future__ import annotations + +from typing import Any + +from shared.services.chunks.chunk_connections import ( + build_resource_target_map, + convert_refs_to_embed_connections, + merge_connections, + normalize_connect_to_targets, + parse_relationship_refs, +) + + +class ZipChunkSchemaBuilder: + def calculate_statistics(self, chunks: list[dict[str, Any]]) -> dict[str, Any]: + total_chunks = len(chunks) + text_chunks = 0 + image_chunks = 0 + table_chunks = 0 + + for chunk in chunks: + chunk_type = chunk.get("type", "") + normalized_type = _normalize_chunk_type(chunk_type) + if normalized_type == "image": + image_chunks += 1 + elif normalized_type == "table": + table_chunks += 1 + else: + text_chunks += 1 + + return { + "total_chunks": total_chunks, + "text_chunks": text_chunks, + "image_chunks": image_chunks, + "table_chunks": table_chunks, + "total_pages": None, + } + + def format_chunks( + self, + chunks: list[dict[str, Any]], + image_files_map: dict[str, dict[str, Any]], + table_files_map: dict[str, dict[str, Any]], + ) -> list[dict[str, Any]]: + resource_target_map = build_resource_target_map( + chunks, + image_files_map=image_files_map, + table_files_map=table_files_map, + ) + + formatted: list[dict[str, Any]] = [] + for chunk in chunks: + chunk_id = str(chunk.get("chunk_id") or chunk.get("know_id")) + chunk_type_str = chunk.get("type", "") + normalized_type = _normalize_chunk_type(chunk_type_str) + image_info = image_files_map.get(chunk_id) + + if normalized_type == "image": + chunk_type = "image" + elif normalized_type == "table": + chunk_type = "table" + else: + chunk_type = "text" + + content = chunk.get("text") or chunk.get("content", "") + path = chunk.get("path", "") + existing_metadata = chunk.get("metadata", {}) + metadata = _base_chunk_metadata(existing_metadata, chunk, content) + + if chunk_type == "text": + metadata.update( + _format_text_metadata( + chunk=chunk, + chunk_type_str=chunk_type_str, + content=str(content), + existing_metadata=existing_metadata, + resource_target_map=resource_target_map, + ) + ) + elif chunk_type == "image": + if image_info: + metadata["file_path"] = image_info["file_path"] + metadata["keywords"] = existing_metadata.get("keywords") or chunk.get( + "keywords", [] + ) + metadata["tokens"] = [] + elif chunk_type == "table": + metadata["file_path"] = _resolve_table_file_path( + chunk=chunk, + chunk_id=chunk_id, + path=str(path), + existing_metadata=existing_metadata, + table_files_map=table_files_map, + ) + metadata["keywords"] = existing_metadata.get("keywords") or chunk.get( + "keywords", [] + ) + metadata["tokens"] = [] + + formatted.append( + { + "chunk_id": chunk_id, + "type": chunk_type, + "content": content, + "path": path, + "metadata": metadata, + } + ) + + return formatted + + +def _normalize_chunk_type(value: Any) -> str: + raw_type = str(value).strip() + return raw_type.split("\n", 1)[0].lower() + + +def _base_chunk_metadata( + existing_metadata: dict[str, Any], + chunk: dict[str, Any], + content: Any, +) -> dict[str, Any]: + metadata = { + "length": existing_metadata.get("length") or len(content), + "summary": existing_metadata.get("summary") or chunk.get("summary", ""), + "page_nums": existing_metadata.get("page_nums", []), + } + document_top_summary = str(existing_metadata.get("document_top_summary") or "").strip() + if document_top_summary: + metadata["document_top_summary"] = document_top_summary + return metadata + + +def _format_text_metadata( + *, + chunk: dict[str, Any], + chunk_type_str: Any, + content: str, + existing_metadata: dict[str, Any], + resource_target_map: dict[str, str], +) -> dict[str, Any]: + relationship_refs = parse_relationship_refs( + chunk.get("type_raw") or chunk_type_str, + content, + ) + embed_connections = convert_refs_to_embed_connections( + relationship_refs, + resource_target_map, + ) + related_connections = normalize_connect_to_targets( + existing_metadata.get("connect_to") + or chunk.get("connect_to") + or chunk.get("connectto"), + resource_target_map, + ) + return { + "tokens": existing_metadata.get("tokens") or chunk.get("tokens", 0), + "keywords": existing_metadata.get("keywords") or chunk.get("keywords", []), + "connect_to": merge_connections(embed_connections, related_connections), + } + + +def _resolve_table_file_path( + *, + chunk: dict[str, Any], + chunk_id: str, + path: str, + existing_metadata: dict[str, Any], + table_files_map: dict[str, dict[str, Any]], +) -> Any: + file_path = existing_metadata.get("file_path") + if file_path: + return file_path + + table_info = table_files_map.get(chunk_id) + if table_info: + return table_info["file_path"] + + table_name = path.split("/")[-1] if "/" in path else f"table_{chunk_id}.html" + return f"tables/{table_name}" diff --git a/packages/shared-python/shared/services/storage/zip_doc_navigation.py b/packages/shared-python/shared/services/storage/zip_doc_navigation.py new file mode 100644 index 000000000..e4e2f8f13 --- /dev/null +++ b/packages/shared-python/shared/services/storage/zip_doc_navigation.py @@ -0,0 +1,183 @@ +"""Document navigation projection for Knowhere ZIP result packages.""" + +from __future__ import annotations + +from typing import Any + +from shared.services.chunks.document_path import split_document_path +from shared.utils.text_utils import truncate_content_preview + + +class ZipDocNavigationBuilder: + def build_hierarchy_dict( + self, + sections: list[dict[str, Any]], + ) -> dict[str, Any]: + hierarchy: dict[str, Any] = {} + title_counts: dict[str, int] = {} + + for section in sections: + raw_title = str(section.get("title") or "").strip() + if not raw_title: + continue + + title_counts[raw_title] = title_counts.get(raw_title, 0) + 1 + title = ( + raw_title + if title_counts[raw_title] == 1 + else f"{raw_title} ({title_counts[raw_title]})" + ) + hierarchy[title] = self.build_hierarchy_dict( + section.get("children") or [] + ) + + return hierarchy + + def build_doc_nav( + self, + formatted_chunks: list[dict[str, Any]], + source_file_name: str, + ) -> dict[str, Any]: + text_chunks: list[dict[str, Any]] = [] + image_resources: list[dict[str, Any]] = [] + table_resources: list[dict[str, Any]] = [] + + stats = { + "total_chunks": 0, + "text_chunks": 0, + "image_chunks": 0, + "table_chunks": 0, + "max_depth": 0, + } + + for formatted_chunk in formatted_chunks: + chunk_type = formatted_chunk.get("type", "text") + path = formatted_chunk.get("path", "") + metadata = formatted_chunk.get("metadata") or {} + summary_raw = (metadata.get("summary") or "").strip() + content_raw = (formatted_chunk.get("content") or "").strip() + summary = " ".join(summary_raw.split()) if summary_raw else "" + content_preview = truncate_content_preview(content_raw) if content_raw else "" + + stats["total_chunks"] += 1 + if chunk_type == "image": + stats["image_chunks"] += 1 + image_resources.append( + { + "path": path, + "summary": summary or content_preview, + } + ) + elif chunk_type == "table": + stats["table_chunks"] += 1 + table_resources.append( + { + "path": path, + "summary": summary or content_preview, + } + ) + else: + stats["text_chunks"] += 1 + text_chunks.append( + { + "path": path, + "summary": summary or content_preview, + } + ) + + sections = self._build_section_tree( + text_chunks, + source_file_name=source_file_name, + ) + stats["max_depth"] = _max_depth(sections) + return { + "version": "1.0", + "file_name": source_file_name or "", + "stats": stats, + "sections": sections, + "resources": { + "images": image_resources, + "tables": table_resources, + }, + } + + def _build_section_tree( + self, + text_chunks: list[dict[str, Any]], + *, + source_file_name: str, + ) -> list[dict[str, Any]]: + root_children: dict[str, dict[str, Any]] = {} + + for chunk in text_chunks: + path = chunk.get("path", "") + root_parts, section_parts = split_document_path( + path, + source_file_name=source_file_name, + ) + + if not section_parts: + key = "__root__" + if key not in root_children: + root_children[key] = { + "title": "Root", + "path": "/".join(root_parts) if root_parts else path, + "summary": chunk.get("summary", ""), + "chunk_count": 0, + "_children_map": {}, + } + root_children[key]["chunk_count"] += 1 + if not root_children[key]["summary"]: + root_children[key]["summary"] = chunk.get("summary", "") + continue + + current_level = root_children + full_section_path_parts = list(root_parts) + for index, part in enumerate(section_parts): + full_section_path_parts.append(part) + if part not in current_level: + current_level[part] = { + "title": part, + "path": "/".join(full_section_path_parts), + "summary": "", + "chunk_count": 0, + "_children_map": {}, + } + node = current_level[part] + if index == len(section_parts) - 1: + node["chunk_count"] += 1 + if not node["summary"]: + node["summary"] = chunk.get("summary", "") + current_level = node["_children_map"] + + return _section_tree_to_output(root_children) + + +def _max_depth(nodes: list[dict[str, Any]], depth: int = 1) -> int: + max_depth = depth if nodes else 0 + for node in nodes: + max_depth = max(max_depth, _max_depth(node.get("children", []), depth + 1)) + return max_depth + + +def _section_tree_to_output( + children_map: dict[str, dict[str, Any]], + level: int = 1, +) -> list[dict[str, Any]]: + result: list[dict[str, Any]] = [] + for node in children_map.values(): + children = _section_tree_to_output(node["_children_map"], level + 1) + total_chunks = node["chunk_count"] + sum( + child.get("chunk_count", 0) for child in children + ) + result.append( + { + "title": node["title"], + "path": node["path"], + "level": level, + "summary": node["summary"], + "chunk_count": total_chunks, + "children": children, + } + ) + return result diff --git a/packages/shared-python/shared/services/storage/zip_manifest_schema.py b/packages/shared-python/shared/services/storage/zip_manifest_schema.py new file mode 100644 index 000000000..19afd84b9 --- /dev/null +++ b/packages/shared-python/shared/services/storage/zip_manifest_schema.py @@ -0,0 +1,42 @@ +"""Manifest projection for Knowhere ZIP result packages.""" + +from __future__ import annotations + +from typing import Any + +from shared.utils.utc_now import utc_now_naive + + +class ZipManifestBuilder: + def generate_manifest( + self, + *, + job_id: str, + data_id: str | None, + source_file_name: str, + statistics: dict[str, Any], + job_metadata: dict[str, Any], + hierarchy: dict[str, Any] | None = None, + ) -> dict[str, Any]: + return { + "version": "2.0", + "job_id": job_id, + "data_id": data_id, + "source_file_name": source_file_name, + "processing_date": utc_now_naive().isoformat() + "Z", + "processing": { + "page_count": job_metadata.get("page_count"), + "billing_status": job_metadata.get("billing_status"), + "cost": { + "micro_dollars": job_metadata.get("billing_amount_micro_dollars"), + "credits": job_metadata.get("billing_credits"), + }, + "timing": { + "started_at": job_metadata.get("processing_started_at"), + "completed_at": job_metadata.get("processing_completed_at"), + "duration_ms": job_metadata.get("processing_duration_ms"), + }, + }, + "statistics": statistics, + "HIERARCHY": hierarchy or {}, + } diff --git a/packages/shared-python/shared/services/storage/zip_package_writer.py b/packages/shared-python/shared/services/storage/zip_package_writer.py new file mode 100644 index 000000000..42c78532e --- /dev/null +++ b/packages/shared-python/shared/services/storage/zip_package_writer.py @@ -0,0 +1,112 @@ +"""Physical ZIP writing for Knowhere result packages.""" + +from __future__ import annotations + +import hashlib +import json +import os +import tempfile +import zipfile +from dataclasses import dataclass +from typing import Any + +from loguru import logger + +from shared.services.storage.zip_result_resources import ZipResourceFileInfo + + +@dataclass(frozen=True) +class ZipPackageWriteRequest: + job_id: str + add_dir: str + formatted_chunks: list[dict[str, Any]] + image_files: tuple[ZipResourceFileInfo, ...] + table_files: tuple[ZipResourceFileInfo, ...] + doc_nav: dict[str, Any] | None + manifest: dict[str, Any] + temp_dir: str | None + + +@dataclass(frozen=True) +class ZipPackageArtifact: + zip_file_path: str + checksum: dict[str, str] + zip_size: int + + +class ZipPackageWriter: + """Write a prepared result package to a ZIP file.""" + + def write(self, request: ZipPackageWriteRequest) -> ZipPackageArtifact: + effective_temp_dir = request.temp_dir or tempfile.gettempdir() + os.makedirs(effective_temp_dir, exist_ok=True) + zip_file_path = os.path.join(effective_temp_dir, f"result_{request.job_id}.zip") + + with zipfile.ZipFile(zip_file_path, "w", zipfile.ZIP_DEFLATED) as zip_file: + chunks_json = json.dumps( + {"chunks": request.formatted_chunks}, + ensure_ascii=False, + indent=2, + ) + zip_file.writestr("chunks.json", chunks_json.encode("utf-8")) + + self._write_optional_file(zip_file, request.add_dir, "full.md") + if self._write_optional_file( + zip_file, + request.add_dir, + "toc_hierarchies.json", + ): + logger.info("Added toc_hierarchies.json to ZIP") + + self._write_resource_files(zip_file, request.image_files, label="Image") + self._write_resource_files(zip_file, request.table_files, label="Table") + + if request.doc_nav is not None: + doc_nav_json = json.dumps(request.doc_nav, ensure_ascii=False, indent=2) + zip_file.writestr("doc_nav.json", doc_nav_json.encode("utf-8")) + logger.info("Added doc_nav.json") + + manifest_json = json.dumps(request.manifest, ensure_ascii=False, indent=2) + zip_file.writestr("manifest.json", manifest_json.encode("utf-8")) + + checksum_value = _calculate_zip_checksum(zip_file_path) + zip_size = os.path.getsize(zip_file_path) + return ZipPackageArtifact( + zip_file_path=zip_file_path, + checksum={"algorithm": "sha256", "value": checksum_value}, + zip_size=zip_size, + ) + + def _write_optional_file( + self, + zip_file: zipfile.ZipFile, + add_dir: str, + filename: str, + ) -> bool: + file_path = os.path.join(add_dir, filename) + if not os.path.exists(file_path): + return False + zip_file.write(file_path, filename) + return True + + def _write_resource_files( + self, + zip_file: zipfile.ZipFile, + resources: tuple[ZipResourceFileInfo, ...], + *, + label: str, + ) -> None: + for resource in resources: + source_path = resource["source_path"] + if os.path.exists(source_path): + zip_file.write(source_path, resource["zip_path"]) + else: + logger.warning(f"{label} file not found: {source_path}") + + +def _calculate_zip_checksum(zip_file_path: str) -> str: + sha256_hash = hashlib.sha256() + with open(zip_file_path, "rb") as file_obj: + for byte_block in iter(lambda: file_obj.read(4096), b""): + sha256_hash.update(byte_block) + return sha256_hash.hexdigest().lower() diff --git a/packages/shared-python/shared/services/storage/zip_result_resources.py b/packages/shared-python/shared/services/storage/zip_result_resources.py new file mode 100644 index 000000000..e84e0cb5a --- /dev/null +++ b/packages/shared-python/shared/services/storage/zip_result_resources.py @@ -0,0 +1,324 @@ +"""Resource discovery for Knowhere ZIP result packages.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Any + +from loguru import logger +from PIL import Image + +from shared.core.exceptions.domain_exceptions import StorageServiceException + +ZipResourceFileInfo = dict[str, Any] + + +@dataclass(frozen=True) +class ZipPackageResources: + image_files: tuple[ZipResourceFileInfo, ...] + table_files: tuple[ZipResourceFileInfo, ...] + + @property + def image_files_map(self) -> dict[str, ZipResourceFileInfo]: + return {str(image["id"]): image for image in self.image_files} + + @property + def table_files_map(self) -> dict[str, ZipResourceFileInfo]: + return {str(table["id"]): table for table in self.table_files} + + +class ZipResourceCollector: + """Resolve chunk resource references to files that belong in the result ZIP.""" + + def collect( + self, + *, + chunks: list[dict[str, Any]], + add_dir: str, + ) -> ZipPackageResources: + images_dir = os.path.join(add_dir, "images") + tables_dir = os.path.join(add_dir, "tables") + return ZipPackageResources( + image_files=tuple(self._collect_image_files(chunks, images_dir)), + table_files=tuple(self._collect_table_files(chunks, tables_dir)), + ) + + def _collect_image_files( + self, + chunks: list[dict[str, Any]], + images_dir: str, + ) -> list[ZipResourceFileInfo]: + image_files: list[ZipResourceFileInfo] = [] + if not os.path.exists(images_dir): + has_image_chunks = any(chunk.get("type", "") == "image" for chunk in chunks) + if has_image_chunks: + raise StorageServiceException( + internal_message=( + "Image directory not found for ZIP packaging: " + f"images_dir={images_dir}" + ), + operation="collect_image_files", + ) + return image_files + + image_files_map = _collect_files_by_name(images_dir) + + for chunk in chunks: + chunk_id = chunk.get("chunk_id") or chunk.get("know_id") + chunk_type = chunk.get("type", "") + + if chunk_type != "image": + continue + + metadata = chunk.get("metadata", {}) + candidate_names: list[str] = [] + metadata_file_path = "" + if metadata and isinstance(metadata, dict): + metadata_file_path = str(metadata.get("file_path") or "").strip() + _add_candidate(candidate_names, metadata.get("file_path")) + _add_candidate(candidate_names, metadata.get("original_name")) + + if metadata_file_path.startswith("images/"): + original_name = os.path.basename(metadata_file_path) + else: + original_name = None + + original_path = chunk.get("path", "") + if original_path: + _add_candidate(candidate_names, original_path) + normalized_path = original_path.replace("-->", "/") + original_name = os.path.basename(normalized_path) + + source_path, matched_name, ext = _resolve_image_source_path( + image_files_map, + candidate_names, + ) + if matched_name: + original_name = matched_name + + if not source_path: + raise StorageServiceException( + internal_message=( + "Cannot resolve image file for ZIP packaging: " + f"chunk_id={chunk_id}, candidates={candidate_names}" + ), + operation="collect_image_files", + ) + + width = None + height = None + try: + with Image.open(source_path) as img: + width, height = img.size + except Exception as exc: + logger.debug( + f"Failed to read image dimensions for ZIP resource {source_path}: {exc}" + ) + + file_size = os.path.getsize(source_path) + zip_path = _resolve_image_zip_path( + metadata=metadata, + original_name=original_name, + chunk_id=str(chunk_id), + ext=ext, + ) + + image_files.append( + { + "id": str(chunk_id), + "file_path": zip_path, + "original_name": original_name or f"image_{chunk_id}.{ext}", + "size_bytes": file_size, + "format": ext.lower(), + "width": width, + "height": height, + "source_path": source_path, + "zip_path": zip_path, + } + ) + + return image_files + + def _collect_table_files( + self, + chunks: list[dict[str, Any]], + tables_dir: str, + ) -> list[ZipResourceFileInfo]: + table_files: list[ZipResourceFileInfo] = [] + if not os.path.exists(tables_dir): + return table_files + + table_files_map = { + filename: file_path + for filename, file_path in _collect_files_by_name(tables_dir).items() + if filename.endswith(".html") + } + + for chunk in chunks: + chunk_id = chunk.get("chunk_id") or chunk.get("know_id") + chunk_type = chunk.get("type", "") + + if chunk_type != "table": + continue + + metadata = chunk.get("metadata", {}) + candidate_names: list[str] = [] + if metadata and isinstance(metadata, dict): + _add_candidate(candidate_names, metadata.get("file_path")) + _add_candidate(candidate_names, metadata.get("original_name")) + + original_path = chunk.get("path", "") + if original_path: + normalized_path = original_path.replace("-->", "/") + original_name = os.path.basename(normalized_path) + _add_candidate(candidate_names, original_path) + else: + original_name = None + + source_path, matched_name = _resolve_table_source_path( + table_files_map, + candidate_names, + ) + if matched_name: + original_name = matched_name + + if not source_path: + raise StorageServiceException( + internal_message=( + "Cannot resolve table file for ZIP packaging: " + f"chunk_id={chunk_id}, candidates={candidate_names}" + ), + operation="collect_table_files", + ) + + file_size = os.path.getsize(source_path) + zip_path = _resolve_table_zip_path( + metadata=metadata, + original_name=original_name, + chunk_id=str(chunk_id), + ) + + table_files.append( + { + "id": str(chunk_id), + "file_path": zip_path, + "original_name": original_name or f"table_{chunk_id}.html", + "size_bytes": file_size, + "format": "html", + "source_path": source_path, + "zip_path": zip_path, + } + ) + + return table_files + + +def _collect_files_by_name(directory_path: str) -> dict[str, str]: + files: dict[str, str] = {} + for filename in os.listdir(directory_path): + file_path = os.path.join(directory_path, filename) + if os.path.isfile(file_path): + files[filename] = file_path + return files + + +def _add_candidate(candidates: list[str], value: str | None) -> None: + if not value: + return + candidate = os.path.basename(str(value).strip().replace("-->", "/")) + if not candidate: + return + if candidate.startswith("[") and candidate.endswith("]"): + candidate = candidate[1:-1].strip() + if candidate and candidate not in candidates: + candidates.append(candidate) + + +def _resolve_image_source_path( + image_files_map: dict[str, str], + candidates: list[str], +) -> tuple[str | None, str | None, str]: + for candidate in candidates: + if candidate in image_files_map: + _, ext = os.path.splitext(candidate) + return image_files_map[candidate], candidate, ext.lstrip(".") or "jpg" + + stem, _ = os.path.splitext(candidate) + if stem: + stem_matches = [ + filename + for filename in image_files_map + if os.path.splitext(filename)[0] == stem + ] + if len(stem_matches) == 1: + matched = stem_matches[0] + _, matched_ext = os.path.splitext(matched) + return ( + image_files_map[matched], + matched, + matched_ext.lstrip(".") or "jpg", + ) + + return None, None, "jpg" + + +def _resolve_table_source_path( + table_files_map: dict[str, str], + candidates: list[str], +) -> tuple[str | None, str | None]: + for candidate in candidates: + if candidate in table_files_map: + return table_files_map[candidate], candidate + + stem, _ = os.path.splitext(candidate) + if stem: + stem_matches = [ + filename + for filename in table_files_map + if os.path.splitext(filename)[0] == stem + ] + if len(stem_matches) == 1: + matched = stem_matches[0] + return table_files_map[matched], matched + + return None, None + + +def _resolve_image_zip_path( + *, + metadata: Any, + original_name: str | None, + chunk_id: str, + ext: str, +) -> str: + if metadata and isinstance(metadata, dict): + zip_file_path = metadata.get("file_path") + if zip_file_path and zip_file_path.startswith("images/"): + return zip_file_path + if original_name: + return f"images/{original_name}" + return f"images/{chunk_id}.{ext}" + + if original_name: + return f"images/{original_name}" + return f"images/{chunk_id}.{ext}" + + +def _resolve_table_zip_path( + *, + metadata: Any, + original_name: str | None, + chunk_id: str, +) -> str: + if metadata and isinstance(metadata, dict): + zip_file_path = metadata.get("file_path") + if zip_file_path and zip_file_path.startswith("tables/"): + return zip_file_path + if original_name: + return f"tables/{original_name}" + return f"tables/{chunk_id}.html" + + if original_name: + return f"tables/{original_name}" + return f"tables/{chunk_id}.html" diff --git a/packages/shared-python/shared/services/storage/zip_result_schema.py b/packages/shared-python/shared/services/storage/zip_result_schema.py new file mode 100644 index 000000000..e444475a7 --- /dev/null +++ b/packages/shared-python/shared/services/storage/zip_result_schema.py @@ -0,0 +1,72 @@ +"""Schema projection facade for Knowhere ZIP result packages.""" + +from __future__ import annotations + +from typing import Any + +from shared.services.storage.zip_chunk_schema import ZipChunkSchemaBuilder +from shared.services.storage.zip_doc_navigation import ZipDocNavigationBuilder +from shared.services.storage.zip_manifest_schema import ZipManifestBuilder + + +class ZipResultSchemaBuilder: + def __init__( + self, + *, + chunk_schema: ZipChunkSchemaBuilder | None = None, + doc_navigation: ZipDocNavigationBuilder | None = None, + manifest_builder: ZipManifestBuilder | None = None, + ) -> None: + self._chunk_schema = chunk_schema or ZipChunkSchemaBuilder() + self._doc_navigation = doc_navigation or ZipDocNavigationBuilder() + self._manifest_builder = manifest_builder or ZipManifestBuilder() + + def calculate_statistics(self, chunks: list[dict[str, Any]]) -> dict[str, Any]: + return self._chunk_schema.calculate_statistics(chunks) + + def format_chunks( + self, + chunks: list[dict[str, Any]], + image_files_map: dict[str, dict[str, Any]], + table_files_map: dict[str, dict[str, Any]], + ) -> list[dict[str, Any]]: + return self._chunk_schema.format_chunks( + chunks, + image_files_map, + table_files_map, + ) + + def generate_manifest( + self, + *, + job_id: str, + data_id: str | None, + source_file_name: str, + statistics: dict[str, Any], + job_metadata: dict[str, Any], + hierarchy: dict[str, Any] | None = None, + ) -> dict[str, Any]: + return self._manifest_builder.generate_manifest( + job_id=job_id, + data_id=data_id, + source_file_name=source_file_name, + statistics=statistics, + job_metadata=job_metadata, + hierarchy=hierarchy, + ) + + def build_hierarchy_dict( + self, + sections: list[dict[str, Any]], + ) -> dict[str, Any]: + return self._doc_navigation.build_hierarchy_dict(sections) + + def build_doc_nav( + self, + formatted_chunks: list[dict[str, Any]], + source_file_name: str, + ) -> dict[str, Any]: + return self._doc_navigation.build_doc_nav( + formatted_chunks, + source_file_name, + ) diff --git a/packages/shared-python/shared/services/storage/zip_result_service.py b/packages/shared-python/shared/services/storage/zip_result_service.py index 68eb1b0fc..b2fee6e17 100644 --- a/packages/shared-python/shared/services/storage/zip_result_service.py +++ b/packages/shared-python/shared/services/storage/zip_result_service.py @@ -1,26 +1,14 @@ """ -ZIP Result Package Generation Service -Generates ZIP packages according to Knowhere-API-ZIP-Spec.md specification +ZIP Result Package Generation Service. + +Generates ZIP packages according to the Knowhere API ZIP result format. """ -import hashlib -import json -import os -import tempfile -import zipfile -from typing import Any, Dict, List, Optional, Tuple +from __future__ import annotations -from loguru import logger -from PIL import Image +from typing import Any -from shared.services.chunks.chunk_connections import ( - build_resource_target_map, - convert_refs_to_embed_connections, - merge_connections, - normalize_connect_to_targets, - parse_relationship_refs, -) -from shared.utils.text_utils import truncate_content_preview +from loguru import logger import pandas as pd @@ -28,849 +16,116 @@ KnowhereException, StorageServiceException, ) -from shared.utils.utc_now import utc_now_naive +from shared.services.storage.zip_package_writer import ( + ZipPackageWriter, + ZipPackageWriteRequest, +) +from shared.services.storage.zip_result_resources import ZipResourceCollector +from shared.services.storage.zip_result_schema import ZipResultSchemaBuilder class ZipResultService: - """ZIP Result Package Generation Service""" + """Generate a result ZIP while preserving the worker-facing package contract.""" - def __init__(self): - pass + def __init__( + self, + *, + schema_builder: ZipResultSchemaBuilder | None = None, + resource_collector: ZipResourceCollector | None = None, + package_writer: ZipPackageWriter | None = None, + ) -> None: + self._schema = schema_builder or ZipResultSchemaBuilder() + self._resources = resource_collector or ZipResourceCollector() + self._writer = package_writer or ZipPackageWriter() def generate_zip_package( self, job_id: str, - chunks: List[Dict[str, Any]], + chunks: list[dict[str, Any]], add_dir: str, source_file_name: str, - data_id: Optional[str], - job_metadata: Dict[str, Any], - parsed_df: Optional["pd.DataFrame"] = None, - temp_dir: Optional[str] = None, - ) -> Tuple[str, Dict[str, str], Dict[str, Any], int]: + data_id: str | None, + job_metadata: dict[str, Any], + parsed_df: pd.DataFrame | None = None, + temp_dir: str | None = None, + ) -> tuple[str, dict[str, str], dict[str, Any], int]: """ - Generate ZIP result package - - Args: - job_id: Job ID - chunks: List of chunks data - add_dir: Parsed directory path (contains images/ and tables/ directories) - source_file_name: Source file name - data_id: User-defined ID - job_metadata: Job metadata - parsed_df: Optional, parsed DataFrame (legacy, unused after doc_nav.json migration) - temp_dir: Optional directory for the generated ZIP file + Generate ZIP result package. Returns: Tuple[zip_file_path, checksum, statistics, zip_size]: - zip_file_path: ZIP file path - checksum: {"algorithm": "sha256", "value": "..."} - - statistics: {"total_chunks": int, "text_chunks": int, "image_chunks": int, "table_chunks": int, "total_pages": Optional[int]} + - statistics: chunk and page statistics persisted with the Job Result - zip_size: ZIP file size in bytes """ try: - # Create temporary ZIP file - effective_temp_dir = temp_dir or tempfile.gettempdir() - os.makedirs(effective_temp_dir, exist_ok=True) - zip_file_path = os.path.join(effective_temp_dir, f"result_{job_id}.zip") - - # Collect image and table file info (must be done before formatting chunks as file info is needed) - images_dir = os.path.join(add_dir, "images") - tables_dir = os.path.join(add_dir, "tables") - image_files_info = self._collect_image_files(chunks, images_dir) - table_files_info = self._collect_table_files(chunks, tables_dir) - - # Create image and table file mappings (chunk_id -> file_info) - image_files_map = {img["id"]: img for img in image_files_info} - table_files_map = {tb["id"]: tb for tb in table_files_info} - - # Convert chunks data format (using file info) - formatted_chunks = self._format_chunks( - chunks, image_files_map, table_files_map + resources = self._resources.collect(chunks=chunks, add_dir=add_dir) + formatted_chunks = self._schema.format_chunks( + chunks, + resources.image_files_map, + resources.table_files_map, ) - statistics = self._calculate_statistics(formatted_chunks) - - doc_nav: Dict[str, Any] = {} - hierarchy: Dict[str, Any] = {} - - # Create ZIP package - with zipfile.ZipFile(zip_file_path, "w", zipfile.ZIP_DEFLATED) as zip_file: - # 1. Generate chunks.json (full version) - chunks_json = json.dumps( - {"chunks": formatted_chunks}, ensure_ascii=False, indent=2 - ) - zip_file.writestr("chunks.json", chunks_json.encode("utf-8")) - - # 2. Try to add full.md (if exists) - full_md_path = os.path.join(add_dir, "full.md") - if os.path.exists(full_md_path): - zip_file.write(full_md_path, "full.md") - - # 2b. Try to add toc_hierarchies.json (if exists) - toc_path = os.path.join(add_dir, "toc_hierarchies.json") - if os.path.exists(toc_path): - zip_file.write(toc_path, "toc_hierarchies.json") - logger.info("Added toc_hierarchies.json to ZIP") - - # 3. Add image files - for img_info in image_files_info: - source_path = img_info["source_path"] - if os.path.exists(source_path): - zip_file.write( - source_path, - img_info["zip_path"], - ) - else: - logger.warning(f"Image file not found: {source_path}") - - # 4. Add table files - for table_info in table_files_info: - source_path = table_info["source_path"] - if os.path.exists(source_path): - zip_file.write( - source_path, - table_info["zip_path"], - ) - else: - logger.warning(f"Table file not found: {source_path}") + statistics = self._schema.calculate_statistics(formatted_chunks) - # 5. Generate doc_nav.json — unified navigation file - try: - doc_nav = self._build_doc_nav(formatted_chunks, source_file_name) - hierarchy = self._build_hierarchy_dict(doc_nav.get("sections", [])) - doc_nav_json = json.dumps(doc_nav, ensure_ascii=False, indent=2) - zip_file.writestr("doc_nav.json", doc_nav_json.encode("utf-8")) - logger.info("Added doc_nav.json") - except Exception as e: - logger.warning(f"generate doc_nav.json fail {e}") - - # 6. Generate manifest.json (checksum not included, stored in database) - manifest = self._generate_manifest( + doc_nav, hierarchy = self._build_navigation_outputs( + formatted_chunks=formatted_chunks, + source_file_name=source_file_name, + ) + manifest = self._schema.generate_manifest( + job_id=job_id, + data_id=data_id, + source_file_name=source_file_name, + statistics=statistics, + job_metadata=job_metadata, + hierarchy=hierarchy, + ) + artifact = self._writer.write( + ZipPackageWriteRequest( job_id=job_id, - data_id=data_id, - source_file_name=source_file_name, - statistics=statistics, - job_metadata=job_metadata, - hierarchy=hierarchy, + add_dir=add_dir, + formatted_chunks=formatted_chunks, + image_files=resources.image_files, + table_files=resources.table_files, + doc_nav=doc_nav, + manifest=manifest, + temp_dir=temp_dir, ) - manifest_json = json.dumps(manifest, ensure_ascii=False, indent=2) - zip_file.writestr("manifest.json", manifest_json.encode("utf-8")) - - # Calculate ZIP package SHA-256 - checksum_value = self._calculate_zip_checksum(zip_file_path) - checksum = {"algorithm": "sha256", "value": checksum_value} - - # Get ZIP file size - zip_size = os.path.getsize(zip_file_path) + ) logger.info( - f"ZIP package generated successfully: job_id={job_id}, size={zip_size}, checksum={checksum_value[:16]}..." + "ZIP package generated successfully: " + f"job_id={job_id}, size={artifact.zip_size}, " + f"checksum={artifact.checksum['value'][:16]}..." ) - return zip_file_path, checksum, statistics, zip_size + return ( + artifact.zip_file_path, + artifact.checksum, + statistics, + artifact.zip_size, + ) except KnowhereException: raise - except Exception as e: - logger.error(f"Failed to generate ZIP package: {e}") + except Exception as exc: + logger.error(f"Failed to generate ZIP package: {exc}") raise StorageServiceException( - internal_message=f"Failed to generate ZIP package: {str(e)}", + internal_message=f"Failed to generate ZIP package: {str(exc)}", operation="generate_zip_package", - original_exception=e, - ) - - def _calculate_statistics(self, chunks: List[Dict[str, Any]]) -> Dict[str, Any]: - """Calculate statistics""" - total_chunks = len(chunks) - text_chunks = 0 - image_chunks = 0 - table_chunks = 0 - - for chunk in chunks: - chunk_type = chunk.get("type", "") - raw_type = str(chunk_type).strip() - normalized_type = raw_type.split("\n", 1)[0].lower() - if normalized_type == "image": - image_chunks += 1 - elif normalized_type == "table": - table_chunks += 1 - else: - text_chunks += 1 - - return { - "total_chunks": total_chunks, - "text_chunks": text_chunks, - "image_chunks": image_chunks, - "table_chunks": table_chunks, - "total_pages": None, # Cannot determine page count at this point - } - - def _format_chunks( - self, - chunks: List[Dict[str, Any]], - image_files_map: Dict[str, Dict[str, Any]], - table_files_map: Dict[str, Dict[str, Any]], - ) -> List[Dict[str, Any]]: - """Convert chunks data to ZIP specification format""" - resource_target_map = build_resource_target_map( - chunks, - image_files_map=image_files_map, - table_files_map=table_files_map, - ) - - formatted = [] - for chunk in chunks: - chunk_id = str(chunk.get("chunk_id") or chunk.get("know_id")) - chunk_type_str = chunk.get("type", "") - raw_type = str(chunk_type_str).strip() - normalized_type = raw_type.split("\n", 1)[0].lower() - img_info = image_files_map.get(chunk_id) - - # Determine chunk type - if normalized_type == "image": - chunk_type = "image" - elif normalized_type == "table": - chunk_type = "table" - else: - chunk_type = "text" - - # Get content - content = chunk.get("text") or chunk.get("content", "") - - # Use original path directly to match kb.csv - path = chunk.get("path", "") - - # Get or build base metadata - existing_metadata = chunk.get("metadata", {}) - metadata = { - "length": existing_metadata.get("length") or len(content), - "summary": existing_metadata.get("summary") or chunk.get("summary", ""), - "page_nums": existing_metadata.get("page_nums", []), - } - document_top_summary = str( - existing_metadata.get("document_top_summary") or "" - ).strip() - if document_top_summary: - metadata["document_top_summary"] = document_top_summary - - # Add type-specific fields - if chunk_type == "text": - metadata["tokens"] = existing_metadata.get("tokens") or chunk.get( - "tokens", 0 - ) - metadata["keywords"] = existing_metadata.get("keywords") or chunk.get( - "keywords", [] - ) - - # Convert in-text resource refs into embeds edges. - relationship_refs = parse_relationship_refs( - chunk.get("type_raw") or chunk_type_str, - str(content), - ) - embed_connections = convert_refs_to_embed_connections( - relationship_refs, resource_target_map - ) - related_connections = normalize_connect_to_targets( - existing_metadata.get("connect_to") - or chunk.get("connect_to") - or chunk.get("connectto"), - resource_target_map, - ) - metadata["connect_to"] = merge_connections( - embed_connections, related_connections - ) - - elif chunk_type == "image": - if img_info: - metadata["file_path"] = img_info["file_path"] - # Unified schema: include keywords and tokens for all chunk types - metadata["keywords"] = existing_metadata.get("keywords") or chunk.get( - "keywords", [] - ) - metadata["tokens"] = [] - - elif chunk_type == "table": - # Get table info from existing_metadata or table_files_map - file_path = existing_metadata.get("file_path") - - if not file_path: - # Get table info from table_files_map - tb_info = table_files_map.get(chunk_id) - if tb_info: - file_path = tb_info["file_path"] - else: - # Extract from path or use default - tbl_name = ( - path.split("/")[-1] - if "/" in path - else f"table_{chunk_id}.html" - ) - file_path = f"tables/{tbl_name}" - - metadata["file_path"] = file_path - # Unified schema: include keywords and tokens for all chunk types - metadata["keywords"] = existing_metadata.get("keywords") or chunk.get( - "keywords", [] - ) - metadata["tokens"] = [] - - formatted_chunk = { - "chunk_id": chunk_id, - "type": chunk_type, - "content": content, - "path": path, - "metadata": metadata, - } - formatted.append(formatted_chunk) - - return formatted - - def _clean_path(self, path: str) -> str: - """Clean path, keep only logical path""" - if not path: - return "/" - - # Remove filesystem path prefix - # Example: .-->users-->KB_DATA_xxx-->dir-->file.pdf-->chapter-->section - # Should extract: chapter-->section - - # Find the last .pdf, .docx, etc. file extension - import re - - # Match filename pattern (with extension) - file_pattern = r"[^/]+\.(pdf|docx|doc|txt|md|xlsx|xls|pptx|ppt)" - match = re.search(file_pattern, path, re.IGNORECASE) - - if match: - # Extract the part after filename - path_after_file = path[match.end() :] - # Clean path separators - path_after_file = path_after_file.replace("-->", "/").strip("/") - if path_after_file: - return path_after_file - - # If no file pattern found, try to clean common prefixes - path = path.replace("-->", "/") - # Remove leading path separators and empty segments - path = "/".join( - [p for p in path.split("/") if p and p not in ["", ".", "users"]] - ) - return path if path else "/" - - def _collect_image_files( - self, chunks: List[Dict[str, Any]], images_dir: str - ) -> List[Dict[str, Any]]: - """Collect image file information""" - image_files = [] - if not os.path.exists(images_dir): - has_image_chunks = any(chunk.get("type", "") == "image" for chunk in chunks) - if has_image_chunks: - raise StorageServiceException( - internal_message=( - "Image directory not found for ZIP packaging: " - f"images_dir={images_dir}" - ), - operation="collect_image_files", - ) - return image_files - - # Get all image files - image_files_map = {} - for filename in os.listdir(images_dir): - file_path = os.path.join(images_dir, filename) - if os.path.isfile(file_path): - image_files_map[filename] = file_path - - def add_candidate(candidates: List[str], value: Optional[str]) -> None: - if not value: - return - candidate = os.path.basename(str(value).strip().replace("-->", "/")) - if not candidate: - return - if candidate.startswith("[") and candidate.endswith("]"): - candidate = candidate[1:-1].strip() - if candidate and candidate not in candidates: - candidates.append(candidate) - - def resolve_source_path( - candidates: List[str], chunk_id: str - ) -> Tuple[Optional[str], Optional[str], str]: - for candidate in candidates: - if candidate in image_files_map: - _, ext = os.path.splitext(candidate) - return ( - image_files_map[candidate], - candidate, - ext.lstrip(".") or "jpg", - ) - - stem, ext = os.path.splitext(candidate) - if stem: - stem_matches = [ - filename - for filename in image_files_map - if os.path.splitext(filename)[0] == stem - ] - if len(stem_matches) == 1: - matched = stem_matches[0] - _, matched_ext = os.path.splitext(matched) - return ( - image_files_map[matched], - matched, - matched_ext.lstrip(".") or "jpg", - ) - - return None, None, "jpg" - - # Match images from chunks - for chunk in chunks: - chunk_id = chunk.get("chunk_id") or chunk.get("know_id") - chunk_type = chunk.get("type", "") - - if chunk_type != "image": - continue - - metadata = chunk.get("metadata", {}) - candidate_names: List[str] = [] - metadata_file_path = "" - if metadata and isinstance(metadata, dict): - metadata_file_path = str(metadata.get("file_path") or "").strip() - add_candidate(candidate_names, metadata.get("file_path")) - add_candidate(candidate_names, metadata.get("original_name")) - - # Try to get original filename from chunk's path field - if metadata_file_path.startswith("images/"): - original_name = os.path.basename(metadata_file_path) - else: - original_name = None - - original_path = chunk.get("path", "") - if original_path: - add_candidate(candidate_names, original_path) - # Normalize path separators: replace --> with /, then extract filename - normalized_path = original_path.replace("-->", "/") - original_name = os.path.basename(normalized_path) - - source_path, matched_name, ext = resolve_source_path( - candidate_names, str(chunk_id) - ) - if matched_name: - original_name = matched_name - - if not source_path: - raise StorageServiceException( - internal_message=( - "Cannot resolve image file for ZIP packaging: " - f"chunk_id={chunk_id}, candidates={candidate_names}" - ), - operation="collect_image_files", - ) - - # Get image dimensions - width = None - height = None - try: - with Image.open(source_path) as img: - width, height = img.size - except Exception: - pass - - file_size = os.path.getsize(source_path) - - # Priority: use file_path from metadata, then original_name, finally chunk_id - if metadata and isinstance(metadata, dict): - # metadata.file_path format: "images/xxx.jpg" - zip_file_path = metadata.get("file_path") - if zip_file_path and zip_file_path.startswith("images/"): - # Use complete path from metadata - zip_path = zip_file_path - # Extract filename as original_name - if not original_name: - original_name = metadata.get( - "original_name" - ) or os.path.basename(zip_file_path) - else: - # If metadata has no file_path, use original_name or chunk_id - if original_name: - zip_path = f"images/{original_name}" - else: - zip_path = f"images/{chunk_id}.{ext}" - else: - # If no metadata, use original_name or chunk_id - if original_name: - zip_path = f"images/{original_name}" - else: - zip_path = f"images/{chunk_id}.{ext}" - - image_files.append( - { - "id": str(chunk_id), - "file_path": zip_path, - "original_name": original_name or f"image_{chunk_id}.{ext}", - "size_bytes": file_size, - "format": ext.lower(), - "width": width, - "height": height, - "source_path": source_path, - "zip_path": zip_path, - } - ) - - return image_files - - def _collect_table_files( - self, chunks: List[Dict[str, Any]], tables_dir: str - ) -> List[Dict[str, Any]]: - """Collect table file information""" - table_files = [] - if not os.path.exists(tables_dir): - return table_files - - # Get all table files - table_files_map = {} - for filename in os.listdir(tables_dir): - file_path = os.path.join(tables_dir, filename) - if os.path.isfile(file_path) and filename.endswith(".html"): - table_files_map[filename] = file_path - - def add_candidate(candidates: List[str], value: Optional[str]) -> None: - if not value: - return - candidate = os.path.basename(str(value).strip().replace("-->", "/")) - if not candidate: - return - if candidate.startswith("[") and candidate.endswith("]"): - candidate = candidate[1:-1].strip() - if candidate and candidate not in candidates: - candidates.append(candidate) - - def resolve_source_path( - candidates: List[str], - ) -> Tuple[Optional[str], Optional[str]]: - for candidate in candidates: - if candidate in table_files_map: - return table_files_map[candidate], candidate - - stem, _ = os.path.splitext(candidate) - if stem: - stem_matches = [ - filename - for filename in table_files_map - if os.path.splitext(filename)[0] == stem - ] - if len(stem_matches) == 1: - matched = stem_matches[0] - return table_files_map[matched], matched - - return None, None - - # Match tables from chunks - for chunk in chunks: - chunk_id = chunk.get("chunk_id") or chunk.get("know_id") - chunk_type = chunk.get("type", "") - - if chunk_type != "table": - continue - - metadata = chunk.get("metadata", {}) - candidate_names: List[str] = [] - if metadata and isinstance(metadata, dict): - add_candidate(candidate_names, metadata.get("file_path")) - add_candidate(candidate_names, metadata.get("original_name")) - - # Try to get original filename from chunk's path field - original_path = chunk.get("path", "") - if original_path: - # Normalize path separators: replace --> with /, then extract filename - normalized_path = original_path.replace("-->", "/") - original_name = os.path.basename(normalized_path) - add_candidate(candidate_names, original_path) - else: - original_name = None - - source_path, matched_name = resolve_source_path(candidate_names) - if matched_name: - original_name = matched_name - - if not source_path: - raise StorageServiceException( - internal_message=( - "Cannot resolve table file for ZIP packaging: " - f"chunk_id={chunk_id}, candidates={candidate_names}" - ), - operation="collect_table_files", - ) - - file_size = os.path.getsize(source_path) - - # Priority: use file_path from metadata, then original_name, finally chunk_id - if metadata and isinstance(metadata, dict): - # metadata.file_path format: "tables/xxx.html" - zip_file_path = metadata.get("file_path") - if zip_file_path and zip_file_path.startswith("tables/"): - # Use complete path from metadata - zip_path = zip_file_path - # Extract filename as original_name - if not original_name: - original_name = metadata.get( - "original_name" - ) or os.path.basename(zip_file_path) - else: - # If metadata has no file_path, use original_name or chunk_id - if original_name: - zip_path = f"tables/{original_name}" - else: - zip_path = f"tables/{chunk_id}.html" - else: - # If no metadata, use original_name or chunk_id - if original_name: - zip_path = f"tables/{original_name}" - else: - zip_path = f"tables/{chunk_id}.html" - - table_files.append( - { - "id": str(chunk_id), - "file_path": zip_path, - "original_name": original_name or f"table_{chunk_id}.html", - "size_bytes": file_size, - "format": "html", - "source_path": source_path, - "zip_path": zip_path, - } + original_exception=exc, ) - return table_files - - def _generate_manifest( + def _build_navigation_outputs( self, - job_id: str, - data_id: Optional[str], + *, + formatted_chunks: list[dict[str, Any]], source_file_name: str, - statistics: Dict[str, Any], - job_metadata: Dict[str, Any], - hierarchy: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """Generate manifest.json""" - manifest = { - "version": "2.0", - "job_id": job_id, - "data_id": data_id, - "source_file_name": source_file_name, - "processing_date": utc_now_naive().isoformat() + "Z", - "processing": { - "page_count": job_metadata.get("page_count"), - "billing_status": job_metadata.get("billing_status"), - "cost": { - "micro_dollars": job_metadata.get("billing_amount_micro_dollars"), - "credits": job_metadata.get("billing_credits"), - }, - "timing": { - "started_at": job_metadata.get("processing_started_at"), - "completed_at": job_metadata.get("processing_completed_at"), - "duration_ms": job_metadata.get("processing_duration_ms"), - }, - }, - "statistics": statistics, - "HIERARCHY": hierarchy or {}, - } - - return manifest - - def _calculate_zip_checksum(self, zip_file_path: str) -> str: - """Calculate SHA-256 checksum of ZIP file""" - sha256_hash = hashlib.sha256() - with open(zip_file_path, "rb") as f: - for byte_block in iter(lambda: f.read(4096), b""): - sha256_hash.update(byte_block) - return sha256_hash.hexdigest().lower() - - def _build_hierarchy_dict( - self, - sections: List[Dict[str, Any]], - ) -> Dict[str, Any]: - """Build a title-only nested hierarchy from doc_nav sections.""" - hierarchy: Dict[str, Any] = {} - title_counts: Dict[str, int] = {} - - for section in sections: - raw_title = str(section.get("title") or "").strip() - if not raw_title: - continue - - title_counts[raw_title] = title_counts.get(raw_title, 0) + 1 - title = ( - raw_title - if title_counts[raw_title] == 1 - else f"{raw_title} ({title_counts[raw_title]})" - ) - hierarchy[title] = self._build_hierarchy_dict( - section.get("children") or [] - ) - - return hierarchy - - def _build_doc_nav( - self, - formatted_chunks: List[Dict[str, Any]], - source_file_name: str, - ) -> Dict[str, Any]: - """Build doc_nav.json — unified navigation file. - - Structured file serving both human demo and LLM navigation. - - The output contains: - - ``sections``: tree of text sections with summaries and chunk counts. - - ``resources``: flat lists of image/table chunks with summaries. - - ``stats``: chunk counts by type. - - Each leaf section carries a ``summary`` derived from: - 1. chunk.metadata.summary (LLM-generated, highest quality) - 2. chunk.content[:300] (fallback truncation) - - Non-leaf section summaries are left empty at this stage and are - populated later by ``summary_builder.enrich_doc_nav_summaries``. - """ - # ── Separate text chunks from resource chunks ── - text_chunks: List[Dict[str, Any]] = [] - image_resources: List[Dict[str, Any]] = [] - table_resources: List[Dict[str, Any]] = [] - - stats = {"total_chunks": 0, "text_chunks": 0, "image_chunks": 0, "table_chunks": 0, "max_depth": 0} - - for fc in formatted_chunks: - ctype = fc.get("type", "text") - path = fc.get("path", "") - meta = fc.get("metadata") or {} - summary_raw = (meta.get("summary") or "").strip() - content_raw = (fc.get("content") or "").strip() - # Normalize whitespace - summary = " ".join(summary_raw.split()) if summary_raw else "" - content_preview = truncate_content_preview(content_raw) if content_raw else "" - - stats["total_chunks"] += 1 - - if ctype == "image": - stats["image_chunks"] += 1 - image_resources.append({ - "path": path, - "summary": summary or content_preview, - }) - elif ctype == "table": - stats["table_chunks"] += 1 - table_resources.append({ - "path": path, - "summary": summary or content_preview, - }) - else: - stats["text_chunks"] += 1 - text_chunks.append({ - "path": path, - "summary": summary or content_preview, - }) - - # ── Build section tree from text chunk paths ── - # Each text chunk path looks like: "kb_root/filename.pdf/Section/Subsection" - # We strip the kb_root and filename prefix to get relative section paths. - sections = self._build_section_tree(text_chunks) - - # Compute max depth - def _max_depth(nodes: list, d: int = 1) -> int: - m = d if nodes else 0 - for n in nodes: - m = max(m, _max_depth(n.get("children", []), d + 1)) - return m - - stats["max_depth"] = _max_depth(sections) - - return { - "version": "1.0", - "file_name": source_file_name or "", - "stats": stats, - "sections": sections, - "resources": { - "images": image_resources, - "tables": table_resources, - }, - } - - def _build_section_tree( - self, - text_chunks: List[Dict[str, Any]], - ) -> List[Dict[str, Any]]: - """Build a tree of sections from flat text chunk paths. - - Each text chunk has a ``path`` like ``"kb/file.pdf/Sec1/Sub1"``. - We extract section parts (after kb_root + filename) and build a - tree using ``children`` arrays. - - Returns a list of top-level section nodes. - """ - # Internal tree node: {title, summary, chunk_count, children: {title: node}} - root_children: Dict[str, dict] = {} # ordered dict of top-level titles - - for chunk in text_chunks: - path = chunk.get("path", "") - parts = [p.strip() for p in path.split("/") if p.strip()] - # Skip kb_root + filename → section parts start at index 2 - section_parts = parts[2:] if len(parts) > 2 else [] - - if not section_parts: - # Root-level chunk (no section hierarchy) - key = "__root__" - if key not in root_children: - root_children[key] = { - "title": "Root", - "path": "/".join(parts[:2]) if len(parts) >= 2 else path, - "summary": chunk.get("summary", ""), - "chunk_count": 0, - "_children_map": {}, - } - root_children[key]["chunk_count"] += 1 - # Use the first chunk's summary for root - if not root_children[key]["summary"]: - root_children[key]["summary"] = chunk.get("summary", "") - continue - - # Walk the tree, creating nodes as needed - current_level = root_children - full_section_path_parts = parts[:2] # start with kb_root/filename - for i, part in enumerate(section_parts): - full_section_path_parts.append(part) - if part not in current_level: - current_level[part] = { - "title": part, - "path": "/".join(full_section_path_parts), - "summary": "", - "chunk_count": 0, - "_children_map": {}, - } - node = current_level[part] - if i == len(section_parts) - 1: - # Leaf — this is the chunk's actual section - node["chunk_count"] += 1 - if not node["summary"]: - node["summary"] = chunk.get("summary", "") - current_level = node["_children_map"] - - # Convert internal tree to output format - def _to_output(children_map: Dict[str, dict], level: int = 1) -> List[Dict[str, Any]]: - result = [] - for node in children_map.values(): - children = _to_output(node["_children_map"], level + 1) - # Compute total chunk_count including descendants - total_chunks = node["chunk_count"] + sum( - c.get("chunk_count", 0) for c in children - ) - out = { - "title": node["title"], - "path": node["path"], - "level": level, - "summary": node["summary"], - "chunk_count": total_chunks, - "children": children, - } - result.append(out) - return result - - return _to_output(root_children) + ) -> tuple[dict[str, Any] | None, dict[str, Any]]: + try: + doc_nav = self._schema.build_doc_nav(formatted_chunks, source_file_name) + hierarchy = self._schema.build_hierarchy_dict(doc_nav.get("sections", [])) + return doc_nav, hierarchy + except Exception as exc: + logger.warning(f"generate doc_nav.json fail {exc}") + return None, {} diff --git a/packages/shared-python/shared/services/webhook/delivery_client.py b/packages/shared-python/shared/services/webhook/delivery_client.py new file mode 100644 index 000000000..2e33aca20 --- /dev/null +++ b/packages/shared-python/shared/services/webhook/delivery_client.py @@ -0,0 +1,151 @@ +"""Pinned HTTP delivery for outbound webhooks.""" + +import asyncio +import time +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from loguru import logger + +from shared.services.http.pinned_outbound import send_pinned_outbound_request +from shared.services.http.url_security import validate_http_url_and_resolve_ip_async + +HTTP_TIMEOUT_SECONDS = 10 + + +@dataclass(frozen=True) +class WebhookDeliveryTarget: + target_url: str + pinned_ip: str + + +@dataclass(frozen=True) +class WebhookDeliveryResult: + success: bool + status_code: int | None + duration_ms: int + error_message: str | None + + +@dataclass(frozen=True) +class WebhookTargetValidation: + target: WebhookDeliveryTarget | None + failure: WebhookDeliveryResult | None + + +class WebhookDeliveryClient: + """Validate and send direct webhook HTTP requests with DNS pinning.""" + + async def validate_target( + self, *, event_id: str, target_url: str + ) -> WebhookTargetValidation: + validation = await validate_http_url_and_resolve_ip_async(target_url) + + if not validation.is_valid: + logger.warning( + f"SSRF validation failed: event_id={event_id}, " + f"error={validation.error_message}" + ) + return WebhookTargetValidation( + target=None, + failure=WebhookDeliveryResult( + success=False, + status_code=400, + duration_ms=0, + error_message=f"SSRF: {validation.error_message}", + ), + ) + + if not validation.validated_ip: + return WebhookTargetValidation( + target=None, + failure=WebhookDeliveryResult( + success=False, + status_code=400, + duration_ms=0, + error_message="SSRF validation did not return a pinned IP", + ), + ) + + return WebhookTargetValidation( + target=WebhookDeliveryTarget( + target_url=target_url, + pinned_ip=validation.validated_ip, + ), + failure=None, + ) + + async def post_json( + self, + *, + event_id: str, + target: WebhookDeliveryTarget, + payload: Mapping[str, Any], + headers: Mapping[str, str], + ) -> WebhookDeliveryResult: + start_time = time.time() + + try: + response = await send_pinned_outbound_request( + method="POST", + url=target.target_url, + pinned_ip=target.pinned_ip, + timeout_seconds=HTTP_TIMEOUT_SECONDS, + headers=headers, + json_body=payload, + ) + duration_ms = int((time.time() - start_time) * 1000) + + if 200 <= response.status < 300: + logger.info( + f"Webhook delivered: event_id={event_id}, status={response.status}" + ) + return WebhookDeliveryResult( + success=True, + status_code=response.status, + duration_ms=duration_ms, + error_message=None, + ) + + if 300 <= response.status < 400: + logger.warning( + f"Webhook redirect blocked (SSRF protection): " + f"event_id={event_id}, status={response.status}" + ) + return WebhookDeliveryResult( + success=False, + status_code=response.status, + duration_ms=duration_ms, + error_message=f"Redirect blocked: HTTP {response.status}", + ) + + logger.warning( + f"Webhook failed: event_id={event_id}, status={response.status}" + ) + return WebhookDeliveryResult( + success=False, + status_code=response.status, + duration_ms=duration_ms, + error_message=f"HTTP {response.status}", + ) + + except asyncio.TimeoutError: + duration_ms = int((time.time() - start_time) * 1000) + logger.error(f"Webhook timeout: event_id={event_id}") + return WebhookDeliveryResult( + success=False, + status_code=None, + duration_ms=duration_ms, + error_message="Connection timeout", + ) + + except Exception as error: + duration_ms = int((time.time() - start_time) * 1000) + logger.error(f"Webhook error: event_id={event_id}, error={error}") + return WebhookDeliveryResult( + success=False, + status_code=None, + duration_ms=duration_ms, + error_message=str(error), + ) diff --git a/packages/shared-python/shared/services/webhook/dispatcher.py b/packages/shared-python/shared/services/webhook/dispatcher.py index e5ee4610d..0f3c4b953 100644 --- a/packages/shared-python/shared/services/webhook/dispatcher.py +++ b/packages/shared-python/shared/services/webhook/dispatcher.py @@ -1,44 +1,25 @@ """ Webhook Dispatcher Service -Dispatches webhook events via HTTP requests with HMAC signing and delivery logging. -Called by Celery task for async processing. +Dispatches webhook events with retry policy. Direct HTTP delivery details live +behind WebhookEventDelivery. """ -import asyncio -import hashlib -import hmac -import json import threading -import time -import uuid from datetime import datetime, timezone -from typing import Any, Dict, Optional, Tuple +from typing import Optional from loguru import logger from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -# Use standard db context - run_async_task handles the loop reuse from shared.core.database import get_db_context -from shared.core.exceptions.domain_exceptions import ( - SystemSettingInvalidException, - SystemSettingMissingException, -) from shared.core.exceptions.webhook_exceptions import WebhookDeliveryException -from shared.models.database.job import Job from shared.models.database.webhook import WebhookEvent, WebhookEventStatus -from shared.models.database.webhook_log import WebhookLog -from shared.utils.pinned_outbound_http import ( - send_pinned_outbound_request, -) -from shared.utils.url_security import ( - HTTPURLValidationResult, - validate_http_url_and_resolve_ip_async, -) +from shared.services.webhook.delivery_client import WebhookDeliveryResult +from shared.services.webhook.event_delivery import WebhookEventDelivery # Constants -HTTP_TIMEOUT_SECONDS = 10 MAX_ATTEMPTS = 6 @@ -54,6 +35,9 @@ class WebhookDispatcher: 5. On failure, signals the caller to schedule retry """ + def __init__(self, event_delivery: WebhookEventDelivery | None = None) -> None: + self._event_delivery = event_delivery or WebhookEventDelivery() + async def dispatch(self, event_id: str) -> bool: """ Dispatch a webhook event. @@ -87,27 +71,24 @@ async def dispatch(self, event_id: str) -> bool: await self._mark_failed(db, event) return True # ACK - # 4. Dispatch the webhook - # Logging is now handled inside _send_webhook - success, status_code, duration_ms, error_message = await self._send_webhook( + delivery_result = await self._event_delivery.send( db=db, event=event, is_manual=False ) - # 6. Handle result (Logging already done) - if success: + if delivery_result.success: await self._mark_delivered(db, event) return True # Success else: # Determine if error is retryable # Retryable: 5xx, timeout (None), 429 (rate limit) # NOT retryable: 4xx (except 429) - client errors won't be fixed by retrying - is_retryable = self._is_retryable_error(status_code) + is_retryable = self._is_retryable_error(delivery_result.status_code) if not is_retryable: # Permanent failure - don't retry logger.warning( f"WebhookEvent permanent failure (non-retryable): " - f"event_id={event_id}, status={status_code}" + f"event_id={event_id}, status={delivery_result.status_code}" ) await self._mark_failed(db, event) return True # ACK - no point retrying @@ -121,9 +102,11 @@ async def dispatch(self, event_id: str) -> bool: # Raise exception so Celery task will retry raise WebhookDeliveryException( - internal_message=f"Webhook delivery failed: {error_message}", + internal_message=( + f"Webhook delivery failed: {delivery_result.error_message}" + ), retryable=True, - status_code=status_code, + status_code=delivery_result.status_code, ) async def mark_event_failed(self, event_id: str) -> None: @@ -147,226 +130,11 @@ async def _fetch_event( ) return result.scalar_one_or_none() - async def _send_webhook( - self, db: AsyncSession, event: WebhookEvent, is_manual: bool = False - ) -> Tuple[bool, Optional[int], int, Optional[str]]: - """ - Send HTTP POST request to webhook target and log the attempt. - - Args: - db: Database session - event: WebhookEvent object - is_manual: True if manually triggered (adds 'trigger': 'manual' to payload) - - Returns: - Tuple of (success, status_code, duration_ms, error_message) - """ - - # Generate attempt ID - attempt_id = str(uuid.uuid4()) - - # SSRF Protection - validation: HTTPURLValidationResult = await validate_http_url_and_resolve_ip_async( - event.target_url, - ) - if not validation.is_valid: - logger.warning( - f"SSRF validation failed: event_id={event.id}, error={validation.error_message}" - ) - return False, 400, 0, f"SSRF: {validation.error_message}" - - # Enrich payload with job result data at delivery time - enriched_payload = await self._enrich_payload(event) - - # Add manual mark if requested - if is_manual: - enriched_payload["trigger"] = "manual" - - # Helper to get user_id from job - async def _get_job_owner(job_id: str) -> Optional[str]: - result = await db.execute(select(Job.user_id).where(Job.job_id == job_id)) - return result.scalar_one_or_none() - - # Resolve secret (Lazy creation) - secret = None - try: - user_id = await _get_job_owner(event.job_id) - if user_id: - secret = await self._resolve_secret(db, user_id, event.target_url) - else: - logger.warning( - f"Could not resolve secret: Job {event.job_id} has no user_id" - ) - except (SystemSettingMissingException, SystemSettingInvalidException) as e: - logger.error(f"Configuration error during secret resolution: {e}") - # Return 424 (Failed Dependency) to ensure it's treated as a non-retryable error - return False, 424, 0, f"Configuration Error: {e}" - except Exception as e: - logger.error(f"Secret resolution failed: {e}") - - if not secret: - logger.error(f"No secret found or created/resolved for event {event.id}") - # Default to non-retryable error for any secret resolution failure - return False, 424, 0, "Secret resolution failed" - - # Sign payload - signature = self._sign_payload(enriched_payload, secret) - - # Build headers - headers = { - "Content-Type": "application/json", - "X-Knowhere-Signature": signature, - "X-Knowhere-Attempt-ID": attempt_id, - "User-Agent": "Knowhere-Webhook/1.0", - } - - start_time = time.time() - status_code = None - error_message = None - success = False - - try: - pinned_ip = validation.validated_ip - if not pinned_ip: - return False, 400, 0, "SSRF validation did not return a pinned IP" - - response = await send_pinned_outbound_request( - method="POST", - url=event.target_url, - pinned_ip=pinned_ip, - timeout_seconds=HTTP_TIMEOUT_SECONDS, - headers=headers, - json_body=enriched_payload, - ) - duration_ms = int((time.time() - start_time) * 1000) - status_code = response.status - - if 200 <= response.status < 300: - logger.info( - f"Webhook delivered: event_id={event.id}, status={response.status}" - ) - success = True - elif 300 <= response.status < 400: - logger.warning( - f"Webhook redirect blocked (SSRF protection): " - f"event_id={event.id}, status={response.status}" - ) - error_message = f"Redirect blocked: HTTP {response.status}" - success = False - else: - logger.warning( - f"Webhook failed: event_id={event.id}, status={response.status}" - ) - error_message = f"HTTP {response.status}" - success = False - - except asyncio.TimeoutError: - duration_ms = int((time.time() - start_time) * 1000) - logger.error(f"Webhook timeout: event_id={event.id}") - error_message = "Connection timeout" - success = False - - except Exception as e: - duration_ms = int((time.time() - start_time) * 1000) - logger.error(f"Webhook error: event_id={event.id}, error={e}") - error_message = str(e) - success = False - - # Log delivery attempt - # If manual, event_id is None to avoid FK violation - log_event_id = None if is_manual else event.id - - try: - # Combine headers and payload - combined_payload = {"header": headers, "payload": enriched_payload} - - log = WebhookLog( - job_id=event.job_id, - event_id=log_event_id, - webhook_url=event.target_url, - attempt_number=event.attempts + 1, - request_payload=combined_payload, - signature=signature, - idempotency_key=str(uuid.uuid4()), - response_status_code=status_code, - error_message=error_message, - duration_ms=duration_ms, - ) - db.add(log) - # If auto-commit is needed? - # Dispatcher.dispatch uses passed 'db' session which is managed by 'async with get_db_context()'. - # It commits inside _mark_delivered etc. - # We should probably commit/flush here to persist log even if update fails? - await db.commit() - - except Exception as e: - logger.error(f"Failed to log webhook delivery: {e}") - - return success, status_code, duration_ms, error_message - - async def _enrich_payload(self, event: WebhookEvent) -> Dict[str, Any]: - """ - Enrich webhook payload with job result data at delivery time. - - For job.completed events: - - Adds result_url (fresh download URL for result zip) - - Adds result (inline payload with checksum/statistics) - - This ensures download URLs are generated fresh (they expire) - and data is current at delivery time. - """ - payload = dict(event.payload) # Copy to avoid mutating stored payload - - # Only enrich completion events - if payload.get("event") != "job.completed": - return payload - - try: - # Fetch job with result - from sqlalchemy.orm import selectinload - - from shared.models.database.job import Job - - async with get_db_context() as db: - result = await db.execute( - select(Job) - .options(selectinload(Job.job_result)) - .where(Job.job_id == event.job_id) - ) - job = result.scalar_one_or_none() - - if not job or not job.job_result: - logger.warning( - f"Job or result not found for enrichment: job_id={event.job_id}" - ) - return payload - - job_result = job.job_result - - # Add result_url (fresh download link) - if job_result.result_s3_key: - from shared.services.storage.file_upload_service import ( - FileUploadService, - ) - - upload_service = FileUploadService() - url_info = await upload_service.generate_download_url( - job_result.result_s3_key - ) - payload["result_url"] = url_info["download_url"] - logger.debug( - f"Enriched payload with result_url for job {event.job_id}" - ) - - # Add result (inline payload) - if job_result.inline_payload: - payload["result"] = job_result.inline_payload - - except Exception as e: - logger.error(f"Failed to enrich payload for event {event.id}: {e}") - # Continue with original payload if enrichment fails - - return payload + async def send_manual_webhook( + self, db: AsyncSession, event: WebhookEvent + ) -> WebhookDeliveryResult: + """Send a webhook immediately for an operator-triggered retry.""" + return await self._event_delivery.send(db=db, event=event, is_manual=True) def _is_retryable_error(self, status_code: Optional[int]) -> bool: """ @@ -402,62 +170,6 @@ def _is_retryable_error(self, status_code: Optional[int]) -> bool: # Examples: 400 Bad Request, 401 Unauthorized, 404 Not Found return False - async def _resolve_secret( - self, db: AsyncSession, user_id: str, endpoint: str - ) -> Optional[str]: - """ - Resolve webhook secret using repository (Lazy creation). - - 1. Try to get existing active secret for user/endpoint. - 2. If not found, create a new one. - 3. Decrypt and return the raw secret string. - """ - try: - # Import here to avoid circular dependency with WebhookDispatcher - from shared.repositories.webhook_secret_repository import ( - WebhookSecretRepository, - ) - - repo = WebhookSecretRepository() - secret_obj = await repo.get_or_create_secret(db, user_id, endpoint=endpoint) - - # Update usage timestamp - if secret_obj: - secret_obj.last_used_at = datetime.now(timezone.utc).replace( - tzinfo=None - ) - db.add(secret_obj) - # We don't commit here to avoid side effects if the caller aborts, - # but the session will eventually be committed by the caller. - - # Decrypt - return repo.decrypt_secret(secret_obj) - except (SystemSettingMissingException, SystemSettingInvalidException): - # Re-raise configuration errors so they can be handled as non-retryable - raise - except Exception as e: - logger.error(f"Failed to resolve/create secret for user {user_id}: {e}") - return None - - def _sign_payload(self, payload: Dict[str, Any], secret: str) -> str: - """ - Generate timestamped HMAC-SHA256 signature. - - Format: t=,v1= - Signed content: "{timestamp}.{json_payload}" - - This prevents replay attacks by binding the signature to the current time. - """ - timestamp = int(time.time()) - payload_str = json.dumps(payload, separators=(",", ":")) - signed_content = f"{timestamp}.{payload_str}" - - signature = hmac.new( - secret.encode("utf-8"), signed_content.encode("utf-8"), hashlib.sha256 - ).hexdigest() - - return f"t={timestamp},v1={signature}" - async def _mark_delivered(self, db: AsyncSession, event: WebhookEvent) -> None: """Mark event as delivered.""" event.status = WebhookEventStatus.DELIVERED diff --git a/packages/shared-python/shared/services/webhook/event_delivery.py b/packages/shared-python/shared/services/webhook/event_delivery.py new file mode 100644 index 000000000..21b5509fa --- /dev/null +++ b/packages/shared-python/shared/services/webhook/event_delivery.py @@ -0,0 +1,127 @@ +"""Direct WebhookEvent delivery attempt orchestration.""" + +import uuid +from typing import Any + +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + SystemSettingInvalidException, + SystemSettingMissingException, +) +from shared.models.database.webhook import WebhookEvent +from shared.models.database.webhook_log import WebhookLog +from shared.services.webhook.delivery_client import ( + WebhookDeliveryClient, + WebhookDeliveryResult, +) +from shared.services.webhook.payload_enrichment import WebhookPayloadEnricher +from shared.services.webhook.secret_resolver import WebhookSecretResolver +from shared.services.webhook.signing import build_webhook_headers + + +class WebhookEventDelivery: + """Send one direct webhook attempt and persist its delivery log.""" + + def __init__( + self, + *, + client: WebhookDeliveryClient | None = None, + enricher: WebhookPayloadEnricher | None = None, + secret_resolver: WebhookSecretResolver | None = None, + ) -> None: + self._client = client or WebhookDeliveryClient() + self._enricher = enricher or WebhookPayloadEnricher() + self._secret_resolver = secret_resolver or WebhookSecretResolver() + + async def send( + self, *, db: AsyncSession, event: WebhookEvent, is_manual: bool = False + ) -> WebhookDeliveryResult: + attempt_id = str(uuid.uuid4()) + target_validation = await self._client.validate_target( + event_id=event.id, + target_url=event.target_url, + ) + if target_validation.failure: + return target_validation.failure + if not target_validation.target: + return WebhookDeliveryResult( + success=False, + status_code=400, + duration_ms=0, + error_message="Webhook target validation failed", + ) + + payload = await self._enricher.enrich(event) + if is_manual: + payload["trigger"] = "manual" + + secret, secret_error = await self._resolve_secret(db, event) + if not secret: + logger.error(f"No secret found or created/resolved for event {event.id}") + return WebhookDeliveryResult( + success=False, + status_code=424, + duration_ms=0, + error_message=secret_error or "Secret resolution failed", + ) + + headers = build_webhook_headers( + payload=payload, + secret=secret, + attempt_id=attempt_id, + ) + result = await self._client.post_json( + event_id=event.id, + target=target_validation.target, + payload=payload, + headers=headers, + ) + await self._log_attempt( + db=db, + event=event, + is_manual=is_manual, + headers=headers, + payload=payload, + result=result, + ) + return result + + async def _resolve_secret( + self, db: AsyncSession, event: WebhookEvent + ) -> tuple[str | None, str | None]: + try: + return await self._secret_resolver.resolve_for_event(db, event), None + except (SystemSettingMissingException, SystemSettingInvalidException) as error: + logger.error(f"Configuration error during secret resolution: {error}") + return None, f"Configuration Error: {error}" + + async def _log_attempt( + self, + *, + db: AsyncSession, + event: WebhookEvent, + is_manual: bool, + headers: dict[str, str], + payload: dict[str, Any], + result: WebhookDeliveryResult, + ) -> None: + try: + log = WebhookLog( + job_id=event.job_id, + event_id=None if is_manual else event.id, + webhook_url=event.target_url, + attempt_number=event.attempts + 1, + request_payload={"header": headers, "payload": payload}, + signature=headers["X-Knowhere-Signature"], + idempotency_key=str(uuid.uuid4()), + response_status_code=result.status_code, + error_message=result.error_message, + duration_ms=result.duration_ms, + ) + db.add(log) + await db.commit() + + except Exception as error: + logger.error(f"Failed to log webhook delivery: {error}") diff --git a/packages/shared-python/shared/services/webhook/payload_enrichment.py b/packages/shared-python/shared/services/webhook/payload_enrichment.py new file mode 100644 index 000000000..2e8d1b73d --- /dev/null +++ b/packages/shared-python/shared/services/webhook/payload_enrichment.py @@ -0,0 +1,50 @@ +"""Delivery-time webhook payload enrichment.""" + +from collections.abc import Mapping +from typing import Any, cast + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.orm import selectinload + +from shared.core.database import get_db_context +from shared.models.database.job import Job +from shared.models.database.webhook import WebhookEvent +from shared.services.jobs.result_delivery import JobResultDeliveryResolver + + +class WebhookPayloadEnricher: + """Add fresh Job Result delivery metadata to webhook payloads.""" + + def __init__(self, resolver: JobResultDeliveryResolver | None = None) -> None: + self._resolver = resolver or JobResultDeliveryResolver() + + async def enrich(self, event: WebhookEvent) -> dict[str, Any]: + payload = dict(cast(Mapping[str, Any], event.payload)) + + if payload.get("event") != "job.completed": + return payload + + try: + async with get_db_context() as db: + result = await db.execute( + select(Job) + .options(selectinload(Job.job_result)) + .where(Job.job_id == event.job_id) + ) + job = result.scalar_one_or_none() + + if not job or not job.job_result: + logger.warning( + f"Job or result not found for enrichment: job_id={event.job_id}" + ) + return payload + + return self._resolver.enrich_payload( + payload, + job_result=job.job_result, + ) + + except Exception as error: + logger.error(f"Failed to enrich payload for event {event.id}: {error}") + return payload diff --git a/packages/shared-python/shared/services/webhook/qstash_client.py b/packages/shared-python/shared/services/webhook/qstash_client.py new file mode 100644 index 000000000..78547d04a --- /dev/null +++ b/packages/shared-python/shared/services/webhook/qstash_client.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any, Optional + +from loguru import logger + +from shared.core.config import app_config +from shared.core.exceptions.domain_exceptions import QStashServiceException +from shared.models.database.webhook import WebhookEventStatus + + +@dataclass(frozen=True) +class QStashDeliveryStatus: + """Terminal delivery status observed from QStash logs.""" + + status: str + response_status_code: Optional[int] + response_body: Optional[str] + error_message: Optional[str] + + +class QStashClientAdapter: + """Upstash QStash client adapter for webhook publication and log lookup.""" + + def __init__(self) -> None: + self._client: Any = None + + def get_client(self) -> Any: + """Lazily initialize the QStash client.""" + if self._client is None: + try: + from qstash import QStash + except ImportError as exc: + raise QStashServiceException( + internal_message=( + "qstash package is required for QStash webhook delivery. " + "Install it with: pip install qstash" + ), + operation="initialize_client", + original_exception=exc, + ) from exc + + token = app_config.QSTASH_TOKEN + if not token: + raise QStashServiceException( + internal_message="QSTASH_TOKEN is not configured", + operation="initialize_client", + ) + + self._client = QStash(token, base_url=app_config.QSTASH_BASE_URL) + return self._client + + def publish_webhook( + self, + *, + target_url: str, + payload: dict[str, Any], + signature: str, + event_id: str, + ) -> Optional[str]: + """Call the QStash publish API.""" + headers = { + "Content-Type": "application/json", + "X-Knowhere-Signature": signature, + "X-Knowhere-Event-ID": event_id, + "User-Agent": "Knowhere-Webhook/1.0", + } + + callback_url = app_config.qstash_callback_url + failure_callback_url = app_config.qstash_failure_callback_url + if not callback_url or not failure_callback_url: + raise QStashServiceException( + internal_message=( + "QSTASH_CALLBACK_BASE_URL must be configured for QStash " + "webhook delivery" + ), + operation="publish_webhook", + ) + + publish_kwargs: dict[str, Any] = { + "url": target_url, + "body": json.dumps(payload, separators=(",", ":")), + "headers": headers, + "retries": app_config.QSTASH_MAX_RETRIES, + "content_type": "application/json", + "retry_delay": _get_retry_delay_expression(), + "callback": callback_url, + "failure_callback": failure_callback_url, + "deduplication_id": event_id, + "label": "knowhere-webhook", + } + + response = self.get_client().message.publish(**publish_kwargs) + + message_id = getattr(response, "message_id", None) + if message_id is None and isinstance(response, dict): + message_id = response.get("messageId") or response.get("message_id") + + return message_id + + def get_terminal_delivery_status( + self, + qstash_message_id: str, + ) -> Optional[QStashDeliveryStatus]: + """Read QStash logs for a terminal destination delivery state.""" + try: + from qstash.log import LogState + + response = self.get_client().log.list( + filter={"message_id": qstash_message_id}, + count=20, + ) + except Exception as exc: + logger.warning( + f"QStash delivery status lookup failed: " + f"message_id={qstash_message_id}, error={exc}" + ) + return None + + terminal_logs = sorted(response.logs, key=lambda log: log.time, reverse=True) + for log in terminal_logs: + if log.state == LogState.DELIVERED: + return QStashDeliveryStatus( + status=WebhookEventStatus.DELIVERED, + response_status_code=log.response_status, + response_body=log.response_body, + error_message=log.error, + ) + + if log.state == LogState.FAILED: + return QStashDeliveryStatus( + status=WebhookEventStatus.FAILED, + response_status_code=log.response_status, + response_body=log.response_body, + error_message=log.error, + ) + + return None + + +def _get_retry_delay_expression() -> str: + # Approximate exponential backoff: 1m, 10m, ~100m, ~100m, ~100m. + return "pow(10, min(retried, 2)) * 60000" diff --git a/packages/shared-python/shared/services/webhook/qstash_payload.py b/packages/shared-python/shared/services/webhook/qstash_payload.py new file mode 100644 index 000000000..e454c076b --- /dev/null +++ b/packages/shared-python/shared/services/webhook/qstash_payload.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +from typing import Any + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.orm import selectinload + +from shared.models.database.job import Job +from shared.services.jobs.result_delivery import JobResultDeliveryResolver + + +class QStashPayloadEnricher: + """Sync payload enricher for QStash webhook publication.""" + + def __init__(self, resolver: JobResultDeliveryResolver | None = None) -> None: + self._resolver = resolver or JobResultDeliveryResolver() + + def enrich(self, db: Any, event: Any) -> dict[str, Any]: + payload = dict(event.payload) + if payload.get("event") != "job.completed": + return payload + + try: + result = db.execute( + select(Job) + .options(selectinload(Job.job_result)) + .where(Job.job_id == event.job_id) + ) + job = result.scalar_one_or_none() + if not job or not job.job_result: + return payload + + return self._resolver.enrich_payload( + payload, + job_result=job.job_result, + ) + except Exception as exc: + logger.error(f"Failed to enrich payload for event {event.id}: {exc}") + return payload diff --git a/packages/shared-python/shared/services/webhook/qstash_publisher.py b/packages/shared-python/shared/services/webhook/qstash_publisher.py index 8150100b2..810f35def 100644 --- a/packages/shared-python/shared/services/webhook/qstash_publisher.py +++ b/packages/shared-python/shared/services/webhook/qstash_publisher.py @@ -11,63 +11,39 @@ from __future__ import annotations -import hashlib -import hmac -import json -import time -from dataclasses import dataclass -from typing import Any, Dict, Optional +from typing import Optional from loguru import logger - -from shared.core.config import app_config -from shared.core.exceptions.domain_exceptions import QStashServiceException -from shared.models.database.webhook import WebhookEventStatus -from shared.utils.url_security import ( +from sqlalchemy import select + +from shared.core.database_sync import get_sync_db_context +from shared.models.database.job import Job +from shared.models.database.webhook import WebhookEvent, WebhookEventStatus +from shared.services.webhook.qstash_client import ( + QStashClientAdapter, + QStashDeliveryStatus, +) +from shared.services.webhook.qstash_payload import QStashPayloadEnricher +from shared.services.webhook.qstash_secret_resolver import QStashSecretResolver +from shared.services.webhook.signing import sign_webhook_payload +from shared.services.http.url_security import ( validate_http_url_and_resolve_ip, ) -@dataclass(frozen=True) -class QStashDeliveryStatus: - """Terminal delivery status observed from QStash logs.""" - - status: str - response_status_code: Optional[int] - response_body: Optional[str] - error_message: Optional[str] - - class QStashWebhookPublisher: """Publishes webhook events to customer endpoints via QStash.""" - def __init__(self) -> None: - self._client: Any = None - - def _get_client(self) -> Any: - """Lazily initialize the QStash client.""" - if self._client is None: - try: - from qstash import QStash - except ImportError as exc: - raise QStashServiceException( - internal_message=( - "qstash package is required for QStash webhook delivery. " - "Install it with: pip install qstash" - ), - operation="initialize_client", - original_exception=exc, - ) from exc - - token = app_config.QSTASH_TOKEN - if not token: - raise QStashServiceException( - internal_message="QSTASH_TOKEN is not configured", - operation="initialize_client", - ) - - self._client = QStash(token, base_url=app_config.QSTASH_BASE_URL) - return self._client + def __init__( + self, + *, + client_adapter: QStashClientAdapter | None = None, + payload_enricher: QStashPayloadEnricher | None = None, + secret_resolver: QStashSecretResolver | None = None, + ) -> None: + self._client_adapter = client_adapter or QStashClientAdapter() + self._payload_enricher = payload_enricher or QStashPayloadEnricher() + self._secret_resolver = secret_resolver or QStashSecretResolver() def publish_event(self, event_id: str) -> Optional[str]: """Publish a webhook event via QStash. @@ -77,12 +53,6 @@ def publish_event(self, event_id: str) -> Optional[str]: Returns the QStash message_id on success, or None on failure. """ - from sqlalchemy import select - - from shared.core.database_sync import get_sync_db_context - from shared.models.database.job import Job - from shared.models.database.webhook import WebhookEvent - with get_sync_db_context() as db: event = db.execute( select(WebhookEvent).where(WebhookEvent.id == event_id) @@ -109,10 +79,8 @@ def publish_event(self, event_id: str) -> Optional[str]: db.commit() return None - # Enrich payload (presigned S3 URL for completed jobs) - payload = self._enrich_payload(db, event) + payload = self._payload_enricher.enrich(db, event) - # Resolve signing secret user_id = db.execute( select(Job.user_id).where(Job.job_id == event.job_id) ).scalar_one_or_none() @@ -123,7 +91,11 @@ def publish_event(self, event_id: str) -> Optional[str]: db.commit() return None - secret = self._resolve_secret(db, str(user_id), event.target_url) + secret = self._secret_resolver.resolve( + db, + user_id=str(user_id), + endpoint=event.target_url, + ) if not secret: logger.error( f"QStash publish: secret resolution failed for event {event_id}" @@ -132,12 +104,10 @@ def publish_event(self, event_id: str) -> Optional[str]: db.commit() return None - # Sign payload with our HMAC - signature = self._sign_payload(payload, secret) + signature = sign_webhook_payload(payload, secret) - # Publish to QStash try: - message_id = self._publish_to_qstash( + message_id = self._client_adapter.publish_webhook( target_url=event.target_url, payload=payload, signature=signature, @@ -159,210 +129,12 @@ def publish_event(self, event_id: str) -> Optional[str]: ) return message_id - def _publish_to_qstash( - self, - target_url: str, - payload: Dict[str, Any], - signature: str, - event_id: str, - ) -> Optional[str]: - """Call the QStash publish API.""" - headers = { - "Content-Type": "application/json", - "X-Knowhere-Signature": signature, - "X-Knowhere-Event-ID": event_id, - "User-Agent": "Knowhere-Webhook/1.0", - } - - # Approximate exponential backoff: 1m, 10m, ~100m, ~100m, ~100m - # pow(10, min(retried, 2)) * 60000 → 60s, 600s, 6000s capped - retry_delay_expression = "pow(10, min(retried, 2)) * 60000" - - callback_url = app_config.qstash_callback_url - failure_callback_url = app_config.qstash_failure_callback_url - if not callback_url or not failure_callback_url: - raise QStashServiceException( - internal_message=( - "QSTASH_CALLBACK_BASE_URL must be configured for QStash " - "webhook delivery" - ), - operation="publish_webhook", - ) - - client = self._get_client() - - publish_kwargs: Dict[str, Any] = { - "url": target_url, - "body": json.dumps(payload, separators=(",", ":")), - "headers": headers, - "retries": app_config.QSTASH_MAX_RETRIES, - "content_type": "application/json", - "retry_delay": retry_delay_expression, - "callback": callback_url, - "failure_callback": failure_callback_url, - "deduplication_id": event_id, - "label": "knowhere-webhook", - } - - response = client.message.publish(**publish_kwargs) - - message_id = getattr(response, "message_id", None) - if message_id is None and isinstance(response, dict): - message_id = response.get("messageId") or response.get("message_id") - - return message_id - def get_terminal_delivery_status( self, qstash_message_id: str, ) -> Optional[QStashDeliveryStatus]: """Read QStash logs for a terminal destination delivery state.""" - try: - from qstash.log import LogState - - response = self._get_client().log.list( - filter={"message_id": qstash_message_id}, - count=20, - ) - except Exception as exc: - logger.warning( - f"QStash delivery status lookup failed: " - f"message_id={qstash_message_id}, error={exc}" - ) - return None - - terminal_logs = sorted(response.logs, key=lambda log: log.time, reverse=True) - for log in terminal_logs: - if log.state == LogState.DELIVERED: - return QStashDeliveryStatus( - status=WebhookEventStatus.DELIVERED, - response_status_code=log.response_status, - response_body=log.response_body, - error_message=log.error, - ) - - if log.state == LogState.FAILED: - return QStashDeliveryStatus( - status=WebhookEventStatus.FAILED, - response_status_code=log.response_status, - response_body=log.response_body, - error_message=log.error, - ) - - return None - - def _enrich_payload(self, db: Any, event: Any) -> Dict[str, Any]: - """Enrich the webhook payload (e.g., generate fresh presigned S3 URL).""" - from sqlalchemy import select - from sqlalchemy.orm import selectinload - - from shared.models.database.job import Job - - payload = dict(event.payload) - if payload.get("event") != "job.completed": - return payload - - try: - result = db.execute( - select(Job) - .options(selectinload(Job.job_result)) - .where(Job.job_id == event.job_id) - ) - job = result.scalar_one_or_none() - if not job or not job.job_result: - return payload - - job_result = job.job_result - if job_result.result_s3_key: - payload["result_url"] = app_config.get_storage_adapter().generate_presigned_url( - job_result.result_s3_key, - expiration=3600, - method="GET", - ) - - if job_result.inline_payload: - payload["result"] = job_result.inline_payload - except Exception as exc: - logger.error(f"Failed to enrich payload for event {event.id}: {exc}") - - return payload - - def _resolve_secret(self, db: Any, user_id: str, endpoint: str) -> Optional[str]: - """Resolve the webhook signing secret for a user/endpoint.""" - from datetime import datetime, timezone - - from sqlalchemy import and_, select - - from shared.core.exceptions.domain_exceptions import ( - SystemSettingInvalidException, - SystemSettingMissingException, - ) - from shared.models.database.webhook_secret import ( - WebhookSecret, - WebhookSecretStatus, - ) - from shared.services.encryption import get_fernet_service - - try: - fernet = get_fernet_service() - except (SystemSettingMissingException, SystemSettingInvalidException) as exc: - logger.error(f"Configuration error during secret resolution: {exc}") - return None - - # Try endpoint-specific secret first, then global - secret_obj = None - if endpoint: - result = db.execute( - select(WebhookSecret).where( - and_( - WebhookSecret.user_id == user_id, - WebhookSecret.endpoint == endpoint, - WebhookSecret.status == WebhookSecretStatus.ACTIVE, - ) - ) - ) - secret_obj = result.scalar_one_or_none() - - if secret_obj is None: - result = db.execute( - select(WebhookSecret).where( - and_( - WebhookSecret.user_id == user_id, - WebhookSecret.endpoint.is_(None), - WebhookSecret.status == WebhookSecretStatus.ACTIVE, - ) - ) - ) - secret_obj = result.scalar_one_or_none() - - if secret_obj is None: - raw_secret = fernet.generate_webhook_secret() - secret_obj = WebhookSecret( - user_id=user_id, - endpoint=endpoint, - secret_encrypted=fernet.encrypt(raw_secret), - status=WebhookSecretStatus.ACTIVE, - ) - db.add(secret_obj) - db.commit() - db.refresh(secret_obj) - - secret_obj.last_used_at = datetime.now(timezone.utc).replace(tzinfo=None) - db.add(secret_obj) - return fernet.decrypt(secret_obj.secret_encrypted) - - @staticmethod - def _sign_payload(payload: Dict[str, Any], secret: str) -> str: - """Generate HMAC-SHA256 signature matching the existing Knowhere format.""" - timestamp = int(time.time()) - payload_str = json.dumps(payload, separators=(",", ":")) - signed_content = f"{timestamp}.{payload_str}" - sig = hmac.new( - secret.encode("utf-8"), - signed_content.encode("utf-8"), - hashlib.sha256, - ).hexdigest() - return f"t={timestamp},v1={sig}" + return self._client_adapter.get_terminal_delivery_status(qstash_message_id) # Module-level singleton diff --git a/packages/shared-python/shared/services/webhook/qstash_secret_resolver.py b/packages/shared-python/shared/services/webhook/qstash_secret_resolver.py new file mode 100644 index 000000000..fb8174883 --- /dev/null +++ b/packages/shared-python/shared/services/webhook/qstash_secret_resolver.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any, Optional + +from loguru import logger +from sqlalchemy import and_, select + +from shared.core.exceptions.domain_exceptions import ( + SystemSettingInvalidException, + SystemSettingMissingException, +) +from shared.models.database.webhook_secret import ( + WebhookSecret, + WebhookSecretStatus, +) +from shared.services.encryption import get_fernet_service + + +class QStashSecretResolver: + """Sync webhook secret resolver for QStash publication.""" + + def resolve(self, db: Any, *, user_id: str, endpoint: str) -> Optional[str]: + try: + fernet = get_fernet_service() + except (SystemSettingMissingException, SystemSettingInvalidException) as exc: + logger.error(f"Configuration error during secret resolution: {exc}") + return None + + secret = self._find_active_secret(db, user_id=user_id, endpoint=endpoint) + if secret is None: + raw_secret = fernet.generate_webhook_secret() + secret = WebhookSecret( + user_id=user_id, + endpoint=endpoint, + secret_encrypted=fernet.encrypt(raw_secret), + status=WebhookSecretStatus.ACTIVE, + ) + db.add(secret) + db.commit() + db.refresh(secret) + + secret.last_used_at = datetime.now(timezone.utc).replace(tzinfo=None) + db.add(secret) + return fernet.decrypt(secret.secret_encrypted) + + def _find_active_secret( + self, + db: Any, + *, + user_id: str, + endpoint: str, + ) -> WebhookSecret | None: + if endpoint: + result = db.execute( + select(WebhookSecret).where( + and_( + WebhookSecret.user_id == user_id, + WebhookSecret.endpoint == endpoint, + WebhookSecret.status == WebhookSecretStatus.ACTIVE, + ) + ) + ) + secret = result.scalar_one_or_none() + if secret is not None: + return secret + + result = db.execute( + select(WebhookSecret).where( + and_( + WebhookSecret.user_id == user_id, + WebhookSecret.endpoint.is_(None), + WebhookSecret.status == WebhookSecretStatus.ACTIVE, + ) + ) + ) + return result.scalar_one_or_none() diff --git a/packages/shared-python/shared/services/webhook/secret_resolver.py b/packages/shared-python/shared/services/webhook/secret_resolver.py new file mode 100644 index 000000000..d0a648ef3 --- /dev/null +++ b/packages/shared-python/shared/services/webhook/secret_resolver.py @@ -0,0 +1,60 @@ +"""Webhook secret resolution for direct deliveries.""" + +from datetime import datetime, timezone + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + SystemSettingInvalidException, + SystemSettingMissingException, +) +from shared.models.database.job import Job +from shared.models.database.webhook import WebhookEvent +from shared.repositories.webhook_secret_repository import WebhookSecretRepository + + +class WebhookSecretResolver: + """Resolve the active endpoint secret for a WebhookEvent delivery.""" + + def __init__(self, repository: WebhookSecretRepository | None = None) -> None: + self._repository = repository or WebhookSecretRepository() + + async def resolve_for_event( + self, db: AsyncSession, event: WebhookEvent + ) -> str | None: + user_id = await self._get_job_owner(db, event.job_id) + if not user_id: + logger.warning(f"Could not resolve secret: Job {event.job_id} has no user_id") + return None + + return await self.resolve_for_endpoint( + db, + user_id=user_id, + endpoint=event.target_url, + ) + + async def resolve_for_endpoint( + self, db: AsyncSession, *, user_id: str, endpoint: str + ) -> str | None: + try: + secret = await self._repository.get_or_create_secret( + db, user_id, endpoint=endpoint + ) + + if secret: + secret.last_used_at = datetime.now(timezone.utc).replace(tzinfo=None) + db.add(secret) + + return self._repository.decrypt_secret(secret) + + except (SystemSettingMissingException, SystemSettingInvalidException): + raise + except Exception as error: + logger.error(f"Failed to resolve/create secret for user {user_id}: {error}") + return None + + async def _get_job_owner(self, db: AsyncSession, job_id: str) -> str | None: + result = await db.execute(select(Job.user_id).where(Job.job_id == job_id)) + return result.scalar_one_or_none() diff --git a/packages/shared-python/shared/services/webhook/signing.py b/packages/shared-python/shared/services/webhook/signing.py new file mode 100644 index 000000000..15c521946 --- /dev/null +++ b/packages/shared-python/shared/services/webhook/signing.py @@ -0,0 +1,34 @@ +"""Webhook request signing.""" + +import hashlib +import hmac +import json +import time +from collections.abc import Mapping +from typing import Any + + +def sign_webhook_payload(payload: Mapping[str, Any], secret: str) -> str: + """Generate the timestamped Knowhere webhook HMAC signature.""" + timestamp = int(time.time()) + payload_text = json.dumps(payload, separators=(",", ":")) + signed_content = f"{timestamp}.{payload_text}" + signature = hmac.new( + secret.encode("utf-8"), + signed_content.encode("utf-8"), + hashlib.sha256, + ).hexdigest() + + return f"t={timestamp},v1={signature}" + + +def build_webhook_headers( + *, payload: Mapping[str, Any], secret: str, attempt_id: str +) -> dict[str, str]: + """Build signed HTTP headers for a direct webhook delivery attempt.""" + return { + "Content-Type": "application/json", + "X-Knowhere-Signature": sign_webhook_payload(payload, secret), + "X-Knowhere-Attempt-ID": attempt_id, + "User-Agent": "Knowhere-Webhook/1.0", + } diff --git a/packages/shared-python/shared/testing/contract_runtime.py b/packages/shared-python/shared/testing/contract_runtime.py index 11aa49efb..ce029a8c6 100644 --- a/packages/shared-python/shared/testing/contract_runtime.py +++ b/packages/shared-python/shared/testing/contract_runtime.py @@ -7,6 +7,7 @@ import subprocess import sys from collections.abc import Awaitable, Callable +from dataclasses import dataclass from pathlib import Path from types import ModuleType from typing import Any, Protocol, cast @@ -65,7 +66,7 @@ "shared.core.state_machine.service_sync", "shared.services.billing.credits_sync_service", "shared.services.billing.work_billing_service", - "shared.services.job_lifecycle_sync", + "shared.services.jobs.lifecycle.service", "shared.services.retrieval.app_service", "shared.services.retrieval.publication_service", "shared.services.webhook", @@ -77,23 +78,36 @@ CONTRACT_DEVELOPER_USER_EMAIL: str = "local-dev-user@knowhere.local" CONTRACT_DEVELOPER_USER_TIER: str = "tier_5" CONTRACT_DEVELOPER_API_KEY_NAME: str = "contract-developer-api-key" -_contract_storage_prepared: bool = False _contract_storage_database_url: str | None = None _contract_fake_redis_server: fakeredis.FakeServer = fakeredis.FakeServer() +@dataclass +class ContractStorageRuntimeState: + is_prepared: bool = False + + +_contract_storage_runtime_state: ContractStorageRuntimeState = ( + ContractStorageRuntimeState() +) + + class PostgreSQLProcess(Protocol): @property - def host(self) -> str: ... + def host(self) -> str: + raise NotImplementedError @property - def port(self) -> int: ... + def port(self) -> int: + raise NotImplementedError @property - def user(self) -> str: ... + def user(self) -> str: + raise NotImplementedError @property - def password(self) -> str | None: ... + def password(self) -> str | None: + raise NotImplementedError def _ensure_import_paths() -> None: @@ -113,13 +127,12 @@ def _ensure_test_directories() -> None: def _reset_contract_storage_state(database_url: str) -> None: global _contract_storage_database_url - global _contract_storage_prepared if _contract_storage_database_url == database_url: return _contract_storage_database_url = database_url - _contract_storage_prepared = False + _contract_storage_runtime_state.is_prepared = False def _ensure_contract_postgresql_port(database_url: URL) -> None: @@ -358,6 +371,24 @@ def clear_application_modules() -> None: sys.modules.pop(module_name, None) continue + if module_name == "shared.services.storage" or module_name.startswith( + "shared.services.storage." + ): + sys.modules.pop(module_name, None) + continue + + if module_name == "shared.services.jobs" or module_name.startswith( + "shared.services.jobs." + ): + sys.modules.pop(module_name, None) + continue + + if module_name == "shared.services.webhook" or module_name.startswith( + "shared.services.webhook." + ): + sys.modules.pop(module_name, None) + continue + if module_name == "app" or module_name.startswith("app."): sys.modules.pop(module_name, None) @@ -578,7 +609,6 @@ def drop_contract_database( postgresql_process: PostgreSQLProcess | None = None, ) -> None: global _contract_storage_database_url - global _contract_storage_prepared contract_database_url = get_contract_database_url(postgresql_process) contract_database_name = make_url(contract_database_url).database @@ -593,7 +623,7 @@ def drop_contract_database( if _contract_storage_database_url == contract_database_url: _contract_storage_database_url = None - _contract_storage_prepared = False + _contract_storage_runtime_state.is_prepared = False def _run_contract_migrations() -> None: @@ -661,16 +691,14 @@ async def _create_contract_engine() -> AsyncEngine: async def prepare_contract_storage() -> None: - global _contract_storage_prepared - _ensure_import_paths() - if not _contract_storage_prepared: + if not _contract_storage_runtime_state.is_prepared: _recreate_contract_database() _initialize_contract_database() _run_contract_migrations() _assert_contract_schema_ready() - _contract_storage_prepared = True + _contract_storage_runtime_state.is_prepared = True await reset_contract_database() await reset_contract_redis() diff --git a/packages/shared-python/shared/utils/CommonHelper.py b/packages/shared-python/shared/utils/CommonHelper.py deleted file mode 100644 index 2790c7a6a..000000000 --- a/packages/shared-python/shared/utils/CommonHelper.py +++ /dev/null @@ -1,45 +0,0 @@ -from io import BytesIO -from pathlib import Path - -import httpx -import pandas as pd -from starlette.datastructures import UploadFile as StarletteUploadFile - -from shared.utils.FileDownUpUtils import s3_upload_file - - -def is_remote(path): - """Check whether a path is a remote URL.""" - if path is None: - return False - if not isinstance(path, str): - return False - return path.startswith("http://") or path.startswith("https://") - - -async def load_file_bytes(file_path, *, file_url="", timeout=None): - if isinstance(file_path, str) and is_remote(file_path): - # If file_path is already a full URL, use it directly. - url_to_use = file_path - if not isinstance(file_url, str): - file_url = file_url.geturl() - # Prefer file_url when provided; otherwise keep file_path. - if file_url: - url_to_use = file_url - async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client: - r = await client.get(url_to_use) # Fetch the resolved URL. - r.raise_for_status() - return r.content - else: - p = Path(file_path) - return p.read_bytes() - - -async def upload_dataframe_to_s3(df: pd.DataFrame, filename: str, prefix: str): - # Write the DataFrame into an in-memory BytesIO buffer. - buffer = BytesIO() - df.to_csv(buffer, index=False) - buffer.seek(0) # Reset the cursor to the buffer start. - - upload_file = StarletteUploadFile(file=buffer, filename=filename) - s3_upload_file(upload_file, prefix) diff --git a/packages/shared-python/shared/utils/FileDownUpUtils.py b/packages/shared-python/shared/utils/FileDownUpUtils.py deleted file mode 100644 index 0f124ebe8..000000000 --- a/packages/shared-python/shared/utils/FileDownUpUtils.py +++ /dev/null @@ -1,253 +0,0 @@ -import os -import uuid -import zipfile -from pathlib import Path -from typing import Optional, Union -from urllib.parse import urljoin - -import aiohttp -import requests -from botocore.exceptions import ClientError -from starlette.datastructures import UploadFile - -from shared.core.config import settings -from shared.core.config.storage import get_cached_storage_adapter -from shared.core.exceptions.domain_exceptions import ( - KnowhereException, - NotFoundException, - StorageServiceException, -) -from shared.models.schemas.s3_file import FliesDownload - - -def s3_upload_file(file: UploadFile, prefix: str): - """ - Upload a file object to S3 storage. - :param file: Input file such as ``abc15sa25ww.doc`` - :param prefix: Storage prefix such as ``upload/123`` - :return: Upload result payload - """ - if prefix and not prefix.endswith("/"): - prefix += "/" - object_key = f"{prefix}{file.filename}" - adapter = get_cached_storage_adapter() - try: - # ``upload_fileobj`` streams efficiently and avoids large in-memory copies. - adapter.upload_fileobj( - file.file, object_key, content_type="application/octet-stream" - ) - public_url = ( - f"{settings.S3_PRIVATE_DOMAIN}/{object_key}" - if settings.S3_PRIVATE_DOMAIN - else f"storage/{object_key}" - ) - content = { - "message": "File uploaded successfully", - "bucket": settings.S3_BUCKET_NAME, - "file_key": object_key, - "public_url_for_reference": public_url, - } - return content - - except KnowhereException: - raise - except Exception as e: - # Wrap storage upload failures in a domain exception. - raise StorageServiceException( - internal_message=f"Storage upload failed: {str(e)}", - operation="upload", - original_exception=e, - ) - - -def s3_download_extract_zip( - url: str, - dest_dir: Union[str, os.PathLike], - *, - filename: str = "parsed.zip", - headers: Optional[dict] = None, - timeout: int | None = None, - chunk_size: int | None = None, - keep_exts: tuple[str, ...] = (".md", ".json"), - exclude_patterns: tuple[str, ...] = (), - clean_empty_dirs: bool = True, -): - """ - Download and extract a zip file, keeping only specific file types. - - Args: - exclude_patterns: Tuple of filename patterns to exclude (e.g., ("content_list", "middle.json")) - """ - import fnmatch - - from shared.core.constants import APIConstants, ProcessingConstants - - # Use defaults when optional arguments are omitted. - if timeout is None: - timeout = APIConstants.S3_FILE_DOWNLOAD_TIMEOUT - if chunk_size is None: - chunk_size = ProcessingConstants.IMG_CHUNK_SIZE - - dest_dir = Path(dest_dir).expanduser().resolve() - dest_dir.mkdir(parents=True, exist_ok=True) - zip_path = dest_dir / filename - - # 1) Download to zip_path and extract. - with requests.get( - url, headers=headers or {}, timeout=timeout, stream=True, allow_redirects=True - ) as r: - r.raise_for_status() - with open(zip_path, "wb") as f: - for chunk in r.iter_content(chunk_size=chunk_size): - if chunk: - f.write(chunk) - - with zipfile.ZipFile(zip_path, "r") as zf: - zf.extractall(dest_dir) - - # 2) Remove files outside keep_exts or matching exclude_patterns. - kept_files = [] - for p in dest_dir.rglob("*"): - if p.is_file(): - # Check if file should be excluded by pattern - should_exclude = False - for pattern in exclude_patterns: - if pattern in p.name or fnmatch.fnmatch(p.name, pattern): - should_exclude = True - break - - if should_exclude: - p.unlink() - elif p.suffix.lower() in keep_exts: - kept_files.append(p) - else: - p.unlink() - - # 4) Remove empty directories when requested. - if clean_empty_dirs: - for d in sorted( - [p for p in dest_dir.rglob("*") if p.is_dir()], - key=lambda x: len(x.parts), - reverse=True, - ): - try: - next(d.iterdir()) - except StopIteration: - d.rmdir() - # 5) Delete the downloaded zip file. - zip_path.unlink(missing_ok=True) - - -def s3_get_download_url(file_key: str, expires_in: int = 3600): - """ - Get a file download URL from its storage key. - :param file_key: Full file path and name - :param expires_in: Desired URL lifetime - :return: Signed download payload - """ - s3_client = settings.get_s3_client() - try: - # Generate a pre-signed URL. - presigned_url = s3_client.generate_presigned_url( - "get_object", - Params={"Bucket": settings.S3_BUCKET_NAME, "Key": file_key}, - ExpiresIn=expires_in, # URL lifetime. - ) - fsdl = FliesDownload( - message="URL signed successfully", - file_key=file_key, - download_url=presigned_url, - expires_in_seconds=expires_in, - ) - return fsdl - - except ClientError as e: - # boto3 may still sign missing objects; the resulting URL can later 404. - raise NotFoundException( - resource="File", - resource_id=file_key, - internal_message=( - f"Could not generate the URL. Check whether the file is correct " - f"or the S3 configuration is valid: {str(e)}" - ), - ) - - -def get_url_file(path): - file_sig = s3_get_download_url(path, expires_in=3600) - # Assemble the final URL. - file_url = file_sig.download_url - response = requests.get(file_url, verify=True) - response.raise_for_status() # Ensure the request succeeds. - return response - - -def get_pub_fileurl(path): - """ - Build a public URL from a storage path. - :param path: - :return: Public URL - """ - base_url = settings.S3_PRIVATE_DOMAIN.rstrip("/") - clean_path = path.replace("\\", "/").strip() - full_url = urljoin(base_url + "/", clean_path) - return full_url - - -def s3_public_file_url(file_key: str) -> str: - permanent_url = f"{settings.S3_PRIVATE_DOMAIN}/{settings.S3_BUCKET_NAME}/{file_key}" - return permanent_url - - -async def download_and_upload_image( - img_url: str, prefix: str = "images/", temp_store_path=None -) -> dict: - """ - Download an image, rename it, upload it to S3, and clean up locally. - :param img_url: Image URL - :param prefix: S3 storage prefix - :return: Dict containing upload results and the new download reference - """ - # Generate a unique filename. - unique_filename = f"{uuid.uuid4()}.jpg" - # Temporary directory. - if temp_store_path is None: - temp_store_path = r"/Volumes/U/temp/output/" - local_file_path = Path(f"{settings.S3_TEMP_PATH or '/tmp'}{unique_filename}") - # Path(f"{temp_store_path}{unique_filename}") - try: - # Download the image asynchronously. - async with aiohttp.ClientSession() as session: - async with session.get(img_url) as response: - response.raise_for_status() - with open(local_file_path, "wb") as f: - f.write(await response.read()) - - # Create a temporary UploadFile wrapper. - from fastapi import UploadFile - - upload_file = UploadFile( - filename=unique_filename, file=open(local_file_path, "rb") - ) - - # Upload to S3. - result = s3_upload_file(upload_file, prefix) - - # Close the file handle and delete the local file. - upload_file.file.close() - os.remove(local_file_path) - return result - - except KnowhereException: - if local_file_path.exists(): - os.remove(local_file_path) - raise - except Exception as e: - # Always remove the local file on failure as well. - if local_file_path.exists(): - os.remove(local_file_path) - raise StorageServiceException( - internal_message=f"Failed to download and upload the image: {str(e)}", - operation="download_and_upload", - original_exception=e, - ) diff --git a/packages/shared-python/shared/utils/file_transfer.py b/packages/shared-python/shared/utils/file_transfer.py deleted file mode 100644 index d307e7995..000000000 --- a/packages/shared-python/shared/utils/file_transfer.py +++ /dev/null @@ -1,171 +0,0 @@ -""" -File Transfer Utilities - -Provides reliable file transfer operations for large files using temp files as buffers. -Uses httpx for proper total timeout enforcement. -""" - -import os -import tempfile -from typing import Dict, Optional -from urllib.parse import urlparse - -import httpx -from loguru import logger - -from shared.utils.http_clients import get_sync_client - - -class FileTransferError(Exception): - """Base exception for file transfer operations""" - - def __init__(self, message: str, status_code: Optional[int] = None): - super().__init__(message) - self.status_code = status_code - - -class DownloadError(FileTransferError): - """ - Download failed - typically a client error. - - The source file may be inaccessible, expired, or invalid. - Worker should raise a 4xx (client error) when catching this. - """ - - pass - - -class UploadError(FileTransferError): - """ - Upload failed - typically a server/service error. - - The target service (e.g., MinerU) may be unavailable or experiencing issues. - Worker should raise a 5xx (server error) when catching this. - """ - - pass - - -def stream_download_and_upload( - source_url: str, - target_url: str, - download_timeout: int = 300, - upload_timeout: int = 300, - chunk_size: int = 8192, - upload_method: str = "PUT", - upload_headers: Optional[Dict[str, str]] = None, - upload_retries: int = 3, -) -> httpx.Response: - """ - Download a file from source_url and upload to target_url using a temp file buffer. - - Uses httpx for proper total timeout enforcement. - Retries upload on failure since temp file is preserved on disk. - - Args: - source_url: URL to download the file from - target_url: URL to upload the file to - download_timeout: Total timeout for download in seconds - upload_timeout: Total timeout for upload in seconds - chunk_size: Chunk size for streaming download - upload_method: HTTP method for upload (PUT or POST) - upload_headers: Additional headers for upload request - upload_retries: Number of retry attempts for upload (default 3) - - Returns: - httpx.Response: The upload response - - Raises: - DownloadError: If download fails (source inaccessible) - UploadError: If upload fails after all retries - """ - # Create temp file manually for explicit cleanup control - tmp_fd, tmp_path = tempfile.mkstemp(suffix=".tmp") - source_host = urlparse(source_url).hostname or source_url[:60] - target_host = urlparse(target_url).hostname or target_url[:60] - - try: - # Phase 1: Download to temp file - logger.debug(f"Downloading from {source_url[:100]}...") - try: - client = get_sync_client() - with client.stream("GET", source_url, timeout=download_timeout) as response: - response.raise_for_status() - - downloaded_bytes = 0 - with os.fdopen(tmp_fd, "wb") as tmp_file: - for chunk in response.iter_bytes(chunk_size=chunk_size): - tmp_file.write(chunk) - downloaded_bytes += len(chunk) - - except httpx.TimeoutException as e: - raise DownloadError( - f"Download timed out: host={source_host}, timeout={download_timeout}s" - ) from e - except httpx.HTTPStatusError as e: - raise DownloadError( - f"Download failed: host={source_host}, status={e.response.status_code}" - ) from e - except httpx.RequestError as e: - raise DownloadError( - f"Download failed: host={source_host}, error={e}" - ) from e - - # Get file size - file_size = os.path.getsize(tmp_path) - logger.info(f"Downloaded {file_size} bytes to temp file") - - # Phase 2: Upload from temp file (with retries) - headers = upload_headers or {} - headers["Content-Length"] = str(file_size) - - last_error = None - for attempt in range(1, upload_retries + 1): - try: - logger.info( - f"Uploading {file_size} bytes (attempt {attempt}/{upload_retries}, timeout={upload_timeout}s)..." - ) - - # Stream directly from file without loading to memory - with open(tmp_path, "rb") as f: - client = get_sync_client() - if upload_method.upper() == "PUT": - upload_response = client.put( - target_url, - content=f, - headers=headers, - timeout=upload_timeout, - ) - else: - upload_response = client.post( - target_url, - content=f, - headers=headers, - timeout=upload_timeout, - ) - - logger.info(f"Upload completed: status={upload_response.status_code}") - return upload_response - - except (httpx.TimeoutException, httpx.RequestError) as e: - last_error = e - logger.warning( - f"Upload attempt {attempt} failed: host={target_host}, error={e}" - ) - if attempt < upload_retries: - logger.info("Retrying upload...") - continue - - # All retries exhausted - raise UploadError( - f"Upload failed: host={target_host}, attempts={upload_retries}, last_error={last_error}" - ) from last_error - - finally: - # Manual cleanup of temp file - if os.path.exists(tmp_path): - try: - os.remove(tmp_path) - logger.debug(f"Temp file cleaned up: {tmp_path}") - except OSError as e: - logger.warning(f"Failed to cleanup temp file {tmp_path}: {e}") diff --git a/packages/shared-python/shared/utils/zip_download.py b/packages/shared-python/shared/utils/zip_download.py new file mode 100644 index 000000000..630173c2f --- /dev/null +++ b/packages/shared-python/shared/utils/zip_download.py @@ -0,0 +1,79 @@ +"""Download-and-extract helpers for remote ZIP artifacts.""" + +import os +import zipfile +from pathlib import Path +from collections.abc import Mapping + +import requests + + +def download_and_extract_zip( + url: str, + dest_dir: str | os.PathLike[str], + *, + filename: str = "parsed.zip", + headers: Mapping[str, str] | None = None, + timeout: int | None = None, + chunk_size: int | None = None, + keep_exts: tuple[str, ...] = (".md", ".json"), + exclude_patterns: tuple[str, ...] = (), + clean_empty_dirs: bool = True, +) -> None: + """Download a ZIP file, extract it, and keep only the requested artifacts.""" + import fnmatch + + from shared.core.constants import APIConstants, ProcessingConstants + + if timeout is None: + timeout = APIConstants.S3_FILE_DOWNLOAD_TIMEOUT + if chunk_size is None: + chunk_size = ProcessingConstants.IMG_CHUNK_SIZE + + destination = Path(dest_dir).expanduser().resolve() + destination.mkdir(parents=True, exist_ok=True) + zip_path = destination / filename + + with requests.get( + url, + headers=headers or {}, + timeout=timeout, + stream=True, + allow_redirects=True, + ) as response: + response.raise_for_status() + with open(zip_path, "wb") as zip_file: + for chunk in response.iter_content(chunk_size=chunk_size): + if chunk: + zip_file.write(chunk) + + with zipfile.ZipFile(zip_path, "r") as extracted_zip: + extracted_zip.extractall(destination) + + for extracted_path in destination.rglob("*"): + if not extracted_path.is_file(): + continue + + should_exclude = False + for pattern in exclude_patterns: + if pattern in extracted_path.name or fnmatch.fnmatch(extracted_path.name, pattern): + should_exclude = True + break + + if should_exclude: + extracted_path.unlink() + elif extracted_path.suffix.lower() not in keep_exts: + extracted_path.unlink() + + if clean_empty_dirs: + for directory in sorted( + [path for path in destination.rglob("*") if path.is_dir()], + key=lambda path: len(path.parts), + reverse=True, + ): + try: + next(directory.iterdir()) + except StopIteration: + directory.rmdir() + + zip_path.unlink(missing_ok=True)