diff --git a/.gitignore b/.gitignore index 4dbbc7299..8da972e73 100644 --- a/.gitignore +++ b/.gitignore @@ -17,6 +17,7 @@ __pycache__/ venv/ venv_pptx2md/ .benchmarks/ +.demo-originals/ .pytest_cache/ .ruff_cache/ .mypy_cache/ diff --git a/README.md b/README.md index f6e2513e6..ec054b44c 100644 --- a/README.md +++ b/README.md @@ -222,18 +222,21 @@ full dependency matrix. ./deploy/local-dev/start-dev.sh ``` -5. Start the API and worker in separate terminals: +5. Run migrations with the local management role. The infrastructure script +provisions `knowhere_runtime`; the API/Worker environment examples use that +non-superuser role so shared demo row-level security is enforced. ```bash -cd apps/api && uv run main.py -cd apps/worker && uv run worker.py +cd apps/api +MIGRATION_DATABASE_URL=postgresql+asyncpg://root:root123@localhost:5432/Knowhere \ + uv run alembic upgrade heads ``` -Run API migrations explicitly before starting the API when the database schema needs updating: +6. Start the API and worker in separate terminals: ```bash -cd apps/api -uv run alembic upgrade heads +cd apps/api && uv run main.py +cd apps/worker && uv run worker.py ``` For API-only development without the dashboard, create an API-only user/key diff --git a/apps/api/.env.example b/apps/api/.env.example index bba577426..c1e8dafd9 100644 --- a/apps/api/.env.example +++ b/apps/api/.env.example @@ -43,7 +43,7 @@ TELEMETRY_ENABLED=true # TELEMETRY_AGGREGATE_INTERVAL_SECONDS=300 # Required for local startup: database -DATABASE_URL=postgresql+asyncpg://root:root123@localhost:5432/Knowhere +DATABASE_URL=postgresql+asyncpg://knowhere_runtime:runtime123@localhost:5432/Knowhere DB_SSL_MODE=disable # DB_SSL_MODE options: disable, allow, prefer, require, verify-ca, verify-full # DB_SSL_CERT=/path/to/client-cert.pem @@ -136,12 +136,6 @@ MOESIF_APPLICATION_ID= FREE_PLAN_INITIAL_CREDITS=5 FRONTEND_URL=http://localhost:3000 -# Reuse canonical demo bundles; disable to restore per-job uploads while keeping -# existing canonical artifact URLs readable. Does not change publication strategy. -DEMO_CANONICAL_BUNDLE_ENABLED=false -# Optional per-process tokenization/frequency reuse for versioned demo sources. -DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED=false - # Required for specific features: parsing providers MINERU_API_KEYS= ILOVEAPI_PUBLIC_KEY= @@ -168,3 +162,9 @@ ENTITY_TYPES=person,location,organization # Re-summarize detected page-track figures/charts/tables with numeric extraction # (ยง4.3). Gated separately from detection for safe rollout; defaults off. + +# Shared demo maintainers still require normal API write permission. +DEMO_MAINTAINER_USER_IDS= +# API/Worker DATABASE_URL must use NOSUPERUSER NOBYPASSRLS. +# Use a separate privileged role for migrations. +MIGRATION_DATABASE_URL= diff --git a/apps/api/alembic/demo-corpus-catalog.json b/apps/api/alembic/demo-corpus-catalog.json new file mode 100644 index 000000000..644e1bbbf --- /dev/null +++ b/apps/api/alembic/demo-corpus-catalog.json @@ -0,0 +1,534 @@ +[ + { + "demo_source_id": "demo-tsla-q4-2025", + "document_id": "ddoc_d55868b790cd", + "title": "TSLA-Q4-2025-Update.pdf", + "category_id": "other", + "file_name": "TSLA-Q4-2025-Update.pdf", + "catalog_metadata": { + "examples": [ + { + "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. [[cite:1]]\n\nThe 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. [[cite:1]]", + "citations": [ + { + "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", + "page_number": null + } + ] + }, + { + "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. [[cite:1]]\n\nEnergy gross profit reached a record $1.1 billion, marking the fifth consecutive record quarter. [[cite:1]]\n\nTesla also plans to begin Megapack 3 and Megablock production at Megafactory Houston in 2026. [[cite:1]]", + "citations": [ + { + "section_path": "TSLA-Q4-2025-Update.pdf/MANUFACTURING & HARDWARE/Energy generation and storage", + "description": "Storage deployment growth", + "content": "record quarterly energy storage deployments driven by Megapack", + "page_number": null + } + ] + }, + { + "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. [[cite:1]]\n\nThe same product update also notes that first-generation Optimus production lines are being installed before volume production. [[cite:1]]", + "citations": [ + { + "section_path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Volume", + "description": "2026 production plans", + "content": "Cybercab, Tesla Semi and Megapack 3 are on schedule for volume production starting in 2026.", + "page_number": null + } + ] + } + ] + } + }, + { + "demo_source_id": "demo-spacex-s1", + "document_id": "ddoc_ceb0d9f0c884", + "title": "spacex-s1.pdf", + "category_id": "financial-reports", + "file_name": "spacex-s1.pdf", + "catalog_metadata": { + "library_source_id": "financial-spacex-s1", + "category_id": "financial-reports", + "title": "spacex-s1.pdf", + "source_url": "https://data.olivierroy.dev/spacex-s1.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-spacex-s1", + "examples": [ + { + "id": "demo-spacex-s1-starlink-scale", + "question": "What does the filing say about Starlink's scale?", + "answer": "The filing says SpaceX operates a high-speed, low-latency global broadband network powered by about 9,600 Starlink broadband and mobile satellites in Low-Earth Orbit. [[cite:1]]\n\nIt says that network serves consumer, enterprise, and government customers across 164 countries, territories, and other markets as of March 31, 2026. [[cite:1]]", + "citations": [ + { + "section_path": "spacex-s1.pdf/Root", + "description": "Starlink network scale", + "content": "approximately 9,600 Starlink broadband and mobile satellites", + "page_number": 28 + } + ] + }, + { + "id": "demo-spacex-s1-launch-reusability", + "question": "How does the filing describe SpaceX's launch reusability?", + "answer": "The filing says Falcon 9 reusability gave SpaceX a step-function cost advantage in space access. [[cite:1]]\n\nIt also says Falcon 9 first stages had demonstrated the ability to refly 34 times as of March 31, 2026. [[cite:1]]", + "citations": [ + { + "section_path": "spacex-s1.pdf/Root", + "description": "Falcon 9 booster reuse", + "content": "refly a first-stage 34 times", + "page_number": 32 + } + ] + }, + { + "id": "demo-spacex-s1-offering", + "question": "What offering does the filing describe?", + "answer": "The filing describes an initial public offering of shares of Space Exploration Technologies Corp. Class A common stock. [[cite:1]]", + "citations": [ + { + "section_path": "spacex-s1.pdf/Root", + "description": "Class A common stock offering", + "content": "This is the initial public offering of shares of Class A common stock", + "page_number": 2 + } + ] + } + ] + } + }, + { + "demo_source_id": "demo-financial-nvda-q1-fy27-earnings-call", + "document_id": "ddoc_f36fa528c4ec", + "title": "NVDA Q1 FY27 Earnings Call Transcript.pdf", + "category_id": "financial-reports", + "file_name": "NVDA Q1 FY27 Earnings Call Transcript.pdf", + "catalog_metadata": { + "library_source_id": "financial-nvda-q1-fy27-earnings-call", + "category_id": "financial-reports", + "title": "NVDA Q1 FY27 Earnings Call Transcript.pdf", + "source_url": "https://s201.q4cdn.com/141608511/files/doc_financials/2027/q1/NVDA-Q1-2027-Earnings-Call-20-May-2026-5_00-PM-ET.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-financial-nvda-q1-fy27-earnings-call", + "examples": [] + } + }, + { + "demo_source_id": "demo-financial-nvda-q4-fy26-earnings-call", + "document_id": "ddoc_18789ec7df05", + "title": "NVDA Q4 FY26 Earnings Call Transcript.pdf", + "category_id": "financial-reports", + "file_name": "NVDA Q4 FY26 Earnings Call Transcript.pdf", + "catalog_metadata": { + "library_source_id": "financial-nvda-q4-fy26-earnings-call", + "category_id": "financial-reports", + "title": "NVDA Q4 FY26 Earnings Call Transcript.pdf", + "source_url": "https://s201.q4cdn.com/141608511/files/doc_financials/2026/q4/NVDA-Q4-2026-Earnings-Call-25-February-2026-5_00-PM-ET.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-financial-nvda-q4-fy26-earnings-call", + "examples": [] + } + }, + { + "demo_source_id": "demo-financial-nvda-q1-fy27-cfo-commentary", + "document_id": "ddoc_4abff126ee50", + "title": "NVIDIA Q1 FY27 CFO Commentary.pdf", + "category_id": "financial-reports", + "file_name": "NVIDIA Q1 FY27 CFO Commentary.pdf", + "catalog_metadata": { + "library_source_id": "financial-nvda-q1-fy27-cfo-commentary", + "category_id": "financial-reports", + "title": "NVIDIA Q1 FY27 CFO Commentary.pdf", + "source_url": "https://s201.q4cdn.com/141608511/files/doc_financials/2027/Q127/Q1FY27-CFO-Commentary.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-financial-nvda-q1-fy27-cfo-commentary", + "examples": [] + } + }, + { + "demo_source_id": "demo-financial-nvda-q1-fy27-results", + "document_id": "ddoc_e66754ebcddf", + "title": "NVIDIA Q1 FY27 Results.pdf", + "category_id": "financial-reports", + "file_name": "NVIDIA Q1 FY27 Results.pdf", + "catalog_metadata": { + "library_source_id": "financial-nvda-q1-fy27-results", + "category_id": "financial-reports", + "title": "NVIDIA Q1 FY27 Results.pdf", + "source_url": "https://s201.q4cdn.com/141608511/files/doc_financials/2027/q1/927dc2d6-a76c-4006-9f34-8769b2c665fb.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-financial-nvda-q1-fy27-results", + "examples": [] + } + }, + { + "demo_source_id": "demo-financial-nvda-q1-fy27-presentation", + "document_id": "ddoc_6333176194f6", + "title": "NVIDIA Q1 FY27 Quarterly Presentation.pdf", + "category_id": "financial-reports", + "file_name": "NVIDIA Q1 FY27 Quarterly Presentation.pdf", + "catalog_metadata": { + "library_source_id": "financial-nvda-q1-fy27-presentation", + "category_id": "financial-reports", + "title": "NVIDIA Q1 FY27 Quarterly Presentation.pdf", + "source_url": "https://s201.q4cdn.com/141608511/files/doc_financials/2027/Q127/NVDA-F1Q27-Quarterly-Presentation-FINAL.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-financial-nvda-q1-fy27-presentation", + "examples": [] + } + }, + { + "demo_source_id": "demo-financial-nvda-fy26-annual-report", + "document_id": "ddoc_9b61ace7640d", + "title": "NVIDIA FY26 Annual Report.pdf", + "category_id": "financial-reports", + "file_name": "NVIDIA FY26 Annual Report.pdf", + "catalog_metadata": { + "library_source_id": "financial-nvda-fy26-annual-report", + "category_id": "financial-reports", + "title": "NVIDIA FY26 Annual Report.pdf", + "source_url": "https://s201.q4cdn.com/141608511/files/doc_financials/2026/ar/2026-Annual-Report-Web.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-financial-nvda-fy26-annual-report", + "examples": [] + } + }, + { + "demo_source_id": "demo-financial-micron-report-530bd7ed", + "document_id": "ddoc_9cd099f0392c", + "title": "Micron investor report 530bd7ed.pdf", + "category_id": "financial-reports", + "file_name": "Micron investor report 530bd7ed.pdf", + "catalog_metadata": { + "library_source_id": "financial-micron-report-530bd7ed", + "category_id": "financial-reports", + "title": "Micron investor report 530bd7ed.pdf", + "source_url": "https://investors.micron.com/static-files/530bd7ed-a8c8-4687-af4a-8c129f740e09", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-financial-micron-report-530bd7ed", + "examples": [] + } + }, + { + "demo_source_id": "demo-financial-micron-report-9c0becf5", + "document_id": "ddoc_2da3db527cb0", + "title": "Micron investor report 9c0becf5.pdf", + "category_id": "financial-reports", + "file_name": "Micron investor report 9c0becf5.pdf", + "catalog_metadata": { + "library_source_id": "financial-micron-report-9c0becf5", + "category_id": "financial-reports", + "title": "Micron investor report 9c0becf5.pdf", + "source_url": "https://investors.micron.com/static-files/9c0becf5-df56-4eec-bd67-453dda68b273", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-financial-micron-report-9c0becf5", + "examples": [] + } + }, + { + "demo_source_id": "demo-financial-goog-10-k-2025", + "document_id": "ddoc_cf7924015c73", + "title": "GOOG 10-K 2025.pdf", + "category_id": "financial-reports", + "file_name": "GOOG 10-K 2025.pdf", + "catalog_metadata": { + "library_source_id": "financial-goog-10-k-2025", + "category_id": "financial-reports", + "title": "GOOG 10-K 2025.pdf", + "source_url": "https://s206.q4cdn.com/479360582/files/doc_financials/2025/q4/GOOG-10-K-2025.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-financial-goog-10-k-2025", + "examples": [] + } + }, + { + "demo_source_id": "demo-financial-meta-q1-2026-results", + "document_id": "ddoc_cf2ca1b29b2b", + "title": "Meta Q1 2026 Exhibit 99.1.pdf", + "category_id": "financial-reports", + "file_name": "Meta Q1 2026 Exhibit 99.1.pdf", + "catalog_metadata": { + "library_source_id": "financial-meta-q1-2026-results", + "category_id": "financial-reports", + "title": "Meta Q1 2026 Exhibit 99.1.pdf", + "source_url": "https://s21.q4cdn.com/399680738/files/doc_financials/2026/q1/Meta-03-31-2026-Exhibit-99-1_final.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-financial-meta-q1-2026-results", + "examples": [] + } + }, + { + "demo_source_id": "demo-financial-meta-q1-2026-presentation", + "document_id": "ddoc_8437a87289ca", + "title": "Meta Q1 2026 Earnings Presentation.pdf", + "category_id": "financial-reports", + "file_name": "Meta Q1 2026 Earnings Presentation.pdf", + "catalog_metadata": { + "library_source_id": "financial-meta-q1-2026-presentation", + "category_id": "financial-reports", + "title": "Meta Q1 2026 Earnings Presentation.pdf", + "source_url": "https://s21.q4cdn.com/399680738/files/doc_financials/2026/q1/Earnings-Presentation-Q1-2026.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-financial-meta-q1-2026-presentation", + "examples": [] + } + }, + { + "demo_source_id": "demo-financial-meta-q1-2026-earnings-call", + "document_id": "ddoc_73049be42462", + "title": "Meta Q1 2026 Earnings Call Transcript.pdf", + "category_id": "financial-reports", + "file_name": "Meta Q1 2026 Earnings Call Transcript.pdf", + "catalog_metadata": { + "library_source_id": "financial-meta-q1-2026-earnings-call", + "category_id": "financial-reports", + "title": "Meta Q1 2026 Earnings Call Transcript.pdf", + "source_url": "https://s21.q4cdn.com/399680738/files/doc_financials/2026/q1/META-Q1-2026-Earnings-Call-Transcript.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-financial-meta-q1-2026-earnings-call", + "examples": [] + } + }, + { + "demo_source_id": "demo-financial-microsoft-2025-annual-report", + "document_id": "ddoc_66092bed83d7", + "title": "Microsoft 2025 Annual Report.docx", + "category_id": "financial-reports", + "file_name": "Microsoft 2025 Annual Report.docx", + "catalog_metadata": { + "library_source_id": "financial-microsoft-2025-annual-report", + "category_id": "financial-reports", + "title": "Microsoft 2025 Annual Report.docx", + "source_url": "https://cdn-dynmedia-1.microsoft.com/is/content/microsoftcorp/2025_AnnualReport.docx", + "mime_type": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "status": "ready", + "demo_source_id": "demo-financial-microsoft-2025-annual-report", + "examples": [] + } + }, + { + "demo_source_id": "demo-financial-alibaba-fy2026-interim-report", + "document_id": "ddoc_7059f6464955", + "title": "Alibaba Fiscal Year 2026 Interim Report.pdf", + "category_id": "financial-reports", + "file_name": "Alibaba Fiscal Year 2026 Interim Report.pdf", + "catalog_metadata": { + "library_source_id": "financial-alibaba-fy2026-interim-report", + "category_id": "financial-reports", + "title": "Alibaba Fiscal Year 2026 Interim Report.pdf", + "source_url": "https://data.alibabagroup.com/ecms-files/1514443390/2b483071-b7a2-45ea-b6cd-f4d35861bab0/Fiscal%20Year%202026%20Interim%20Report.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-financial-alibaba-fy2026-interim-report", + "examples": [] + } + }, + { + "demo_source_id": "demo-research-attention-is-all-you-need", + "document_id": "ddoc_871e2a7f3883", + "title": "Attention Is All You Need.pdf", + "category_id": "research-papers", + "file_name": "Attention Is All You Need.pdf", + "catalog_metadata": { + "library_source_id": "research-attention-is-all-you-need", + "category_id": "research-papers", + "title": "Attention Is All You Need.pdf", + "source_url": "https://arxiv.org/pdf/1706.03762", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-research-attention-is-all-you-need", + "examples": [] + } + }, + { + "demo_source_id": "demo-research-rag-survey", + "document_id": "ddoc_06dfbe90cf37", + "title": "Retrieval-Augmented Generation for Large Language Models: A Survey.pdf", + "category_id": "research-papers", + "file_name": "Retrieval-Augmented Generation for Large Language Models: A Survey.pdf", + "catalog_metadata": { + "library_source_id": "research-rag-survey", + "category_id": "research-papers", + "title": "Retrieval-Augmented Generation for Large Language Models: A Survey.pdf", + "source_url": "https://arxiv.org/pdf/2312.10997", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-research-rag-survey", + "examples": [] + } + }, + { + "demo_source_id": "demo-research-rag-realized", + "document_id": "ddoc_e200c74e90b0", + "title": "Retrieval-Augmented Generation Realized.pdf", + "category_id": "research-papers", + "file_name": "Retrieval-Augmented Generation Realized.pdf", + "catalog_metadata": { + "library_source_id": "research-rag-realized", + "category_id": "research-papers", + "title": "Retrieval-Augmented Generation Realized.pdf", + "source_url": "https://www.appliedai.de/uploads/files/retrieval-augmented-generation-realized/AppliedAI_White_Paper_Retrieval-augmented-Generation-Realized_FINAL_20240618.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-research-rag-realized", + "examples": [] + } + }, + { + "demo_source_id": "demo-research-toolformer", + "document_id": "ddoc_53cfd9dde93e", + "title": "Toolformer: Language Models Can Teach Themselves to Use Tools.pdf", + "category_id": "research-papers", + "file_name": "Toolformer: Language Models Can Teach Themselves to Use Tools.pdf", + "catalog_metadata": { + "library_source_id": "research-toolformer", + "category_id": "research-papers", + "title": "Toolformer: Language Models Can Teach Themselves to Use Tools.pdf", + "source_url": "https://arxiv.org/pdf/2302.04761", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-research-toolformer", + "examples": [] + } + }, + { + "demo_source_id": "demo-research-ai-agents-overview", + "document_id": "ddoc_62f4b379bb3f", + "title": "AI Agents Overview.pdf", + "category_id": "research-papers", + "file_name": "AI Agents Overview.pdf", + "catalog_metadata": { + "library_source_id": "research-ai-agents-overview", + "category_id": "research-papers", + "title": "AI Agents Overview.pdf", + "source_url": "https://cseweb.ucsd.edu/~yiying/cse291a-fall25/reading/ai-agents.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-research-ai-agents-overview", + "examples": [] + } + }, + { + "demo_source_id": "demo-stem-jurafsky-transformers", + "document_id": "ddoc_b0c768360858", + "title": "Speech and Language Processing, Chapter 8: Transformers.pdf", + "category_id": "stem-books", + "file_name": "Speech and Language Processing, Chapter 8: Transformers.pdf", + "catalog_metadata": { + "library_source_id": "stem-jurafsky-transformers", + "category_id": "stem-books", + "title": "Speech and Language Processing, Chapter 8: Transformers.pdf", + "source_url": "https://web.stanford.edu/~jurafsky/slp3/8.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-stem-jurafsky-transformers", + "examples": [] + } + }, + { + "demo_source_id": "demo-stem-transformers-tutorial", + "document_id": "ddoc_2c336b703af1", + "title": "Transformers Tutorial.pdf", + "category_id": "stem-books", + "file_name": "Transformers Tutorial.pdf", + "catalog_metadata": { + "library_source_id": "stem-transformers-tutorial", + "category_id": "stem-books", + "title": "Transformers Tutorial.pdf", + "source_url": "https://aihub.csic.es/wp-content/uploads/2024/07/Transformers-tutorial.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-stem-transformers-tutorial", + "examples": [] + } + }, + { + "demo_source_id": "demo-stem-deep-learning-transformer-network", + "document_id": "ddoc_5e1989f36292", + "title": "Deep Learning - Transformer Network.pdf", + "category_id": "stem-books", + "file_name": "Deep Learning - Transformer Network.pdf", + "catalog_metadata": { + "library_source_id": "stem-deep-learning-transformer-network", + "category_id": "stem-books", + "title": "Deep Learning - Transformer Network.pdf", + "source_url": "https://homel.vsb.cz/~pla06/files/dl/dl_10.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-stem-deep-learning-transformer-network", + "examples": [] + } + }, + { + "demo_source_id": "demo-stem-probability-theory-lecture-notes", + "document_id": "ddoc_0ea07bf4aa34", + "title": "Probability Theory Lecture Notes.pdf", + "category_id": "stem-books", + "file_name": "Probability Theory Lecture Notes.pdf", + "catalog_metadata": { + "library_source_id": "stem-probability-theory-lecture-notes", + "category_id": "stem-books", + "title": "Probability Theory Lecture Notes.pdf", + "source_url": "https://www.math.union.edu/~marianop/ma41600sum19/Notes.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-stem-probability-theory-lecture-notes", + "examples": [] + } + }, + { + "demo_source_id": "demo-stem-introduction-statistical-learning-theory", + "document_id": "ddoc_68dc69595b03", + "title": "Introduction to Statistical Learning Theory.pdf", + "category_id": "stem-books", + "file_name": "Introduction to Statistical Learning Theory.pdf", + "catalog_metadata": { + "library_source_id": "stem-introduction-statistical-learning-theory", + "category_id": "stem-books", + "title": "Introduction to Statistical Learning Theory.pdf", + "source_url": "https://cciliber.github.io/intro-slt/slides/lec1.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-stem-introduction-statistical-learning-theory", + "examples": [] + } + }, + { + "demo_source_id": "demo-stem-statistical-learning", + "document_id": "ddoc_fcffc765aeaf", + "title": "Statistical Learning Notes.pdf", + "category_id": "stem-books", + "file_name": "Statistical Learning Notes.pdf", + "catalog_metadata": { + "library_source_id": "stem-statistical-learning", + "category_id": "stem-books", + "title": "Statistical Learning Notes.pdf", + "source_url": "http://www.statslab.cam.ac.uk/~rds37/teaching/stat_learning/notes.pdf", + "mime_type": "application/pdf", + "status": "ready", + "demo_source_id": "demo-stem-statistical-learning", + "examples": [] + } + } +] diff --git a/apps/api/alembic/demo-corpus-policies.sql b/apps/api/alembic/demo-corpus-policies.sql new file mode 100644 index 000000000..7e59f1bab --- /dev/null +++ b/apps/api/alembic/demo-corpus-policies.sql @@ -0,0 +1,55 @@ +CREATE OR REPLACE FUNCTION is_demo_revision_ready(document_identity text, revision_identity text) +RETURNS boolean LANGUAGE sql STABLE SET search_path = public, pg_temp AS $$ + SELECT EXISTS ( + SELECT 1 FROM demo_documents d JOIN job_results r ON r.demo_document_id = d.document_id + JOIN demo_document_map_unit_indexes i ON i.document_id = d.document_id AND i.job_result_id = r.id + WHERE d.document_id = document_identity AND r.id = revision_identity AND d.status = 'active' + ) +$$; +ALTER TABLE demo_documents ENABLE ROW LEVEL SECURITY; +ALTER TABLE demo_documents FORCE ROW LEVEL SECURITY; +CREATE POLICY demo_maintainer_write ON demo_documents FOR ALL USING (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> '') WITH CHECK (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> ''); +CREATE POLICY demo_public_read ON demo_documents FOR SELECT USING (status = 'active'); +ALTER TABLE demo_retrieval_namespace_generations ENABLE ROW LEVEL SECURITY; +ALTER TABLE demo_retrieval_namespace_generations FORCE ROW LEVEL SECURITY; +CREATE POLICY demo_maintainer_write ON demo_retrieval_namespace_generations FOR ALL USING (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> '') WITH CHECK (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> ''); +CREATE POLICY demo_public_read ON demo_retrieval_namespace_generations FOR SELECT USING (true); +ALTER TABLE demo_retrieval_namespace_map_snapshots ENABLE ROW LEVEL SECURITY; +ALTER TABLE demo_retrieval_namespace_map_snapshots FORCE ROW LEVEL SECURITY; +CREATE POLICY demo_maintainer_write ON demo_retrieval_namespace_map_snapshots FOR ALL USING (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> '') WITH CHECK (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> ''); +CREATE POLICY demo_public_read ON demo_retrieval_namespace_map_snapshots FOR SELECT USING (true); +ALTER TABLE demo_retrieval_hit_stats ENABLE ROW LEVEL SECURITY; +ALTER TABLE demo_retrieval_hit_stats FORCE ROW LEVEL SECURITY; +CREATE POLICY demo_caller_stats ON demo_retrieval_hit_stats FOR ALL USING (user_id = NULLIF(current_setting('knowhere.demo_caller', true), '')) WITH CHECK (user_id = NULLIF(current_setting('knowhere.demo_caller', true), '') AND namespace = '__knowhere_demo__'); +ALTER TABLE demo_document_map_unit_indexes ENABLE ROW LEVEL SECURITY; +ALTER TABLE demo_document_map_unit_indexes FORCE ROW LEVEL SECURITY; +CREATE POLICY demo_maintainer_write ON demo_document_map_unit_indexes FOR ALL USING (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> '') WITH CHECK (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> ''); +CREATE POLICY demo_public_read ON demo_document_map_unit_indexes FOR SELECT USING (EXISTS (SELECT 1 FROM job_results r JOIN demo_documents d ON d.document_id = r.demo_document_id WHERE r.id = demo_document_map_unit_indexes.job_result_id AND d.document_id = demo_document_map_unit_indexes.document_id AND d.status = 'active')); +ALTER TABLE demo_document_sections ENABLE ROW LEVEL SECURITY; +ALTER TABLE demo_document_sections FORCE ROW LEVEL SECURITY; +CREATE POLICY demo_maintainer_write ON demo_document_sections FOR ALL USING (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> '') WITH CHECK (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> ''); +CREATE POLICY demo_public_read ON demo_document_sections FOR SELECT USING (is_demo_revision_ready(document_id, job_result_id)); +ALTER TABLE demo_graph_nodes ENABLE ROW LEVEL SECURITY; +ALTER TABLE demo_graph_nodes FORCE ROW LEVEL SECURITY; +CREATE POLICY demo_maintainer_write ON demo_graph_nodes FOR ALL USING (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> '') WITH CHECK (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> ''); +CREATE POLICY demo_public_read ON demo_graph_nodes FOR SELECT USING (is_demo_revision_ready(owner_document_id, job_result_id)); +ALTER TABLE demo_retrieval_serving_revision_manifests ENABLE ROW LEVEL SECURITY; +ALTER TABLE demo_retrieval_serving_revision_manifests FORCE ROW LEVEL SECURITY; +CREATE POLICY demo_maintainer_write ON demo_retrieval_serving_revision_manifests FOR ALL USING (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> '') WITH CHECK (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> ''); +CREATE POLICY demo_public_read ON demo_retrieval_serving_revision_manifests FOR SELECT USING (is_demo_revision_ready(document_id, job_result_id)); +ALTER TABLE demo_document_chunks ENABLE ROW LEVEL SECURITY; +ALTER TABLE demo_document_chunks FORCE ROW LEVEL SECURITY; +CREATE POLICY demo_maintainer_write ON demo_document_chunks FOR ALL USING (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> '') WITH CHECK (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> ''); +CREATE POLICY demo_public_read ON demo_document_chunks FOR SELECT USING (is_demo_revision_ready(document_id, job_result_id)); +ALTER TABLE demo_document_map_units ENABLE ROW LEVEL SECURITY; +ALTER TABLE demo_document_map_units FORCE ROW LEVEL SECURITY; +CREATE POLICY demo_maintainer_write ON demo_document_map_units FOR ALL USING (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> '') WITH CHECK (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> ''); +CREATE POLICY demo_public_read ON demo_document_map_units FOR SELECT USING (is_demo_revision_ready(document_id, job_result_id)); +ALTER TABLE demo_graph_edges ENABLE ROW LEVEL SECURITY; +ALTER TABLE demo_graph_edges FORCE ROW LEVEL SECURITY; +CREATE POLICY demo_maintainer_write ON demo_graph_edges FOR ALL USING (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> '') WITH CHECK (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> ''); +CREATE POLICY demo_public_read ON demo_graph_edges FOR SELECT USING (is_demo_revision_ready(owner_document_id, job_result_id)); +ALTER TABLE demo_document_map_unit_tokens ENABLE ROW LEVEL SECURITY; +ALTER TABLE demo_document_map_unit_tokens FORCE ROW LEVEL SECURITY; +CREATE POLICY demo_maintainer_write ON demo_document_map_unit_tokens FOR ALL USING (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> '') WITH CHECK (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> ''); +CREATE POLICY demo_public_read ON demo_document_map_unit_tokens FOR SELECT USING (EXISTS (SELECT 1 FROM demo_document_map_units u WHERE u.id = map_unit_id)); diff --git a/apps/api/alembic/demo-corpus-schema.sql b/apps/api/alembic/demo-corpus-schema.sql new file mode 100644 index 000000000..9af5a276a --- /dev/null +++ b/apps/api/alembic/demo-corpus-schema.sql @@ -0,0 +1,287 @@ +ALTER TABLE job_results ADD COLUMN IF NOT EXISTS demo_document_id varchar(36); +ALTER TABLE job_results ADD CONSTRAINT ck_job_results_corpus_target CHECK (document_id IS NULL OR demo_document_id IS NULL); +ALTER TABLE job_results ADD CONSTRAINT uq_job_results_demo_revision UNIQUE (id, demo_document_id); + +CREATE TABLE IF NOT EXISTS demo_documents ( + document_id VARCHAR(36) NOT NULL, + demo_source_id VARCHAR(128) NOT NULL, + title TEXT NOT NULL, + category_id VARCHAR(128) NOT NULL, + catalog_metadata JSON NOT NULL, + user_id TEXT NOT NULL, + namespace VARCHAR(255) NOT NULL, + status VARCHAR(32) NOT NULL, + current_job_result_id VARCHAR(36), + source_file_name TEXT, + document_metadata JSON, + parse_track VARCHAR(32) NOT NULL, + created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + updated_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + archived_at TIMESTAMP WITHOUT TIME ZONE, + PRIMARY KEY (document_id), + CONSTRAINT ck_demo_documents_namespace CHECK (namespace = '__knowhere_demo__'), + UNIQUE (demo_source_id) +) + +; +CREATE INDEX IF NOT EXISTS idx_demo_documents_current_job_result ON demo_documents (current_job_result_id); +CREATE INDEX IF NOT EXISTS idx_demo_documents_user_namespace_status ON demo_documents (user_id, namespace, status); + +CREATE TABLE IF NOT EXISTS demo_retrieval_namespace_generations ( + id VARCHAR(100) NOT NULL, + user_id TEXT NOT NULL, + namespace VARCHAR(255) NOT NULL, + generation BIGINT NOT NULL, + updated_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + PRIMARY KEY (id), + CONSTRAINT uq_demo_retrieval_namespace_generations_scope UNIQUE (user_id, namespace) +) + +; + +CREATE TABLE IF NOT EXISTS demo_retrieval_namespace_map_snapshots ( + id VARCHAR(100) NOT NULL, + user_id TEXT NOT NULL, + namespace VARCHAR(255) NOT NULL, + generation BIGINT NOT NULL, + format_version INTEGER NOT NULL, + payload_zlib BYTEA NOT NULL, + checksum VARCHAR(64) NOT NULL, + updated_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + PRIMARY KEY (id), + CONSTRAINT uq_demo_retrieval_namespace_map_snapshots_scope UNIQUE (user_id, namespace) +) + +; + +CREATE TABLE IF NOT EXISTS demo_retrieval_hit_stats ( + id VARCHAR(36) NOT NULL, + user_id TEXT NOT NULL, + namespace VARCHAR(255) NOT NULL, + hit_kind VARCHAR(32) NOT NULL, + document_id VARCHAR(36) NOT NULL, + chunk_id VARCHAR(64), + hit_count INTEGER NOT NULL, + last_hit_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + updated_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + PRIMARY KEY (id), + FOREIGN KEY(document_id) REFERENCES demo_documents (document_id) ON DELETE CASCADE +) + +; +CREATE INDEX IF NOT EXISTS idx_demo_retrieval_hit_stats_chunk ON demo_retrieval_hit_stats (chunk_id); +CREATE INDEX IF NOT EXISTS idx_demo_retrieval_hit_stats_document ON demo_retrieval_hit_stats (document_id); +CREATE INDEX IF NOT EXISTS idx_demo_retrieval_hit_stats_scope_kind ON demo_retrieval_hit_stats (user_id, namespace, hit_kind); +CREATE UNIQUE INDEX IF NOT EXISTS uq_demo_retrieval_hit_stats_chunk_key ON demo_retrieval_hit_stats (user_id, namespace, hit_kind, document_id, chunk_id) WHERE chunk_id IS NOT NULL; +CREATE UNIQUE INDEX IF NOT EXISTS uq_demo_retrieval_hit_stats_document_key ON demo_retrieval_hit_stats (user_id, namespace, hit_kind, document_id) WHERE chunk_id IS NULL; + +CREATE TABLE IF NOT EXISTS demo_document_map_unit_indexes ( + id VARCHAR(100) NOT NULL, + document_id VARCHAR(36) NOT NULL, + job_result_id VARCHAR(36) NOT NULL, + format_version INTEGER NOT NULL, + unit_count INTEGER NOT NULL, + token_count INTEGER NOT NULL, + average_idf_path FLOAT NOT NULL, + average_idf_content FLOAT NOT NULL, + path_document_count INTEGER, + path_total_length INTEGER, + content_document_count INTEGER, + content_total_length INTEGER, + created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + PRIMARY KEY (id), + CONSTRAINT fk_demo_documentmapunitindex_revision FOREIGN KEY(job_result_id, document_id) REFERENCES job_results (id, demo_document_id) ON DELETE RESTRICT, + CONSTRAINT uq_demo_document_map_unit_indexes_revision UNIQUE (document_id, job_result_id), + FOREIGN KEY(document_id) REFERENCES demo_documents (document_id) ON DELETE CASCADE, + FOREIGN KEY(job_result_id) REFERENCES job_results (id) ON DELETE RESTRICT +) + +; +CREATE INDEX IF NOT EXISTS idx_demo_document_map_unit_indexes_revision ON demo_document_map_unit_indexes (document_id, job_result_id); + +CREATE TABLE IF NOT EXISTS demo_document_sections ( + section_id VARCHAR(36) NOT NULL, + user_id TEXT NOT NULL, + namespace VARCHAR(255) NOT NULL, + document_id VARCHAR(36) NOT NULL, + job_result_id VARCHAR(36) NOT NULL, + parent_section_id VARCHAR(36), + section_path TEXT NOT NULL, + section_title TEXT, + section_level INTEGER NOT NULL, + summary TEXT, + section_metadata JSON, + sort_order INTEGER NOT NULL, + created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + PRIMARY KEY (section_id), + CONSTRAINT uq_demo_sections_identity UNIQUE (section_id, document_id, job_result_id), + CONSTRAINT fk_demo_sections_parent_revision FOREIGN KEY(parent_section_id, document_id, job_result_id) REFERENCES demo_document_sections (section_id, document_id, job_result_id) DEFERRABLE INITIALLY DEFERRED, + CONSTRAINT fk_demo_documentsection_revision FOREIGN KEY(job_result_id, document_id) REFERENCES job_results (id, demo_document_id) ON DELETE RESTRICT, + CONSTRAINT uq_demo_document_sections_revision_path UNIQUE (document_id, job_result_id, section_path), + FOREIGN KEY(document_id) REFERENCES demo_documents (document_id) ON DELETE CASCADE, + FOREIGN KEY(job_result_id) REFERENCES job_results (id) ON DELETE RESTRICT, + FOREIGN KEY(parent_section_id) REFERENCES demo_document_sections (section_id) ON DELETE SET NULL +) + +; +CREATE INDEX IF NOT EXISTS idx_demo_document_sections_doc_revision ON demo_document_sections (document_id, job_result_id); +CREATE INDEX IF NOT EXISTS idx_demo_document_sections_revision_snapshot_order ON demo_document_sections (document_id, job_result_id, sort_order, section_id); +CREATE INDEX IF NOT EXISTS idx_demo_document_sections_scope ON demo_document_sections (user_id, namespace); + +CREATE TABLE IF NOT EXISTS demo_graph_nodes ( + node_id VARCHAR(128) NOT NULL, + user_id TEXT NOT NULL, + namespace VARCHAR(255) NOT NULL, + node_kind VARCHAR(32) NOT NULL, + owner_document_id VARCHAR(36) NOT NULL, + job_result_id VARCHAR(36) NOT NULL, + ref_document_id VARCHAR(36), + ref_section_id VARCHAR(36), + properties JSON, + created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + updated_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + PRIMARY KEY (node_id), + CONSTRAINT fk_demo_graphnode_revision FOREIGN KEY(job_result_id, owner_document_id) REFERENCES job_results (id, demo_document_id) ON DELETE RESTRICT, + FOREIGN KEY(owner_document_id) REFERENCES demo_documents (document_id) ON DELETE CASCADE, + FOREIGN KEY(job_result_id) REFERENCES job_results (id) ON DELETE RESTRICT +) + +; +CREATE INDEX IF NOT EXISTS idx_demo_graph_nodes_owner_revision ON demo_graph_nodes (owner_document_id, job_result_id); +CREATE INDEX IF NOT EXISTS idx_demo_graph_nodes_ref_document ON demo_graph_nodes (ref_document_id); +CREATE INDEX IF NOT EXISTS idx_demo_graph_nodes_ref_section ON demo_graph_nodes (ref_section_id); +CREATE INDEX IF NOT EXISTS idx_demo_graph_nodes_scope ON demo_graph_nodes (user_id, namespace, node_kind); +CREATE INDEX IF NOT EXISTS idx_demo_graph_nodes_top_entities_gin ON demo_graph_nodes USING gin ((properties::jsonb -> 'top_entities')); +CREATE INDEX IF NOT EXISTS idx_demo_graph_nodes_top_keywords_gin ON demo_graph_nodes USING gin ((properties::jsonb -> 'top_keywords')); + +CREATE TABLE IF NOT EXISTS demo_retrieval_serving_revision_manifests ( + id VARCHAR(100) NOT NULL, + user_id TEXT NOT NULL, + namespace VARCHAR(255) NOT NULL, + document_id VARCHAR(36) NOT NULL, + job_result_id VARCHAR(36) NOT NULL, + format_version INTEGER NOT NULL, + payload_zlib BYTEA NOT NULL, + checksum VARCHAR(64) NOT NULL, + created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + PRIMARY KEY (id), + CONSTRAINT fk_demo_retrievalservingrevisionmanifest_revision FOREIGN KEY(job_result_id, document_id) REFERENCES job_results (id, demo_document_id) ON DELETE RESTRICT, + CONSTRAINT uq_demo_retrieval_serving_revision_manifests_revision UNIQUE (document_id, job_result_id), + FOREIGN KEY(document_id) REFERENCES demo_documents (document_id) ON DELETE CASCADE, + FOREIGN KEY(job_result_id) REFERENCES job_results (id) ON DELETE RESTRICT +) + +; +CREATE INDEX IF NOT EXISTS idx_demo_retrieval_serving_revision_manifests_scope ON demo_retrieval_serving_revision_manifests (user_id, namespace, document_id, job_result_id); + +CREATE TABLE IF NOT EXISTS demo_document_chunks ( + id VARCHAR(36) NOT NULL, + chunk_id VARCHAR(64) NOT NULL, + user_id TEXT NOT NULL, + namespace VARCHAR(255) NOT NULL, + document_id VARCHAR(36) NOT NULL, + job_result_id VARCHAR(36) NOT NULL, + section_id VARCHAR(36), + chunk_type VARCHAR(64) NOT NULL, + content TEXT, + content_lexical_text TEXT, + path_lexical_text TEXT, + content_search_text TEXT, + path_search_text TEXT, + term_search_text TEXT, + source_chunk_path TEXT, + file_path TEXT, + chunk_metadata JSON, + sort_order INTEGER NOT NULL, + created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + PRIMARY KEY (id), + CONSTRAINT fk_demo_documentchunk_section_revision FOREIGN KEY(section_id, document_id, job_result_id) REFERENCES demo_document_sections (section_id, document_id, job_result_id) DEFERRABLE INITIALLY DEFERRED, + CONSTRAINT fk_demo_documentchunk_revision FOREIGN KEY(job_result_id, document_id) REFERENCES job_results (id, demo_document_id) ON DELETE RESTRICT, + CONSTRAINT uq_demo_document_chunks_revision_path UNIQUE (document_id, job_result_id, source_chunk_path), + FOREIGN KEY(document_id) REFERENCES demo_documents (document_id) ON DELETE CASCADE, + FOREIGN KEY(job_result_id) REFERENCES job_results (id) ON DELETE RESTRICT, + FOREIGN KEY(section_id) REFERENCES demo_document_sections (section_id) ON DELETE SET NULL +) + +; +CREATE INDEX IF NOT EXISTS idx_demo_document_chunks_chunk_id ON demo_document_chunks (chunk_id); +CREATE INDEX IF NOT EXISTS idx_demo_document_chunks_doc_revision ON demo_document_chunks (document_id, job_result_id); +CREATE INDEX IF NOT EXISTS idx_demo_document_chunks_revision_section_order ON demo_document_chunks (document_id, job_result_id, section_id, sort_order, chunk_id, id); +CREATE INDEX IF NOT EXISTS idx_demo_document_chunks_revision_snapshot_order ON demo_document_chunks (document_id, job_result_id, sort_order, chunk_id, id); +CREATE INDEX IF NOT EXISTS idx_demo_document_chunks_scope ON demo_document_chunks (user_id, namespace); +CREATE INDEX IF NOT EXISTS idx_demo_document_chunks_section ON demo_document_chunks (section_id); + +CREATE TABLE IF NOT EXISTS demo_document_map_units ( + id VARCHAR(160) NOT NULL, + document_id VARCHAR(36) NOT NULL, + job_result_id VARCHAR(36) NOT NULL, + unit_id VARCHAR(128) NOT NULL, + section_id VARCHAR(36) NOT NULL, + unit_kind VARCHAR(32) NOT NULL, + path_token_count INTEGER NOT NULL, + content_token_count INTEGER NOT NULL, + has_image BOOLEAN NOT NULL, + has_table BOOLEAN NOT NULL, + sort_order INTEGER NOT NULL, + created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + PRIMARY KEY (id), + CONSTRAINT fk_demo_documentmapunit_section_revision FOREIGN KEY(section_id, document_id, job_result_id) REFERENCES demo_document_sections (section_id, document_id, job_result_id) DEFERRABLE INITIALLY DEFERRED, + CONSTRAINT fk_demo_documentmapunit_revision FOREIGN KEY(job_result_id, document_id) REFERENCES job_results (id, demo_document_id) ON DELETE RESTRICT, + FOREIGN KEY(document_id) REFERENCES demo_documents (document_id) ON DELETE CASCADE, + FOREIGN KEY(job_result_id) REFERENCES job_results (id) ON DELETE RESTRICT +) + +; +CREATE INDEX IF NOT EXISTS idx_demo_document_map_units_has_image ON demo_document_map_units (document_id, job_result_id) WHERE has_image IS true; +CREATE INDEX IF NOT EXISTS idx_demo_document_map_units_has_table ON demo_document_map_units (document_id, job_result_id) WHERE has_table IS true; +CREATE INDEX IF NOT EXISTS idx_demo_document_map_units_revision_order ON demo_document_map_units (document_id, job_result_id, sort_order, unit_id); +CREATE INDEX IF NOT EXISTS idx_demo_document_map_units_section ON demo_document_map_units (section_id); + +CREATE TABLE IF NOT EXISTS demo_graph_edges ( + edge_id VARCHAR(160) NOT NULL, + user_id TEXT NOT NULL, + namespace VARCHAR(255) NOT NULL, + edge_kind VARCHAR(32) NOT NULL, + source_node_id VARCHAR(128) NOT NULL, + target_node_id VARCHAR(128) NOT NULL, + owner_document_id VARCHAR(36) NOT NULL, + job_result_id VARCHAR(36) NOT NULL, + is_directed BOOLEAN NOT NULL, + weight FLOAT, + properties JSON, + created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + updated_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + PRIMARY KEY (edge_id), + CONSTRAINT fk_demo_graphedge_revision FOREIGN KEY(job_result_id, owner_document_id) REFERENCES job_results (id, demo_document_id) ON DELETE RESTRICT, + FOREIGN KEY(source_node_id) REFERENCES demo_graph_nodes (node_id) ON DELETE CASCADE, + FOREIGN KEY(target_node_id) REFERENCES demo_graph_nodes (node_id) ON DELETE CASCADE, + FOREIGN KEY(owner_document_id) REFERENCES demo_documents (document_id) ON DELETE CASCADE, + FOREIGN KEY(job_result_id) REFERENCES job_results (id) ON DELETE RESTRICT +) + +; +CREATE INDEX IF NOT EXISTS idx_demo_graph_edges_owner_revision ON demo_graph_edges (owner_document_id, job_result_id); +CREATE INDEX IF NOT EXISTS idx_demo_graph_edges_scope ON demo_graph_edges (user_id, namespace, edge_kind); +CREATE INDEX IF NOT EXISTS idx_demo_graph_edges_source ON demo_graph_edges (source_node_id); +CREATE INDEX IF NOT EXISTS idx_demo_graph_edges_target ON demo_graph_edges (target_node_id); + +CREATE TABLE IF NOT EXISTS demo_document_map_unit_tokens ( + id VARCHAR(36) NOT NULL, + map_unit_id VARCHAR(160) NOT NULL, + channel VARCHAR(16) NOT NULL, + token TEXT NOT NULL, + token_hash VARCHAR(64) NOT NULL, + frequency INTEGER NOT NULL, + PRIMARY KEY (id), + FOREIGN KEY(map_unit_id) REFERENCES demo_document_map_units (id) ON DELETE CASCADE DEFERRABLE INITIALLY DEFERRED +) + +; +CREATE INDEX IF NOT EXISTS idx_demo_document_map_unit_tokens_token_lookup_binary ON demo_document_map_unit_tokens (channel, decode(token_hash, 'hex'::text)) INCLUDE (map_unit_id, token, frequency); +CREATE INDEX IF NOT EXISTS idx_demo_document_map_unit_tokens_unit ON demo_document_map_unit_tokens (map_unit_id, channel) INCLUDE (token, frequency); +ALTER TABLE demo_documents ADD CONSTRAINT fk_demo_documents_current_result FOREIGN KEY (current_job_result_id) REFERENCES job_results(id) ON DELETE RESTRICT; +ALTER TABLE job_results ADD CONSTRAINT fk_job_results_demo_document FOREIGN KEY (demo_document_id) REFERENCES demo_documents(document_id) ON DELETE RESTRICT; +CREATE INDEX IF NOT EXISTS ix_job_results_demo_document_id ON job_results (demo_document_id); +CREATE INDEX IF NOT EXISTS idx_demo_document_chunks_term_trgm ON demo_document_chunks USING gin (lower(COALESCE(term_search_text, '')) gin_trgm_ops); +ALTER TABLE demo_documents ADD CONSTRAINT fk_demo_documents_current_revision FOREIGN KEY (current_job_result_id, document_id) REFERENCES job_results(id, demo_document_id) DEFERRABLE INITIALLY DEFERRED; diff --git a/apps/api/alembic/env.py b/apps/api/alembic/env.py index 021cc5071..d8827c988 100644 --- a/apps/api/alembic/env.py +++ b/apps/api/alembic/env.py @@ -12,7 +12,7 @@ from shared.services.auth.user_table_bootstrap import ensure_better_auth_user_table # Build a synchronous database URL by replacing asyncpg with psycopg2. -sync_database_url = settings.DATABASE_URL.replace("asyncpg", "psycopg2") +sync_database_url = (settings.MIGRATION_DATABASE_URL or settings.DATABASE_URL).replace("asyncpg", "psycopg2") # Read SSL connect args from shared settings. ssl_connect_args = settings.get_ssl_connect_args() diff --git a/apps/api/alembic/versions/4d5e6f7a8b9c_add_shared_demo_corpus.py b/apps/api/alembic/versions/4d5e6f7a8b9c_add_shared_demo_corpus.py new file mode 100644 index 000000000..690bde80b --- /dev/null +++ b/apps/api/alembic/versions/4d5e6f7a8b9c_add_shared_demo_corpus.py @@ -0,0 +1,45 @@ +"""Add dedicated shared demo tables and reserve their namespace.""" + +from __future__ import annotations + +import json +from pathlib import Path + +from alembic import op +import sqlalchemy as sa + +revision: str = "4d5e6f7a8b9c" +down_revision: str = "3c4d5e6f7a8b" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + connection = op.get_bind() + if connection.execute(sa.text("SELECT EXISTS(SELECT 1 FROM documents WHERE namespace = '__knowhere_demo__')")).scalar_one(): + raise RuntimeError("Reserved demo namespace conflicts with existing private documents; stop and resolve the conflict before migration.") + directory = Path(__file__).resolve().parent.parent + # Frozen DDL is independent of future ORM changes. Alembic executes this + # migration once; IF NOT EXISTS also tolerates pre-created indexes/tables. + connection.exec_driver_sql((directory / "demo-corpus-schema.sql").read_text()) + connection.exec_driver_sql((directory / "demo-corpus-policies.sql").read_text()) + entries: list[dict[str, object]] = json.loads((directory / "demo-corpus-catalog.json").read_text()) + for entry in entries: + connection.execute(sa.text(""" + INSERT INTO demo_documents + (document_id, demo_source_id, title, category_id, catalog_metadata, + user_id, namespace, status, source_file_name, parse_track, created_at, updated_at) + VALUES (:document_id, :demo_source_id, :title, :category_id, CAST(:catalog_metadata AS json), + '__knowhere_demo__', '__knowhere_demo__', 'active', :file_name, 'chunk', now(), now()) + ON CONFLICT (demo_source_id) DO NOTHING + """), {**entry, "catalog_metadata": json.dumps(entry["catalog_metadata"])}) + connection.exec_driver_sql(""" + CREATE UNIQUE INDEX IF NOT EXISTS uq_jobs_active_demo_source + ON jobs ((job_metadata ->> 'demo_source_id')) + WHERE job_metadata ->> 'corpus_target' = 'DEMO' + AND status IN ('waiting-file', 'pending', 'running', 'converting') + """) + + +def downgrade() -> None: + raise RuntimeError("Shared demo revisions must be retained. Roll back API/Worker images without dropping the additive schema.") diff --git a/apps/api/app/api/v1/routes/demo.py b/apps/api/app/api/v1/routes/demo.py index 288c27d46..2eec36b9c 100644 --- a/apps/api/app/api/v1/routes/demo.py +++ b/apps/api/app/api/v1/routes/demo.py @@ -1,171 +1,56 @@ -"""Notebook demo document catalog routes.""" +"""Public shared demo directory and immutable revision media routes.""" from __future__ import annotations from typing import Any -from app.api.dependencies.auth import require_write_permission -from app.api.dependencies.current_user import with_current_user -from app.services.demo.document_service import DemoDocumentService -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, field_validator +from fastapi import APIRouter, Depends, Query, Response +from fastapi.responses import RedirectResponse from sqlalchemy.ext.asyncio import AsyncSession +from app.services.demo.shared_document_service import SharedDemoDocumentService +from app.services.documents.lifecycle_service import DocumentService 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 +from shared.core.exceptions.demo_materialization_removed import DemoMaterializationRemovedException router = APIRouter(tags=["Demo Documents"]) - -_demo_service = DemoDocumentService() - - -class DemoMaterializeRequest(BaseModel): - """Request to copy selected canonical demo sources into a 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) +_demo_service: SharedDemoDocumentService = SharedDemoDocumentService() +_document_service: DocumentService = DocumentService() @router.get("/catalog") -async def get_demo_catalog() -> dict[str, Any]: - """Return API-owned canonical demo source metadata and curated Q/A.""" - return _demo_service.get_catalog() +async def get_demo_catalog(response: Response, db: AsyncSession = Depends(get_db)) -> dict[str, Any]: + response.headers["Cache-Control"] = "no-store" + return await _demo_service.get_catalog(db) @router.get("/sources/{demo_source_id}/chunks") -async def list_demo_source_chunks( - demo_source_id: str, - page: int = Query(1, ge=1, description="Page number"), - 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_service.list_chunks( - demo_source_id=demo_source_id, - page=page, - page_size=page_size, - ) - if response is None: - raise _demo_source_not_found(demo_source_id) - return response +async def list_demo_source_chunks(demo_source_id: str, job_result_id: str | None = Query(None, max_length=36), page: int = Query(1, ge=1), page_size: int = Query(50, ge=1, le=200), db: AsyncSession = Depends(get_db)) -> dict[str, Any] | None: + document, revision = await _demo_service.resolve_source(db, source_id=demo_source_id, job_result_id=job_result_id) + response = await _document_service.list_document_chunks(db, user_id="", document_id=document.document_id, job_result_id=revision.id, page=page, page_size=page_size, chunk_type=None, include_asset_urls=True) + return {**(response or {}), "demo_source_id": demo_source_id, "canonical_document_id": document.document_id} @router.get("/sources/{demo_source_id}/chunks/{demo_chunk_id}") -async def get_demo_source_chunk( - demo_source_id: str, - demo_chunk_id: str, -) -> dict[str, Any]: - """Return one canonical demo chunk for citation focusing.""" - response = _demo_service.get_chunk( - demo_source_id=demo_source_id, - demo_chunk_id=demo_chunk_id, - ) +async def get_demo_source_chunk(demo_source_id: str, demo_chunk_id: str, job_result_id: str | None = Query(None, max_length=36), db: AsyncSession = Depends(get_db)) -> dict[str, Any] | None: + document, revision = await _demo_service.resolve_source(db, source_id=demo_source_id, job_result_id=job_result_id) + response = await _document_service.get_document_chunk(db, user_id="", document_id=document.document_id, document_chunk_id=demo_chunk_id, job_result_id=revision.id, include_asset_urls=True) if response is None: - raise NotFoundException( - resource="Demo document chunk", - resource_id=demo_chunk_id, - internal_message="Demo document chunk not found", - ) - return response + from shared.core.exceptions.domain_exceptions import NotFoundException + raise NotFoundException(resource="Demo chunk", resource_id=demo_chunk_id) + return {**response, "demo_source_id": demo_source_id, "canonical_document_id": document.document_id} @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_service.get_original_file_path( - demo_source_id=demo_source_id, - ) - if file_path is None: - raise _demo_source_not_found(demo_source_id) - - return FileResponse( - path=file_path, - media_type="application/pdf", - filename=file_path.name, - content_disposition_type="inline", - ) +async def get_demo_source_original(demo_source_id: str, job_result_id: str | None = Query(None, max_length=36), db: AsyncSession = Depends(get_db)) -> RedirectResponse: + return RedirectResponse(await _demo_service.get_media_url(db, source_id=demo_source_id, job_result_id=job_result_id), status_code=307, headers={"Cache-Control": "no-store"}) @router.get("/sources/{demo_source_id}/assets/{asset_path:path}") -async def get_demo_source_asset( - demo_source_id: str, - asset_path: str, -) -> FileResponse: - """Return a canonical parsed media or table asset for preview.""" - file_path = _demo_service.get_asset_file_path( - demo_source_id=demo_source_id, - asset_path=asset_path, - ) - if file_path is None: - raise _demo_source_not_found(demo_source_id) - - return FileResponse( - path=file_path, - filename=file_path.name, - content_disposition_type="inline", - ) - - -@router.post("/materializations") -async def materialize_demo_sources( - payload: DemoMaterializeRequest, - current_user: CurrentUser = Depends(with_current_user), - _write_permission: None = Depends(require_write_permission), - db: AsyncSession = Depends(get_db), -) -> dict[str, Any]: - """Copy canonical demo sources into the authenticated user's namespace.""" - namespace = normalize_retrieval_namespace(payload.namespace) - try: - materialized_sources = await _demo_service.materialize_sources( - db, - user_id=current_user.user_id, - namespace=namespace, - demo_source_ids=payload.demo_source_ids, - ) - except KeyError as error: - raise _demo_source_not_found(str(error.args[0])) from error - - return { - "namespace": namespace, - "sources": [ - { - "demo_source_id": source.demo_source_id, - "document_id": source.document_id, - "status": source.status, - "title": source.title, - "mime_type": source.mime_type, - "size_bytes": source.size_bytes, - "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, - }, - } - for source in materialized_sources - ], - } - - -def _demo_source_not_found(demo_source_id: str) -> NotFoundException: - return NotFoundException( - resource="Demo document source", - resource_id=demo_source_id, - internal_message="Demo document source not found", - ) +async def get_demo_source_asset(demo_source_id: str, asset_path: str, job_result_id: str | None = Query(None, max_length=36), db: AsyncSession = Depends(get_db)) -> RedirectResponse: + return RedirectResponse(await _demo_service.get_media_url(db, source_id=demo_source_id, job_result_id=job_result_id, asset_path=asset_path), status_code=307, headers={"Cache-Control": "no-store"}) + + +@router.post("/materializations", status_code=410) +async def materialize_demo_sources() -> None: + raise DemoMaterializationRemovedException() diff --git a/apps/api/app/api/v1/routes/documents.py b/apps/api/app/api/v1/routes/documents.py index 16a17a93a..c9c4f2444 100644 --- a/apps/api/app/api/v1/routes/documents.py +++ b/apps/api/app/api/v1/routes/documents.py @@ -63,6 +63,7 @@ async def list_documents( @router.get("/{document_id}") async def get_document( document_id: str, + job_result_id: str | None = Query(None, max_length=36), current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), ): @@ -70,6 +71,7 @@ async def get_document( db, user_id=current_user.user_id, document_id=document_id, + job_result_id=job_result_id, ) if document is None: raise NotFoundException( @@ -83,6 +85,7 @@ async def get_document( @router.get("/{document_id}/chunks") async def list_document_chunks( document_id: str, + job_result_id: str | None = Query(None, max_length=36), page: int = Query(1, ge=1, description="Page number"), page_size: int = Query(50, ge=1, le=200, description="Items per page"), chunk_type: DocumentChunkType | None = Query(None, description="Chunk type filter"), @@ -97,6 +100,7 @@ async def list_document_chunks( db, user_id=current_user.user_id, document_id=document_id, + job_result_id=job_result_id, page=page, page_size=page_size, chunk_type=chunk_type, @@ -115,6 +119,7 @@ async def list_document_chunks( async def get_document_chunk( document_id: str, document_chunk_id: str, + job_result_id: str | None = Query(None, max_length=36), include_asset_urls: bool = Query( False, description="Generate 7-day asset URLs for image/table chunks when true", @@ -126,6 +131,7 @@ async def get_document_chunk( db, user_id=current_user.user_id, document_id=document_id, + job_result_id=job_result_id, document_chunk_id=document_chunk_id, include_asset_urls=include_asset_urls, ) diff --git a/apps/api/app/api/v1/routes/jobs.py b/apps/api/app/api/v1/routes/jobs.py index 25562ecbd..05e4283ac 100644 --- a/apps/api/app/api/v1/routes/jobs.py +++ b/apps/api/app/api/v1/routes/jobs.py @@ -125,3 +125,25 @@ async def confirm_upload( request_payload=request, user_id=current_user.user_id, ) + + +@router.post("/{job_id}/upload-url", summary="Renew a waiting job upload URL") +async def renew_upload_url( + job_id: str, + current_user: CurrentUser = Depends(with_current_user), + _write_permission: None = Depends(require_write_permission), + db: AsyncSession = Depends(get_db), +): + from app.services.jobs.source_upload_service import renew_source_upload + return await renew_source_upload(db, user_id=current_user.user_id, job_id=job_id) + + +@router.delete("/{job_id}", summary="Reject deletion of retained demo revisions") +async def delete_job( + job_id: str, + current_user: CurrentUser = Depends(with_current_user), + _write_permission: None = Depends(require_write_permission), + db: AsyncSession = Depends(get_db), +): + from app.services.jobs.deletion_guard import reject_job_deletion + await reject_job_deletion(db, user_id=current_user.user_id, job_id=job_id) diff --git a/apps/api/app/api/v2/routes/documents.py b/apps/api/app/api/v2/routes/documents.py index 7bde02b51..1abc3b815 100644 --- a/apps/api/app/api/v2/routes/documents.py +++ b/apps/api/app/api/v2/routes/documents.py @@ -7,7 +7,7 @@ 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 +from fastapi import APIRouter, Depends, Query from sqlalchemy.ext.asyncio import AsyncSession from shared.core.database import get_db @@ -21,6 +21,7 @@ @router.get("/{document_id}/files/page-citation-source") async def get_document_page_citation_source( document_id: str, + job_result_id: str | None = Query(None, max_length=36), current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), ) -> dict[str, Any]: @@ -28,6 +29,7 @@ async def get_document_page_citation_source( db, user_id=current_user.user_id, document_id=document_id, + job_result_id=job_result_id, ) if response is None: raise NotFoundException( diff --git a/apps/api/app/api/v2/routes/jobs.py b/apps/api/app/api/v2/routes/jobs.py index c9116bc05..0baa200a4 100644 --- a/apps/api/app/api/v2/routes/jobs.py +++ b/apps/api/app/api/v2/routes/jobs.py @@ -112,3 +112,25 @@ async def confirm_upload( request_payload=request, user_id=current_user.user_id, ) + + +@router.post("/{job_id}/upload-url", summary="Renew a waiting job upload URL") +async def renew_upload_url( + job_id: str, + current_user: CurrentUser = Depends(with_current_user), + _write_permission: None = Depends(require_write_permission), + db: AsyncSession = Depends(get_db), +): + from app.services.jobs.source_upload_service import renew_source_upload + return await renew_source_upload(db, user_id=current_user.user_id, job_id=job_id) + + +@router.delete("/{job_id}", summary="Reject deletion of retained demo revisions") +async def delete_job( + job_id: str, + current_user: CurrentUser = Depends(with_current_user), + _write_permission: None = Depends(require_write_permission), + db: AsyncSession = Depends(get_db), +): + from app.services.jobs.deletion_guard import reject_job_deletion + await reject_job_deletion(db, user_id=current_user.user_id, job_id=job_id) diff --git a/apps/api/app/repositories/document_repository.py b/apps/api/app/repositories/document_repository.py index 865f12304..110a96fad 100644 --- a/apps/api/app/repositories/document_repository.py +++ b/apps/api/app/repositories/document_repository.py @@ -4,11 +4,15 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage +from app.services.demo.revision_reader import resolve_demo_revision + from datetime import datetime, timezone from typing import Sequence, cast -from sqlalchemy import func, select +from sqlalchemy import func, select, true from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm import noload from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job import Job @@ -28,12 +32,13 @@ async def list_by_user_namespace( limit: int, offset: int, ) -> Sequence[Document]: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) result = await db.execute( - select(Document) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status != "archived") - .order_by(Document.updated_at.desc(), Document.document_id.asc()) + select(corpusStorage.Document) + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.Document.namespace == namespace) + .where(corpusStorage.Document.status != "archived") + .order_by(corpusStorage.Document.updated_at.desc(), corpusStorage.Document.document_id.asc()) .limit(limit) .offset(offset) ) @@ -46,11 +51,12 @@ async def count_by_user_namespace( user_id: str, namespace: str, ) -> int: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) result = await db.execute( - select(func.count(Document.document_id)) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status != "archived") + select(func.count(corpusStorage.Document.document_id)) + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.Document.namespace == namespace) + .where(corpusStorage.Document.status != "archived") ) return int(result.scalar_one()) @@ -61,10 +67,12 @@ async def get_document( document_id: str, user_id: str, ) -> Document | None: + corpusStorage: CorpusStorage = CorpusStorage.resolve_document(document_id) result = await db.execute( - select(Document) - .where(Document.document_id == document_id) - .where(Document.user_id == user_id) + select(corpusStorage.Document) + .where(corpusStorage.Document.document_id == document_id) + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.Document.status != "archived" if corpusStorage.is_demo else true()) ) return result.scalar_one_or_none() @@ -74,16 +82,22 @@ async def get_current_document_job_revision( *, document_id: str, user_id: str, + job_result_id: str | None = None, ) -> DocumentJobRevisionRow | None: + corpusStorage: CorpusStorage = CorpusStorage.resolve_document(document_id) + if corpusStorage.is_demo: + document, result = await resolve_demo_revision(db, document_id=document_id, job_result_id=job_result_id) + job = (await db.execute(select(Job).where(Job.job_id == result.job_id))).scalar_one() + return document, result, job stmt = ( - select(Document, JobResult, Job) - .join(JobResult, JobResult.id == Document.current_job_result_id) + select(corpusStorage.Document, JobResult, Job) + .join(JobResult, JobResult.id == (job_result_id or corpusStorage.Document.current_job_result_id)) .join(Job, Job.job_id == JobResult.job_id) - .where(Document.document_id == document_id) - .where(Document.user_id == user_id) - .where(JobResult.document_id == Document.document_id) + .where(corpusStorage.Document.document_id == document_id) + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(user_id)) + .where(JobResult.document_id == corpusStorage.Document.document_id) .where(Job.user_id == user_id) - .where(Document.status != "archived") + .where(corpusStorage.Document.status != "archived") .limit(1) ) @@ -109,13 +123,14 @@ async def count_current_document_chunks( job_result_id: str, chunk_type: str | None = None, ) -> int: + corpusStorage: CorpusStorage = CorpusStorage.resolve_document(document_id) stmt = ( - select(func.count(DocumentChunk.id)) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) + select(func.count(corpusStorage.DocumentChunk.id)) + .where(corpusStorage.DocumentChunk.document_id == document_id) + .where(corpusStorage.DocumentChunk.job_result_id == job_result_id) ) if chunk_type is not None: - stmt = stmt.where(func.lower(DocumentChunk.chunk_type) == chunk_type) + stmt = stmt.where(func.lower(corpusStorage.DocumentChunk.chunk_type) == chunk_type) result = await db.execute(stmt) return int(result.scalar_one()) @@ -130,25 +145,26 @@ async def list_current_document_chunks( offset: int, chunk_type: str | None = None, ) -> Sequence[DocumentChunkRow]: + corpusStorage: CorpusStorage = CorpusStorage.resolve_document(document_id) stmt = ( - select(DocumentChunk, DocumentSection, JobResult) + select(corpusStorage.DocumentChunk, corpusStorage.DocumentSection, JobResult).options(noload(JobResult.chunks), noload(JobResult.job)) .outerjoin( - DocumentSection, - DocumentSection.section_id == DocumentChunk.section_id, + corpusStorage.DocumentSection, + corpusStorage.DocumentSection.section_id == corpusStorage.DocumentChunk.section_id, ) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) + .join(JobResult, JobResult.id == corpusStorage.DocumentChunk.job_result_id) + .where(corpusStorage.DocumentChunk.document_id == document_id) + .where(corpusStorage.DocumentChunk.job_result_id == job_result_id) .order_by( - DocumentChunk.sort_order.asc(), - DocumentChunk.created_at.asc(), - DocumentChunk.id.asc(), + corpusStorage.DocumentChunk.sort_order.asc(), + corpusStorage.DocumentChunk.created_at.asc(), + corpusStorage.DocumentChunk.id.asc(), ) .limit(limit) .offset(offset) ) if chunk_type is not None: - stmt = stmt.where(func.lower(DocumentChunk.chunk_type) == chunk_type) + stmt = stmt.where(func.lower(corpusStorage.DocumentChunk.chunk_type) == chunk_type) result = await db.execute(stmt) return cast(Sequence[DocumentChunkRow], result.all()) @@ -161,16 +177,17 @@ async def get_current_document_chunk( job_result_id: str, document_chunk_id: str, ) -> DocumentChunkRow | None: + corpusStorage: CorpusStorage = CorpusStorage.resolve_document(document_id) stmt = ( - select(DocumentChunk, DocumentSection, JobResult) + select(corpusStorage.DocumentChunk, corpusStorage.DocumentSection, JobResult).options(noload(JobResult.chunks), noload(JobResult.job)) .outerjoin( - DocumentSection, - DocumentSection.section_id == DocumentChunk.section_id, + corpusStorage.DocumentSection, + corpusStorage.DocumentSection.section_id == corpusStorage.DocumentChunk.section_id, ) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.id == document_chunk_id) + .join(JobResult, JobResult.id == corpusStorage.DocumentChunk.job_result_id) + .where(corpusStorage.DocumentChunk.document_id == document_id) + .where(corpusStorage.DocumentChunk.job_result_id == job_result_id) + .where(corpusStorage.DocumentChunk.id == document_chunk_id) .limit(1) ) diff --git a/apps/api/app/repositories/job_result_repository.py b/apps/api/app/repositories/job_result_repository.py index 58abeac7d..f9b4c947d 100644 --- a/apps/api/app/repositories/job_result_repository.py +++ b/apps/api/app/repositories/job_result_repository.py @@ -98,5 +98,8 @@ async def delete_by_job_id(self, db: AsyncSession, job_id: str) -> None: result = await self.get_by_job_id(db, job_id) if not result: return + if result.demo_document_id is not None: + from shared.core.exceptions.domain_exceptions import ConflictException + raise ConflictException(user_message="Shared demo revision results and their assets must be retained.", resource="Demo revision", resource_id=result.id) await db.delete(result) await db.flush() 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 f0a63bd1c..000000000 --- a/apps/api/app/services/demo/document_service.py +++ /dev/null @@ -1,87 +0,0 @@ -"""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/revision_reader.py b/apps/api/app/services/demo/revision_reader.py new file mode 100644 index 000000000..b31b24979 --- /dev/null +++ b/apps/api/app/services/demo/revision_reader.py @@ -0,0 +1,44 @@ +"""Resolve only complete demo revisions; never serve staged parser payloads.""" + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm import noload + +from shared.core.exceptions.demo_not_ready import DemoNotReadyException +from shared.core.exceptions.domain_exceptions import NotFoundException +from shared.models.database.demo_corpus import DemoDocument, DemoDocumentMapUnitIndex +from shared.models.database.job_result import JobResult + + +async def resolve_demo_revision( + db: AsyncSession, *, document_id: str, job_result_id: str | None = None +) -> tuple[DemoDocument, JobResult]: + document: DemoDocument | None = ( + await db.execute( + select(DemoDocument).where( + DemoDocument.document_id == document_id, DemoDocument.status == "active" + ) + ) + ).scalar_one_or_none() + if document is None: + raise NotFoundException(resource="Demo document", resource_id=document_id) + revisionId: str | None = job_result_id or document.current_job_result_id + if revisionId is None: + raise DemoNotReadyException(is_explicit_source=True) + result: JobResult | None = ( + await db.execute( + select(JobResult) + .options(noload(JobResult.chunks), noload(JobResult.job)) + .join( + DemoDocumentMapUnitIndex, + (DemoDocumentMapUnitIndex.job_result_id == JobResult.id) + & (DemoDocumentMapUnitIndex.document_id == document_id), + ) + .where( + JobResult.id == revisionId, JobResult.demo_document_id == document_id + ) + ) + ).scalar_one_or_none() + if result is None: + raise NotFoundException(resource="Demo revision", resource_id=revisionId) + return document, result diff --git a/apps/api/app/services/demo/shared_document_service.py b/apps/api/app/services/demo/shared_document_service.py new file mode 100644 index 000000000..eb6699c82 --- /dev/null +++ b/apps/api/app/services/demo/shared_document_service.py @@ -0,0 +1,223 @@ +"""Directory and media reads for the single shared, revisioned demo corpus.""" + +from __future__ import annotations + +import asyncio +import mimetypes +from typing import Any +from urllib.parse import quote + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm import noload + +from app.services.demo.official_library_catalog import get_official_library_catalog +from app.services.demo.revision_reader import resolve_demo_revision +from shared.core.exceptions.domain_exceptions import NotFoundException +from shared.models.database.demo_corpus import ( + DemoDocument, + DemoDocumentChunk, + DemoDocumentMapUnitIndex, +) +from shared.models.database.job import Job +from shared.models.database.job_result import JobResult +from shared.services.retrieval.corpus_storage import DEMO_NAMESPACE +from shared.services.storage.job_file_storage import JobFileStorage +from shared.services.storage.demo_asset_signer import DemoAssetSigner + + +class SharedDemoDocumentService: + async def get_catalog(self, db: AsyncSession) -> dict[str, Any]: + sourceExpression = Job.job_metadata["demo_source_id"].as_string() + latestJobs = ( + select(sourceExpression.label("source_id"), Job.status) + .where(Job.job_metadata["corpus_target"].as_string() == "DEMO") + .distinct(sourceExpression) + .order_by(sourceExpression, Job.created_at.desc()) + .subquery() + ) + counts = ( + select( + DemoDocumentChunk.document_id, + DemoDocumentChunk.job_result_id, + func.count(DemoDocumentChunk.id).label("chunk_count"), + ) + .group_by(DemoDocumentChunk.document_id, DemoDocumentChunk.job_result_id) + .subquery() + ) + catalogRows = ( + await db.execute( + select( + DemoDocument, + JobResult, + latestJobs.c.status, + DemoDocumentMapUnitIndex.id, + counts.c.chunk_count, + ) + .options(noload(JobResult.chunks), noload(JobResult.job)) + .outerjoin( + JobResult, JobResult.id == DemoDocument.current_job_result_id + ) + .outerjoin( + latestJobs, latestJobs.c.source_id == DemoDocument.demo_source_id + ) + .outerjoin( + DemoDocumentMapUnitIndex, + (DemoDocumentMapUnitIndex.document_id == DemoDocument.document_id) + & ( + DemoDocumentMapUnitIndex.job_result_id + == DemoDocument.current_job_result_id + ), + ) + .outerjoin( + counts, + (counts.c.document_id == DemoDocument.document_id) + & (counts.c.job_result_id == DemoDocument.current_job_result_id), + ) + .where(DemoDocument.status == "active") + .order_by(DemoDocument.title) + ) + ).all() + readySources: list[dict[str, Any]] = [] + importStatuses: list[dict[str, str]] = [] + for document, revision, latestStatus, marker, count in catalogRows: + state: str = ( + "planned" + if latestStatus is None + else "processing" + if latestStatus in ("waiting-file", "pending", "running", "converting") + else "failed" + if latestStatus == "failed" + else "ready" + if document.current_job_result_id + else "planned" + ) + importStatuses.append( + {"demo_source_id": document.demo_source_id, "status": state} + ) + if revision is None or marker is None: + continue + metadata: dict[str, Any] = revision.document_metadata or {} + mimeType: str = ( + mimetypes.guess_type(document.source_file_name or document.title)[0] + or "application/octet-stream" + ) + source: dict[str, Any] = { + "demo_source_id": document.demo_source_id, + "canonical_document_id": document.document_id, + "job_result_id": revision.id, + "namespace": DEMO_NAMESPACE, + "title": document.title, + "category_id": document.category_id, + "mime_type": mimeType, + "size_bytes": int(metadata.get("source_size_bytes") or 0), + "chunk_count": int(count or 0), + "import_status": state, + "original_file": { + "url": f"/api/v1/demo/sources/{quote(document.demo_source_id, safe='')}/original?job_result_id={revision.id}", + "mime_type": mimeType, + "size_bytes": int(metadata.get("source_size_bytes") or 0), + "can_download": True, + }, + "examples": metadata.get("bound_demo_examples", []), + } + readySources.append(source) + library: dict[str, Any] = get_official_library_catalog( + {str(source["demo_source_id"]): source for source in readySources} + ) + statuses: dict[str, str] = { + entry["demo_source_id"]: entry["status"] for entry in importStatuses + } + for source in library.get("sources", []): + if source.get("demo_source_id") not in { + item["demo_source_id"] for item in readySources + }: + source["status"] = "planned" + source["import_status"] = statuses.get( + str(source.get("demo_source_id") or ""), "planned" + ) + readySource = next( + ( + item + for item in readySources + if item["demo_source_id"] == source.get("demo_source_id") + ), + None, + ) + if readySource is not None: + source["job_result_id"] = readySource["job_result_id"] + source["namespace"] = DEMO_NAMESPACE + return { + "namespace": DEMO_NAMESPACE, + "sources": readySources, + "import_status": importStatuses, + "official_library": library, + } + + async def resolve_source( + self, db: AsyncSession, *, source_id: str, job_result_id: str | None = None + ) -> tuple[DemoDocument, JobResult]: + document: DemoDocument | None = ( + await db.execute( + select(DemoDocument).where( + DemoDocument.demo_source_id == source_id, + DemoDocument.status == "active", + ) + ) + ).scalar_one_or_none() + if document is None: + raise NotFoundException(resource="Demo source", resource_id=source_id) + return await resolve_demo_revision( + db, document_id=document.document_id, job_result_id=job_result_id + ) + + async def get_media_url( + self, + db: AsyncSession, + *, + source_id: str, + job_result_id: str | None = None, + asset_path: str | None = None, + ) -> str: + document, revision = await self.resolve_source( + db, source_id=source_id, job_result_id=job_result_id + ) + job: Job = ( + await db.execute(select(Job).where(Job.job_id == revision.job_id)) + ).scalar_one() + storage: JobFileStorage = JobFileStorage() + if asset_path is not None: + signer: DemoAssetSigner = DemoAssetSigner( + job.job_id, revision.document_metadata or {} + ) + assetUrl: str | None = await asyncio.to_thread( + signer.generate_url, asset_path + ) + if assetUrl is None: + raise NotFoundException(resource="Demo asset", resource_id=asset_path) + return assetUrl + if asset_path is None: + key: str | None = job.s3_key + bucket: str = storage.uploads_bucket + fileName: str = str( + (job.job_metadata or {}).get("source_file_name") + or document.source_file_name + or document.title + ) + if not key: + raise NotFoundException(resource="Demo original", resource_id=source_id) + mimeType: str = mimetypes.guess_type(fileName)[0] or "application/octet-stream" + # Redirects let the storage service handle Range without proxying large + # original files through an API worker. The GET signature overrides MIME. + return await asyncio.to_thread( + storage.storage_adapter.generate_presigned_url, + key, + expiration=3600, + bucket=bucket, + method="GET", + headers={ + "Content-Type": mimeType, + "Content-Disposition": "inline; filename*=UTF-8''" + + quote(fileName, safe=""), + }, + ) diff --git a/apps/api/app/services/demo/source_materializer.py b/apps/api/app/services/demo/source_materializer.py deleted file mode 100644 index 022746046..000000000 --- a/apps/api/app/services/demo/source_materializer.py +++ /dev/null @@ -1,555 +0,0 @@ -"""Demo Source Materialization workflow.""" - -from __future__ import annotations - -import asyncio -import shutil -import tempfile -import time -from collections.abc import Iterable -from contextlib import nullcontext -from dataclasses import dataclass -from datetime import datetime, timedelta, timezone -from hashlib import blake2b -from pathlib import Path -from uuid import uuid4 - -import logfire - -from app.services.demo.source_catalog import DemoSourceCatalog, DemoSourceDefinition -from app.services.demo.canonical_bundle import CanonicalDemoBundleStore -from app.services.demo.canonical_bundle_result import CanonicalDemoBundle -from sqlalchemy.exc import IntegrityError -from sqlalchemy import delete, func, select -from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession - -from shared.core.exceptions.domain_exceptions import ConflictException, ValidationException -from shared.models.database.demo_materialization import DemoMaterialization -from shared.models.database.job import Job -from shared.models.database.job_result import JobResult -from shared.services.jobs.lifecycle.publication_trace import PublicationTrace -from shared.services.jobs.lifecycle.publication_trace_sql import ( - attach_publication_trace, - clear_publication_trace, - install_publication_trace_sql_instrumentation, -) -from shared.services.retrieval.cache_service import invalidate_retrieval_cache_namespaces -from shared.services.retrieval.publication_service import RetrievalPublicationService -from shared.services.retrieval.publication_models import DocumentPublicationScope -from shared.services.retrieval.publication_preparation_cache import PublicationPreparationCache -from shared.services.retrieval.scoring.knowhere_hybrid import MAP_UNIT_INDEX_FORMAT_VERSION -from shared.services.redis import RedisPublicationSemaphore, RedisService, RedisServiceFactory -from shared.core.config import settings -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() - self._redis_service = RedisServiceFactory.get_service() - - 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 - ] - try: - claims = await self._claim_sources( - db, - user_id=user_id, - namespace=namespace, - sources=selected_sources, - ) - except IntegrityError as error: - await db.rollback() - raise ConflictException( - user_message="This demo source is currently being materialized.", - reason="ABORTED", - resource="Demo materialization", - ) from error - await db.commit() - claim_ids = tuple(claim.id for claim in claims.values()) - results: list[MaterializedDemoSource] = [] - for source in selected_sources: - try: - semaphore = RedisPublicationSemaphore( - self._redis_service, - concurrency=settings.MATERIALIZATION_DB_PUBLICATION_CONCURRENCY, - lease_seconds=settings.MATERIALIZATION_DB_PUBLICATION_LEASE_SECONDS, - acquire_timeout_seconds=settings.MATERIALIZATION_DB_PUBLICATION_ACQUIRE_TIMEOUT_SECONDS, - ) - result = await self._materialize_source( - db, - user_id=user_id, - namespace=namespace, - source=source, - claim=claims[source.demo_source_id], - semaphore=semaphore, - ) - results.append(result) - except Exception: - await db.rollback() - await self._release_claims( - db, - claim_ids=claim_ids, - ) - raise - 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, - claim: DemoMaterialization, - semaphore: RedisPublicationSemaphore, - ) -> MaterializedDemoSource: - document_id = f"doc_{uuid4().hex[:12]}" - job_id = f"job_demo_{uuid4().hex[:12]}" - job_result_id = str(uuid4()) - timestamp = _utc_now() - if settings.KNOWHERE_PUBLICATION_TRACE_ENABLED: - bind = db.bind - if bind is None: - raise RuntimeError("Demo publication database session has no bind") - install_publication_trace_sql_instrumentation( - bind.engine if isinstance(bind, AsyncConnection) else bind - ) - stage_started_at = time.perf_counter() - result_bundle = await _upload_demo_result_bundle( - job_id=job_id, - demo_source_id=source.demo_source_id, - source_directory=self._catalog.source_directory(source), - redis_service=self._redis_service, - ) - logfire.info( - "Demo materialization source bundle upload completed", - demo_source_id=source.demo_source_id, - duration_seconds=time.perf_counter() - stage_started_at, - reused=result_bundle.reused, - content_version=result_bundle.content_version, - ) - - 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, - **( - {"result_raw_prefix": result_bundle.raw_prefix} - if result_bundle.content_version - else {} - ), - }, - inline_payload={"source": "canonical_demo"}, - result_s3_key=result_bundle.zip_key, - result_size=result_bundle.zip_size, - created_at=timestamp, - updated_at=timestamp, - ) - ) - chunks = self._catalog.publication_chunks(source) - stage_started_at = time.perf_counter() - wait_seconds = await semaphore.acquire() - logfire.info( - "Demo materialization publication semaphore acquired", - demo_source_id=source.demo_source_id, - wait_seconds=wait_seconds, - ) - trace: PublicationTrace | None = ( - PublicationTrace.start( - attempt_ref=f"attempt_{uuid4().hex}", - owner="async", - job_id=job_id, - job_result_id=job_result_id, - document_id=document_id, - scope_fingerprint=_publication_scope_fingerprint( - user_id=user_id, - namespace=namespace, - ), - job_type="demo_materialization", - parse_track="chunk", - ) - if settings.KNOWHERE_PUBLICATION_TRACE_ENABLED - else None - ) - trace_connection: AsyncConnection | None = None - try: - if trace is not None: - for chunk_type in ("text", "image", "table", "page"): - trace.record_count( - f"input_{chunk_type}_chunks", - sum( - 1 - for chunk in chunks - if str(chunk.get("type") or chunk.get("chunk_type") or "text") - == chunk_type - ), - ) - checkout_started_at = time.perf_counter() - trace_connection = await db.connection() - trace.record_stage( - "pool_checkout_wait", - (time.perf_counter() - checkout_started_at) * 1000, - ) - attach_publication_trace(trace_connection, trace) - - base_rows_started_at = time.perf_counter() - with trace.stage("result_binding") if trace is not None else nullcontext(): - await db.flush() - logfire.info( - "Demo materialization base rows completed", - demo_source_id=source.demo_source_id, - duration_seconds=time.perf_counter() - base_rows_started_at, - ) - with ( - PublicationPreparationCache.scope( - source_id=source.demo_source_id, - content_version=result_bundle.content_version, - index_format_version=MAP_UNIT_INDEX_FORMAT_VERSION, - ) - if settings.DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED - and result_bundle.content_version - else nullcontext() - ): - published_state = 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], - update_namespace_snapshot=False, - trace=trace, - ) - ) - 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, - trace=trace, - ) - ) - await db.flush() - logfire.info( - "Demo materialization sections chunks and map index completed", - demo_source_id=source.demo_source_id, - duration_seconds=time.perf_counter() - stage_started_at, - chunk_count=len(chunks), - ) - - if published_state is None or published_state.document_id != document_id: - raise RuntimeError("Demo publication did not create its requested document") - if published_state.manifest_payload is None: - raise RuntimeError("Demo publication did not create a serving manifest") - manifest_payload = published_state.manifest_payload - stage_started_at = time.perf_counter() - await db.run_sync( - lambda sync_db: self._publication_service.update_namespace_snapshot( - sync_db, - scope=DocumentPublicationScope( - user_id=user_id, - namespace=namespace, - document_id=document_id, - job_result_id=job_result_id, - source_file_name=source.title, - ), - manifest_payload=manifest_payload, - previous_namespace=published_state.previous_namespace, - trace=trace, - ) - ) - logfire.info( - "Demo materialization namespace snapshot completed", - demo_source_id=source.demo_source_id, - duration_seconds=time.perf_counter() - stage_started_at, - ) - claim.document_id = document_id - claim.status = "ready" - claim.claimed_at = None - claim.updated_at = _utc_now() - commit_started_at = time.perf_counter() - with trace.stage("commit") if trace is not None else nullcontext(): - await db.commit() - logfire.info( - "Demo materialization publication commit completed", - demo_source_id=source.demo_source_id, - duration_seconds=time.perf_counter() - commit_started_at, - ) - except Exception: - if trace is not None: - try: - with trace.stage("rollback"): - await db.rollback() - finally: - trace.finish(outcome="rollback") - raise - else: - if trace is not None: - trace.finish(outcome="success") - finally: - if trace_connection is not None: - clear_publication_trace(trace_connection) - await semaphore.release() - - return _materialized_source_payload( - source=source, - document_id=document_id, - status="created", - ) - - async def _claim_sources( - self, - db: AsyncSession, - *, - user_id: str, - namespace: str, - sources: list[DemoSourceDefinition], - ) -> dict[str, DemoMaterialization]: - now = _utc_now() - claims: dict[str, DemoMaterialization] = {} - for source in sorted(sources, key=lambda item: item.demo_source_id): - lock_acquired = await db.scalar( - select( - func.pg_try_advisory_xact_lock( - _materialization_lock_id( - user_id=user_id, - namespace=namespace, - demo_source_id=source.demo_source_id, - ) - ) - ) - ) - if not lock_acquired: - raise ConflictException( - user_message="This demo source is currently being materialized.", - reason="ABORTED", - resource="Demo materialization", - resource_id=source.demo_source_id, - ) - result = await db.execute( - select(DemoMaterialization) - .where(DemoMaterialization.user_id == user_id) - .where(DemoMaterialization.namespace == namespace) - .where(DemoMaterialization.demo_source_id == source.demo_source_id) - ) - existing = result.scalar_one_or_none() - if existing is not None: - if not _is_stale_claim(existing, now=now): - raise _materialization_conflict(existing, source) - claim = existing - claim.status = "materializing" - claim.document_id = None - claim.claimed_at = now - claim.updated_at = now - else: - claim = DemoMaterialization( - user_id=user_id, - namespace=namespace, - demo_source_id=source.demo_source_id, - status="materializing", - document_id=None, - claimed_at=now, - created_at=now, - updated_at=now, - ) - db.add(claim) - claims[source.demo_source_id] = claim - await db.flush() - return claims - - async def _release_claims( - self, - db: AsyncSession, - *, - claim_ids: Iterable[str], - ) -> None: - claim_ids = tuple(claim_ids) - if not claim_ids: - return - await db.execute( - delete(DemoMaterialization).where( - DemoMaterialization.status == "materializing", - DemoMaterialization.id.in_(claim_ids), - ) - ) - await db.commit() - - -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 - - -def _is_stale_claim(materialization: DemoMaterialization, *, now: datetime) -> bool: - return ( - materialization.status == "materializing" - and materialization.claimed_at is not None - and materialization.claimed_at < now - timedelta(minutes=10) - ) - - -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 _publication_scope_fingerprint(*, user_id: str, namespace: str) -> str: - scope = f"{user_id}\0{namespace}".encode("utf-8") - return f"scope-{blake2b(scope, digest_size=16).hexdigest()}" - - -def _materialization_conflict( - materialization: DemoMaterialization, - source: DemoSourceDefinition, -) -> ConflictException: - reason = "ALREADY_EXISTS" if materialization.status == "ready" else "ABORTED" - return ConflictException( - user_message=( - "This demo source has already been materialized." - if reason == "ALREADY_EXISTS" - else "This demo source is currently being materialized." - ), - reason=reason, - resource="Demo materialization", - resource_id=source.demo_source_id, - ) - - -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, - ) - - -async def _upload_demo_result_bundle( - *, - job_id: str, - demo_source_id: str, - source_directory: Path, - redis_service: RedisService, -) -> CanonicalDemoBundle: - if not settings.DEMO_CANONICAL_BUNDLE_ENABLED: - return await asyncio.to_thread( - _upload_job_result_bundle, job_id=job_id, source_directory=source_directory - ) - return await CanonicalDemoBundleStore(redis_service=redis_service).ensure_bundle( - source_id=demo_source_id, source_directory=source_directory - ) - - -def _upload_job_result_bundle( - *, job_id: str, source_directory: Path -) -> CanonicalDemoBundle: - with tempfile.TemporaryDirectory(prefix="knowhere-demo-result-") as temporary: - archive_path: Path = Path(shutil.make_archive( - str(Path(temporary) / job_id), "zip", root_dir=source_directory - )) - zip_size: int = archive_path.stat().st_size - bundle = get_result_storage().upload( - job_id=job_id, result_dir=str(source_directory), - zip_file_path=str(archive_path), - ) - return CanonicalDemoBundle( - zip_key=bundle.zip_key, zip_size=zip_size, raw_prefix=bundle.raw_prefix, - content_version="", reused=False, - ) - - -def _utc_now() -> datetime: - return datetime.now(timezone.utc).replace(tzinfo=None) diff --git a/apps/api/app/services/document_ingestion/confirmation_service.py b/apps/api/app/services/document_ingestion/confirmation_service.py index 4cc3a7b8e..4470fe679 100644 --- a/apps/api/app/services/document_ingestion/confirmation_service.py +++ b/apps/api/app/services/document_ingestion/confirmation_service.py @@ -43,6 +43,9 @@ async def confirm_upload( check_job_permission(job, user_id, job_id) assert job is not None + from shared.services.retrieval.demo_job_scope import validate_job_corpus + await db.run_sync(lambda session: validate_job_corpus(session, job=job)) + 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}") diff --git a/apps/api/app/services/document_ingestion/creation_service.py b/apps/api/app/services/document_ingestion/creation_service.py index 5304a9140..79e2227ca 100644 --- a/apps/api/app/services/document_ingestion/creation_service.py +++ b/apps/api/app/services/document_ingestion/creation_service.py @@ -156,7 +156,7 @@ async def _reject_duplicate_source_file_name( scope: ResolvedDocumentIngestionScope, source_file_name: str, ) -> None: - if payload.document_id: + if payload.document_id or scope.job_metadata.get("corpus_target") == "DEMO": return existing = await find_active_document_by_source_file_name( db, diff --git a/apps/api/app/services/document_ingestion/demo_scope.py b/apps/api/app/services/document_ingestion/demo_scope.py new file mode 100644 index 000000000..58d8fad79 --- /dev/null +++ b/apps/api/app/services/document_ingestion/demo_scope.py @@ -0,0 +1,95 @@ +"""Claim a stable shared source before creating its ordinary ingestion job.""" + +from __future__ import annotations + +from uuid import uuid5, NAMESPACE_URL + +from sqlalchemy import select, text +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + NotFoundException, + ValidationException, +) +from shared.models.database.demo_corpus import DemoDocument +from shared.models.database.job import Job +from shared.models.schemas.job import JobCreateBase +from shared.services.retrieval.corpus_storage import DEMO_NAMESPACE, DEMO_OWNER +from shared.services.retrieval.demo_authorization import authorize_demo_transaction +from app.services.document_ingestion.scope_service import ( + raise_document_ingestion_conflict, +) + + +async def resolve_demo_scope( + db: AsyncSession, *, user_id: str, payload: JobCreateBase +) -> str: + await db.run_sync( + lambda session: authorize_demo_transaction(session, user_id=user_id) + ) + sourceId: str = str(payload.data_id or "").strip() + if not sourceId or len(sourceId) > 128: + raise ValidationException( + user_message="Demo ingestion requires data_id with at most 128 characters.", + violations=[{"field": "data_id", "description": "Required demo source ID"}], + ) + # This transaction lock spans source creation and the JobRepository commit. + # The partial unique index protects claims from any other API process too. + await db.execute( + text("SELECT pg_advisory_xact_lock(hashtextextended(:source, 0))"), + {"source": "demo:" + sourceId}, + ) + document = ( + await db.execute( + select(DemoDocument) + .where(DemoDocument.demo_source_id == sourceId) + .with_for_update() + ) + ).scalar_one_or_none() + if document is not None and document.status == "archived": + raise NotFoundException(resource="Demo source", resource_id=sourceId) + if payload.document_id and ( + document is None or payload.document_id != document.document_id + ): + raise ValidationException( + user_message="document_id must belong to this demo data_id.", + violations=[ + {"field": "document_id", "description": "Demo source mismatch"} + ], + ) + if document is None: + fileName: str = str(payload.file_name or sourceId) + document = DemoDocument( + document_id="ddoc_" + uuid5(NAMESPACE_URL, sourceId).hex[:12], + demo_source_id=sourceId, + title=fileName, + category_id="other", + catalog_metadata={}, + user_id=DEMO_OWNER, + namespace=DEMO_NAMESPACE, + status="active", + source_file_name=fileName, + parse_track="chunk", + ) + db.add(document) + await db.flush() + activeJob: Job | None = ( + ( + await db.execute( + select(Job).where( + Job.job_metadata["corpus_target"].as_string() == "DEMO", + Job.job_metadata["demo_source_id"].as_string() == sourceId, + Job.status.in_( + ("waiting-file", "pending", "running", "converting") + ), + ) + ) + ) + .scalars() + .first() + ) + if activeJob is not None: + raise_document_ingestion_conflict( + document_id=document.document_id, active_job_id=activeJob.job_id + ) + return document.document_id diff --git a/apps/api/app/services/document_ingestion/scope_service.py b/apps/api/app/services/document_ingestion/scope_service.py index 7682bebb3..842bade73 100644 --- a/apps/api/app/services/document_ingestion/scope_service.py +++ b/apps/api/app/services/document_ingestion/scope_service.py @@ -50,7 +50,7 @@ async def find_active_job_for_document( def is_active_document_job_unique_violation(exc: Exception) -> bool: text = f"{exc} {getattr(exc, 'orig', '')}" - return "uq_jobs_user_active_document" in text + return "uq_jobs_user_active_document" in text or "uq_jobs_active_demo_source" in text def raise_document_ingestion_conflict( diff --git a/apps/api/app/services/document_ingestion/service.py b/apps/api/app/services/document_ingestion/service.py index a0f049c29..fc2a0e6dc 100644 --- a/apps/api/app/services/document_ingestion/service.py +++ b/apps/api/app/services/document_ingestion/service.py @@ -29,6 +29,9 @@ from sqlalchemy.ext.asyncio import AsyncSession from shared.core.config import settings +from shared.services.retrieval.corpus_storage import CorpusStorage, DEMO_NAMESPACE +from shared.services.retrieval.demo_authorization import require_demo_maintainer +from app.services.document_ingestion.demo_scope import resolve_demo_scope from shared.core.exceptions.domain_exceptions import ( ConflictException, JobOperationException, @@ -111,6 +114,8 @@ async def _create_job_with_command_factory( build_command: IngestionCommandFactory, ) -> JobResponse: try: + if CorpusStorage.resolve_namespace(payload.namespace).is_demo or str(payload.document_id or "").startswith("ddoc_"): + require_demo_maintainer(current_user.user_id) job_id = f"job_{uuid.uuid4().hex[:12]}" file_extension = await self._validate_create_payload( payload, @@ -135,6 +140,8 @@ async def _create_job_with_command_factory( current_user=current_user, scope=scope, ) + except PermissionDeniedException: + raise except NotFoundException: raise except ValidationException: @@ -279,6 +286,15 @@ async def _resolve_scope( page_memory_config=command.page_memory_config, ), ) + if CorpusStorage.resolve_namespace(payload.namespace).is_demo: + documentId: str = await resolve_demo_scope(db, user_id=current_user.user_id, payload=payload) + JobMetadataHelper.set_document_scope(job_metadata, document_id=documentId, namespace=DEMO_NAMESPACE) + job_metadata["corpus_target"] = "DEMO" + job_metadata["demo_source_id"] = str(payload.data_id).strip() + return ResolvedDocumentIngestionScope(job_metadata=job_metadata, document_id=documentId, namespace=DEMO_NAMESPACE) + if str(payload.document_id or "").startswith("ddoc_"): + raise ValidationException(user_message="Demo documents cannot move to a private namespace.", violations=[{"field": "namespace", "description": "Use the reserved demo namespace"}]) + job_metadata["corpus_target"] = "PRIVATE" requested_document_id = JobMetadataHelper.get_document_id(job_metadata) if requested_document_id: active_job = await find_active_job_for_document( diff --git a/apps/api/app/services/documents/lifecycle_service.py b/apps/api/app/services/documents/lifecycle_service.py index 9e1bd839e..a5b44a4c4 100644 --- a/apps/api/app/services/documents/lifecycle_service.py +++ b/apps/api/app/services/documents/lifecycle_service.py @@ -7,6 +7,9 @@ from typing import Any from app.repositories.document_repository import DocumentRepository +from app.services.demo.revision_reader import resolve_demo_revision +from shared.services.retrieval.corpus_storage import CorpusStorage +from shared.services.retrieval.demo_authorization import authorize_demo_transaction from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession @@ -26,6 +29,7 @@ lock_namespace_generation, ) from shared.services.storage.result_storage import ResultStorage, get_result_storage +from shared.services.storage.demo_asset_signer import DemoAssetSigner _DOCUMENT_CHUNK_ASSET_URL_EXPIRES_SECONDS = 7 * 24 * 60 * 60 _MEDIA_CHUNK_TYPES = frozenset({"image", "table"}) @@ -294,6 +298,7 @@ async def list_document_chunks( *, user_id: str, document_id: str, + job_result_id: str | None = None, page: int, page_size: int, chunk_type: str | None, @@ -307,7 +312,11 @@ async def list_document_chunks( if document is None: return None - job_result_id = document.current_job_result_id + if CorpusStorage.resolve_document(document_id).is_demo: + _, revision = await resolve_demo_revision(db, document_id=document_id, job_result_id=job_result_id) + job_result_id = revision.id + else: + job_result_id = job_result_id or document.current_job_result_id if not job_result_id: return { "document_id": document.document_id, @@ -345,6 +354,7 @@ async def list_document_chunks( section=section, job_id=job_result.job_id, raw_prefix=_result_raw_prefix(job_result.document_metadata), + revision_metadata=job_result.document_metadata, include_asset_urls=include_asset_urls, result_storage=result_storage, ) @@ -372,6 +382,7 @@ async def get_document_chunk( *, user_id: str, document_id: str, + job_result_id: str | None = None, document_chunk_id: str, include_asset_urls: bool, ) -> dict[str, Any] | None: @@ -380,13 +391,20 @@ async def get_document_chunk( user_id=user_id, document_id=document_id, ) - if document is None or not document.current_job_result_id: + if document is None: return None + if CorpusStorage.resolve_document(document_id).is_demo: + _, revision = await resolve_demo_revision(db, document_id=document_id, job_result_id=job_result_id) + job_result_id = revision.id + else: + job_result_id = job_result_id or document.current_job_result_id + if job_result_id is None: + return None row = await self._repository.get_current_document_chunk( db, document_id=document_id, - job_result_id=document.current_job_result_id, + job_result_id=job_result_id, document_chunk_id=document_chunk_id, ) if row is None: @@ -397,13 +415,14 @@ async def get_document_chunk( return { "document_id": document.document_id, "namespace": document.namespace, - "job_result_id": document.current_job_result_id, + "job_result_id": job_result_id, "job_id": job_result.job_id, "chunk": self._chunk_payload( chunk=chunk, section=section, job_id=job_result.job_id, raw_prefix=_result_raw_prefix(job_result.document_metadata), + revision_metadata=job_result.document_metadata, include_asset_urls=include_asset_urls, result_storage=result_storage, ), @@ -415,6 +434,7 @@ async def get_document( *, user_id: str, document_id: str, + job_result_id: str | None = None, ) -> dict[str, Any] | None: document = await self._repository.get_document( db, @@ -423,6 +443,9 @@ async def get_document( ) if document is None: return None + if CorpusStorage.resolve_document(document_id).is_demo: + _, revision = await resolve_demo_revision(db, document_id=document_id, job_result_id=job_result_id) + return {**document_payload(document), "job_result_id": revision.id} return document_payload(document) async def get_document_page_citation_source( @@ -431,17 +454,19 @@ async def get_document_page_citation_source( *, user_id: str, document_id: str, + job_result_id: str | None = None, ) -> dict[str, Any] | None: row = await self._repository.get_current_document_job_revision( db, user_id=user_id, document_id=document_id, + job_result_id=job_result_id, ) if row is None: return None document, job_result, job = row - if document.parse_track != _PAGE_MEMORY_PARSE_TRACK: + if str((job_result.document_metadata or {}).get("parse_track") or document.parse_track) != _PAGE_MEMORY_PARSE_TRACK: return None result_storage = self._result_storage or get_result_storage() @@ -468,6 +493,13 @@ async def get_document_page_citation_source( if not source_url: return None + if CorpusStorage.resolve_document(document_id).is_demo: + from app.services.demo.shared_document_service import SharedDemoDocumentService + source_url = await SharedDemoDocumentService().get_media_url( + db, source_id=document.demo_source_id, job_result_id=job_result.id, + asset_path=_PAGE_CITATION_SOURCE_FILE_NAME, + ) + expires_at = datetime.now(timezone.utc) + timedelta( seconds=_PAGE_CITATION_SOURCE_EXPIRES_SECONDS, ) @@ -492,6 +524,7 @@ def _chunk_payload( raw_prefix: str | None, include_asset_urls: bool, result_storage: ResultStorage | None, + revision_metadata: dict[str, Any] | None = None, ) -> dict[str, Any]: chunk_type = _normalize_chunk_type(chunk.chunk_type) file_path = chunk.file_path @@ -527,6 +560,15 @@ def _chunk_payload( ), "created_at": _datetime_payload(chunk.created_at), } + if include_asset_urls and job_id and CorpusStorage.resolve_document(chunk.document_id).is_demo: + signer: DemoAssetSigner = DemoAssetSigner(job_id, revision_metadata or {}) + payload["asset_url"] = signer.generate_url(file_path) if file_path and chunk_type in _MEDIA_CHUNK_TYPES else None + for asset in page_assets: + asset.pop("asset_url", None) + if url := signer.generate_url(str(asset["artifact_ref"])): + asset["asset_url"] = url + if page_assets: + metadata["page_assets"] = page_assets return payload async def archive_document( @@ -536,6 +578,8 @@ async def archive_document( user_id: str, document_id: str, ) -> dict[str, Any] | None: + if CorpusStorage.resolve_document(document_id).is_demo: + await db.run_sync(lambda session: authorize_demo_transaction(session, user_id=user_id)) document = await self._repository.get_document( db, user_id=user_id, diff --git a/apps/api/app/services/jobs/deletion_guard.py b/apps/api/app/services/jobs/deletion_guard.py new file mode 100644 index 000000000..08486538e --- /dev/null +++ b/apps/api/app/services/jobs/deletion_guard.py @@ -0,0 +1,40 @@ +"""Protect shared originals and revision assets from job deletion requests.""" + +from typing import NoReturn + +from fastapi import HTTPException +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + ConflictException, + NotFoundException, +) +from shared.models.database.job import Job +from shared.models.database.job_result import JobResult + + +async def reject_job_deletion( + db: AsyncSession, *, user_id: str, job_id: str +) -> NoReturn: + job: Job | None = ( + await db.execute( + select(Job).where(Job.job_id == job_id, Job.user_id == user_id) + ) + ).scalar_one_or_none() + if job is None: + raise NotFoundException(resource="Job", resource_id=job_id) + revision: JobResult | None = ( + await db.execute(select(JobResult).where(JobResult.job_id == job_id)) + ).scalar_one_or_none() + if revision is not None and revision.demo_document_id is not None: + raise ConflictException( + user_message="This job owns a retained shared demo revision. Its original and assets cannot be deleted; archive the demo document instead.", + resource="Demo revision", + resource_id=revision.id, + ) + # This checkout retains processing jobs and has no deletion operation. + raise HTTPException( + status_code=405, + detail="Jobs are retained processing records. Use document archive to remove content from retrieval.", + ) diff --git a/apps/api/app/services/jobs/result_projection.py b/apps/api/app/services/jobs/result_projection.py index b39166c9d..ba0424e87 100644 --- a/apps/api/app/services/jobs/result_projection.py +++ b/apps/api/app/services/jobs/result_projection.py @@ -154,6 +154,7 @@ async def build_job_result_response( result, result_url, result_url_expires_at = await _resolve_result_delivery(job) return JobResultResponse( + job_result_id=getattr(getattr(job, "job_result", None), "id", None), job_id=job.job_id, namespace=JobMetadataHelper.get_namespace(job_metadata), document_id=resolve_job_document_id(job, job_metadata), diff --git a/apps/api/app/services/jobs/source_upload_service.py b/apps/api/app/services/jobs/source_upload_service.py new file mode 100644 index 000000000..16505d2bb --- /dev/null +++ b/apps/api/app/services/jobs/source_upload_service.py @@ -0,0 +1,40 @@ +"""Renew a waiting job's source upload capability without changing its claim.""" + +from __future__ import annotations + +import asyncio +from typing import Any + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + ConflictException, + NotFoundException, +) +from shared.models.database.job import Job +from shared.services.retrieval.demo_job_scope import validate_job_corpus +from shared.services.storage.job_file_storage import JobFileStorage + + +async def renew_source_upload( + db: AsyncSession, *, user_id: str, job_id: str +) -> dict[str, Any]: + job: Job | None = ( + await db.execute( + select(Job).where(Job.job_id == job_id, Job.user_id == user_id) + ) + ).scalar_one_or_none() + if job is None: + raise NotFoundException(resource="Job", resource_id=job_id) + await db.run_sync(lambda session: validate_job_corpus(session, job=job)) + if job.status != "waiting-file" or not job.s3_key: + raise ConflictException( + user_message="Only waiting-file jobs can renew source upload URLs.", + resource="Job", + resource_id=job_id, + ) + extension: str = "." + job.s3_key.rsplit(".", 1)[-1] if "." in job.s3_key else "" + return await asyncio.to_thread( + JobFileStorage().generate_upload_url, job_id=job_id, file_extension=extension + ) diff --git a/apps/api/main.py b/apps/api/main.py index 9d8f94675..2ab5d4d4f 100644 --- a/apps/api/main.py +++ b/apps/api/main.py @@ -28,11 +28,7 @@ async def lifespan(app: FastAPI): """ Application lifecycle management """ - from shared.services.retrieval.publication_strategy import ( - announce_publication_strategy, - ) - announce_publication_strategy("knowhere-api") from shared.core.database import prewarm_connection_pool @@ -50,6 +46,8 @@ async def lifespan(app: FastAPI): redis_url = redis_pool_manager.config.get_connection_url() RateLimitConfig.get_instance(redis_url) async with get_db_context() as session: + from shared.services.retrieval.demo_authorization import verify_runtime_database_role + await session.run_sync(verify_runtime_database_role) await load_rules(session) logger.info("rate limit rules loaded at startup; restart the pod to apply changes") diff --git a/apps/api/tests/conftest.py b/apps/api/tests/conftest.py index c9a6b9e75..ce3843cee 100644 --- a/apps/api/tests/conftest.py +++ b/apps/api/tests/conftest.py @@ -12,6 +12,7 @@ from pytest import MonkeyPatch from tests.support.import_environment import configure_import_environment, ensure_import_paths from shared.testing.contract_runtime import ( + configure_runtime_database_role, CONTRACT_POSTGRESQL_PORT_RANGE, PostgreSQLProcess, clear_application_modules, @@ -23,6 +24,7 @@ seed_contract_developer, ) from shared.testing.postgresql_environment import find_executable +from shared.testing.contract_runtime import get_contract_database_url configure_import_environment() ensure_import_paths() @@ -76,11 +78,15 @@ async def _load_api_module( ) -> ModuleType: configure_contract_environment(monkeypatch, postgresql_process) await prepare_contract_storage() + adminUrl: str = get_contract_database_url(postgresql_process) + await configure_runtime_database_role(monkeypatch, adminUrl) _ensure_import_paths() clear_application_modules() return importlib.import_module("main") + + @asynccontextmanager async def _create_api_client( monkeypatch: MonkeyPatch, diff --git a/apps/api/tests/contract/test_agent_evidence_retention_contract.py b/apps/api/tests/contract/test_agent_evidence_retention_contract.py new file mode 100644 index 000000000..18ff3cba1 --- /dev/null +++ b/apps/api/tests/contract/test_agent_evidence_retention_contract.py @@ -0,0 +1,145 @@ +"""Provider/harness contract: complementary reads remain visible until finish.""" + +from __future__ import annotations + +import json +from contextlib import asynccontextmanager +from collections.abc import AsyncIterator + +import pytest +from openai.types.chat import ChatCompletion +from pytest import MonkeyPatch +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.agent_explore.budget import EpisodeBudget +from shared.services.retrieval.agent_explore.harness import openai_harness +from shared.services.retrieval.agent_tools import ToolBudget, ToolResult + + +class EvidenceProvider: + """Read what is missing from the current context; finish when both facts exist.""" + + def __init__(self) -> None: + self.calls: int = 0 + self.hasStaleDiscovery: bool = False + + def chat_completion_raw_with_usage( + self, *, messages: list[dict[str, object]], **options: object + ) -> tuple[ChatCompletion, dict[str, int]]: + self.calls += 1 + visibleText: str = "\n".join( + str(message["content"]) + for message in messages + if message.get("role") == "tool" + ) + name: str = "corpus_read" + arguments: dict[str, object] + if self.calls == 1: + name = "corpus_recall" + arguments = {"query": "launch capabilities"} + elif self.calls == 2: + arguments = { + "refs": [{"section_path": "Launch"}, {"section_path": "Heavy"}] + } + elif "Falcon 9: 23 tonnes" not in visibleText: + arguments = {"refs": [{"section_path": "Launch"}]} + elif "Falcon Heavy: 64 tonnes" not in visibleText: + arguments = {"refs": [{"section_path": "Heavy"}]} + else: + name = "finish" + arguments = { + "refs": [ + {"document_id": "ddoc_launch", "section_path": "Launch"}, + {"document_id": "ddoc_launch", "section_path": "Heavy"}, + ] + } + self.hasStaleDiscovery = "DISCOVERY_LIST" in visibleText + if options.get("tool_choice") != "auto": + name = "finish" + arguments = {"refs": []} + response: ChatCompletion = ChatCompletion.model_validate( + { + "id": f"completion-{self.calls}", + "object": "chat.completion", + "created": 0, + "model": "contract-provider", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "tool_calls": [ + { + "id": f"call-{self.calls}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(arguments), + }, + } + ], + }, + } + ], + } + ) + return response, {"total_tokens": 100} + + +@pytest.mark.asyncio +async def test_complementary_reads_survive_discovery_compaction( + monkeypatch: MonkeyPatch, +) -> None: + provider: EvidenceProvider = EvidenceProvider() + monkeypatch.setattr( + openai_harness, + "_resolve_client_and_model", + lambda: (provider, "contract-provider"), + ) + + async def dispatchTool( + name: str, arguments: dict[str, object], **context: object + ) -> ToolResult: + if name == "corpus.recall": + return ToolResult(text="DISCOVERY_LIST: Launch, Heavy") + requested: str = json.dumps(arguments) + sections: list[str] = [ + section for section in ("Launch", "Heavy") if section in requested + ] + # A multi-section read exposes the first fact but cuts off the second. + text: str = "" + if "Launch" in sections: + text += "Falcon 9: 23 tonnes\n" + "x" * ToolBudget().max_chars + if "Heavy" in sections: + text += "Falcon Heavy: 64 tonnes\n" + return ToolResult( + text=text, + refs=[ + {"document_id": "ddoc_launch", "section_path": section} + for section in sections + ], + ) + + monkeypatch.setattr(openai_harness, "dispatch_tool_call", dispatchTool) + + @asynccontextmanager + async def createDatabase() -> AsyncIterator[AsyncSession]: + raise AssertionError("The provider contract must not access the database") + yield AsyncSession() + + result = await openai_harness.OpenAIHarness().run_episode( + db_factory=createDatabase, + user_id="actual-reader", + namespace="__knowhere_demo__", + query="Compare Falcon 9 and Falcon Heavy launch capabilities", + budget=EpisodeBudget(max_steps=8, wall_clock_seconds=30), + ) + assert result.stop_reason == "finished", [step.tool_name for step in result.steps] + assert provider.calls == 4 + assert not provider.hasStaleDiscovery + assert {reference["section_path"] for reference in result.refs} == { + "Launch", + "Heavy", + } + assert "truncated" in result.steps[1].observation_text diff --git a/apps/api/tests/contract/test_canonical_demo_bundle_contract.py b/apps/api/tests/contract/test_canonical_demo_bundle_contract.py index 443d1647b..89fea590e 100644 --- a/apps/api/tests/contract/test_canonical_demo_bundle_contract.py +++ b/apps/api/tests/contract/test_canonical_demo_bundle_contract.py @@ -13,7 +13,6 @@ from httpx import AsyncClient from pytest import MonkeyPatch -from tests.support.contract_database import ContractDatabase def _write_source(directory: Path, content: str) -> None: @@ -226,99 +225,3 @@ async def test_canonical_bundle_concurrent_callers_publish_once( assert len({result.zip_key for result in results}) == 1 assert len({result.raw_prefix for result in results}) == 1 assert sum(not result.reused for result in results) == 1 - - -@pytest.mark.asyncio -async def test_http_materialization_reuses_canonical_bundle_across_namespaces( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], - monkeypatch: MonkeyPatch, -) -> None: - monkeypatch.setenv("DEMO_CANONICAL_BUNDLE_ENABLED", "true") - async with developer_api_client_factory() as api_client: - import app.services.demo.source_materializer as source_materializer_module - - first = await api_client.post( - "/api/v1/demo/materializations", - json={ - "namespace": "canonical-contract-one", - "demo_source_ids": ["demo-tsla-q4-2025"], - }, - ) - second = await api_client.post( - "/api/v1/demo/materializations", - json={ - "namespace": "canonical-contract-two", - "demo_source_ids": ["demo-tsla-q4-2025"], - }, - ) - assert first.status_code == 200 - assert second.status_code == 200 - first_document_id = str(first.json()["sources"][0]["document_id"]) - second_document_id = str(second.json()["sources"][0]["document_id"]) - chunks_response = await api_client.get( - f"/api/v1/documents/{first_document_id}/chunks?page_size=200" - "&include_asset_urls=true" - ) - retrieval_response = await api_client.post( - "/api/v1/retrieval/query", - json={ - "namespace": "canonical-contract-one", - "query": "Tesla investment", - "top_k": 3, - "use_agentic": False, - }, - ) - assert retrieval_response.status_code == 200 - retrieval_asset_urls: list[str] = [ - str(result["asset_url"]) - for result in retrieval_response.json()["results"] - if result.get("asset_url") - ] - assert retrieval_asset_urls - assert all("demo-canonical" in url for url in retrieval_asset_urls) - rows = await ContractDatabase.fetch_all( - """ - SELECT d.document_id, jr.result_s3_key, jr.result_size, - jr.document_metadata ->> 'result_raw_prefix' AS result_raw_prefix - FROM job_results AS jr - JOIN documents AS d ON d.current_job_result_id = jr.id - WHERE d.document_id IN (:first_document_id, :second_document_id) - ORDER BY d.document_id - """, - { - "first_document_id": first_document_id, - "second_document_id": second_document_id, - }, - ) - - monkeypatch.setattr( - source_materializer_module.settings, - "DEMO_CANONICAL_BUNDLE_ENABLED", - False, - ) - existing_document_response = await api_client.get( - f"/api/v1/documents/{first_document_id}/chunks?page_size=200" - "&include_asset_urls=true" - ) - - chunks_body = chunks_response.json() - asset_urls = [ - str(chunk["asset_url"]) - for chunk in chunks_body["chunks"] - if chunk.get("asset_url") - ] - assert asset_urls - assert any("demo-canonical" in asset_url for asset_url in asset_urls) - assert existing_document_response.status_code == 200 - existing_asset_urls = [ - str(chunk["asset_url"]) - for chunk in existing_document_response.json()["chunks"] - if chunk.get("asset_url") - ] - assert existing_asset_urls == asset_urls - assert len(rows) == 2 - assert len({str(row["result_s3_key"]) for row in rows}) == 1 - assert len({str(row["result_raw_prefix"]) for row in rows}) == 1 - assert all(int(row["result_size"]) > 0 for row in rows) diff --git a/apps/api/tests/contract/test_demo_documents_contract.py b/apps/api/tests/contract/test_demo_documents_contract.py index f283dfc8f..4b1d0e351 100644 --- a/apps/api/tests/contract/test_demo_documents_contract.py +++ b/apps/api/tests/contract/test_demo_documents_contract.py @@ -1,24 +1,18 @@ from __future__ import annotations -import asyncio import importlib import sys -from collections.abc import Callable -from contextlib import AbstractAsyncContextManager from pathlib import Path from types import SimpleNamespace from types import ModuleType -from typing import Any, cast +from typing import Any -import pytest -from httpx import AsyncClient from pytest import MonkeyPatch from tests.support.import_environment import ( configure_import_environment, ensure_import_paths, ) -from tests.support.contract_database import ContractDatabase DEMO_SOURCE_ID = "demo-tsla-q4-2025" @@ -155,565 +149,3 @@ def test_should_preserve_filename_rooted_sections_when_publishing_demo_chunks() assert f"{source.title}/Root/MANAGEMENT DISCUSSION SECTION" in publication_paths assert f"{source.title}/Root/QUESTION AND ANSWER SECTION" in publication_paths - - -@pytest.mark.asyncio -async def test_should_return_demo_catalog_with_resolvable_canonical_citations( - api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], -) -> None: - async with api_client_factory() as api_client: - catalog_response = await api_client.get("/api/v1/demo/catalog") - - assert catalog_response.status_code == 200 - catalog = cast(dict[str, Any], catalog_response.json()) - sources = cast(list[dict[str, Any]], catalog["sources"]) - sources_by_id = {str(source["demo_source_id"]): source for source in sources} - source = sources_by_id[DEMO_SOURCE_ID] - spacex_source = sources_by_id[SPACEX_DEMO_SOURCE_ID] - official_library = cast(dict[str, Any], catalog["official_library"]) - library_sources = cast(list[dict[str, Any]], official_library["sources"]) - library_sources_by_id = { - str(source["library_source_id"]): source for source in library_sources - } - examples = cast(list[dict[str, Any]], source["examples"]) - citations = cast(list[dict[str, Any]], examples[0]["citations"]) - citation = citations[0] - spacex_examples = cast(list[dict[str, Any]], spacex_source["examples"]) - spacex_citations = cast(list[dict[str, Any]], spacex_examples[0]["citations"]) - - assert str(sources[0]["demo_source_id"]) == DEMO_SOURCE_ID - assert SPACEX_DEMO_SOURCE_ID in sources_by_id - assert source["demo_source_id"] == DEMO_SOURCE_ID - assert source["canonical_document_id"] == "demo-doc-tsla-q4-2025" - assert source["chunk_count"] == 71 - assert source["original_file"]["can_download"] is False - assert citation["canonical_document_id"] == "demo-doc-tsla-q4-2025" - assert citation["canonical_chunk_id"].startswith(f"{DEMO_SOURCE_ID}:") - assert "[[cite:1]]" in examples[0]["answer"] - assert citation["page_citation_page_number"] == 12 - assert citation["page_nums"] == [12] - assert str(citation["page_citation_asset_url"]).endswith( - "/page_citation_assets/page-12.png" - ) - assert spacex_source["official_library"]["library_source_id"] == ( - "financial-spacex-s1" - ) - assert spacex_source["canonical_document_id"] == "demo-doc-spacex-s1" - assert spacex_source["chunk_count"] == 227 - assert spacex_citations[0]["canonical_chunk_id"].startswith( - f"{SPACEX_DEMO_SOURCE_ID}:" - ) - assert "[[cite:1]]" in spacex_examples[0]["answer"] - assert spacex_citations[0]["page_citation_page_number"] == 28 - assert spacex_citations[0]["page_nums"] == [28] - assert str(spacex_citations[0]["source"]["section_path"]).endswith("/Root") - assert [ - category["category_id"] - for category in cast(list[dict[str, Any]], official_library["categories"]) - ] == ["financial-reports", "research-papers", "stem-books"] - assert library_sources_by_id["financial-spacex-s1"]["status"] == "ready" - assert library_sources_by_id["financial-spacex-s1"]["demo_source_id"] == ( - SPACEX_DEMO_SOURCE_ID - ) - assert library_sources_by_id["financial-spacex-s1"]["chunk_count"] == 227 - assert library_sources_by_id["stem-statistical-learning"]["status"] == "ready" - assert library_sources_by_id["stem-statistical-learning"]["demo_source_id"] == ( - "demo-stem-statistical-learning" - ) - assert library_sources_by_id["stem-statistical-learning"]["chunk_count"] == 45 - research_demo_sources = { - "research-attention-is-all-you-need": ( - "demo-research-attention-is-all-you-need", - 34, - ), - "research-rag-survey": ("demo-research-rag-survey", 35), - "research-rag-realized": ("demo-research-rag-realized", 59), - "research-toolformer": ("demo-research-toolformer", 41), - "research-ai-agents-overview": ("demo-research-ai-agents-overview", 58), - } - for library_source_id, ( - expected_demo_source_id, - expected_chunk_count, - ) in research_demo_sources.items(): - assert library_sources_by_id[library_source_id]["status"] == "ready" - assert ( - library_sources_by_id[library_source_id]["demo_source_id"] - == expected_demo_source_id - ) - assert ( - library_sources_by_id[library_source_id]["chunk_count"] - == expected_chunk_count - ) - assert ( - library_sources_by_id["financial-nvda-q1-fy27-earnings-call"]["demo_source_id"] - == "demo-financial-nvda-q1-fy27-earnings-call" - ) - assert library_sources_by_id["financial-micron-report-530bd7ed"]["status"] == ( - "ready" - ) - assert ( - library_sources_by_id["financial-micron-report-530bd7ed"]["demo_source_id"] - == "demo-financial-micron-report-530bd7ed" - ) - assert ( - library_sources_by_id["financial-micron-report-530bd7ed"]["chunk_count"] == 77 - ) - assert ( - library_sources_by_id["financial-micron-report-9c0becf5"]["demo_source_id"] - == "demo-financial-micron-report-9c0becf5" - ) - assert ( - library_sources_by_id["financial-micron-report-9c0becf5"]["chunk_count"] == 82 - ) - assert library_sources_by_id["stem-transformers-tutorial"]["demo_source_id"] == ( - "demo-stem-transformers-tutorial" - ) - assert library_sources_by_id["stem-transformers-tutorial"]["chunk_count"] == 301 - assert library_sources_by_id["stem-information-theory"]["status"] == "planned" - - async with api_client_factory() as api_client: - chunks_response = await api_client.get( - f"/api/v1/demo/sources/{DEMO_SOURCE_ID}/chunks?page_size=200" - ) - chunk_response = await api_client.get( - "/api/v1/demo/sources/" - f"{DEMO_SOURCE_ID}/chunks/{citation['canonical_chunk_id']}" - ) - - assert chunks_response.status_code == 200 - assert chunk_response.status_code == 200 - chunks_body = cast(dict[str, Any], chunks_response.json()) - chunk_page = cast(list[dict[str, Any]], chunks_body["chunks"]) - asset_url = next( - str(chunk["asset_url"]) for chunk in chunk_page if chunk.get("asset_url") - ) - chunk_body = cast(dict[str, Any], chunk_response.json()) - chunk = cast(dict[str, Any], chunk_body["chunk"]) - - assert chunk["id"] == citation["canonical_chunk_id"] - assert citation["content"] in chunk["content"] - - async with api_client_factory() as api_client: - asset_response = await api_client.get(asset_url) - original_response = await api_client.get( - f"/api/v1/demo/sources/{DEMO_SOURCE_ID}/original" - ) - internal_asset_response = await api_client.get( - f"/api/v1/demo/sources/{DEMO_SOURCE_ID}/assets/full.md" - ) - - assert asset_response.status_code == 200 - assert asset_response.headers["content-disposition"].startswith("inline") - assert original_response.status_code == 200 - assert original_response.headers["content-disposition"].startswith("inline") - assert internal_asset_response.status_code == 404 - - -@pytest.mark.asyncio -async def test_should_materialize_demo_source_without_parse_or_credit_charge( - 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, - ) - empty_cached_response = await api_client.post( - "/api/v1/retrieval/query", - json={ - "namespace": "contract-demo", - "query": "xAI investment", - "top_k": 5, - "use_agentic": False, - }, - ) - first_response = await api_client.post( - "/api/v1/demo/materializations", - json={ - "namespace": "contract-demo", - "demo_source_ids": [DEMO_SOURCE_ID], - }, - ) - retry_response = await api_client.post( - "/api/v1/demo/materializations", - json={ - "namespace": "contract-demo", - "demo_source_ids": [DEMO_SOURCE_ID], - }, - ) - retrieval_response = await api_client.post( - "/api/v1/retrieval/query", - json={ - "namespace": "contract-demo", - "query": "xAI investment", - "top_k": 5, - "use_agentic": False, - }, - ) - document_chunks_response = await api_client.get( - f"/api/v1/documents/{first_response.json()['sources'][0]['document_id']}" - "/chunks?page_size=200" - ) - - assert empty_cached_response.status_code == 200 - assert first_response.status_code == 200 - assert retry_response.status_code == 409 - assert retrieval_response.status_code == 200 - assert document_chunks_response.status_code == 200 - - empty_cached_body = cast(dict[str, Any], empty_cached_response.json()) - assert cast(list[dict[str, Any]], empty_cached_body["results"]) == [] - - first_source = cast(dict[str, Any], first_response.json()["sources"][0]) - document_id = str(first_source["document_id"]) - - assert first_source["status"] == "created" - assert retry_response.json()["error"]["details"]["reason"] == "ALREADY_EXISTS" - - 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' - AND demo_source_id = :demo_source_id - """, - {"demo_source_id": DEMO_SOURCE_ID}, - ) - document_row = await ContractDatabase.fetch_one( - """ - SELECT document_id, status, source_file_name - FROM documents - WHERE document_id = :document_id - """, - {"document_id": document_id}, - ) - chunk_rows = await ContractDatabase.fetch_all( - """ - SELECT id - FROM document_chunks - WHERE document_id = :document_id - """, - {"document_id": document_id}, - ) - job_rows = await ContractDatabase.fetch_all( - """ - SELECT job_id, status, job_type, credits_charged, billing_status - FROM jobs - WHERE user_id = 'local-dev-user' - AND job_metadata ->> 'demo_source_id' = :demo_source_id - """, - {"demo_source_id": DEMO_SOURCE_ID}, - ) - - assert materialization_rows == [ - {"demo_source_id": DEMO_SOURCE_ID, "document_id": document_id} - ] - assert document_row == { - "document_id": document_id, - "status": "active", - "source_file_name": "TSLA-Q4-2025-Update.pdf", - } - assert len(chunk_rows) == 71 - assert len(job_rows) == 1 - job_row = job_rows[0] - assert job_row["status"] == "done" - assert job_row["job_type"] == "demo_materialization" - assert job_row["credits_charged"] == 0 - assert job_row["billing_status"] == "skipped" - - retrieval_body = cast(dict[str, Any], retrieval_response.json()) - retrieval_results = cast(list[dict[str, Any]], retrieval_body["results"]) - chunk_page_body = cast(dict[str, Any], document_chunks_response.json()) - materialized_chunks = cast(list[dict[str, Any]], chunk_page_body["chunks"]) - media_chunks = [ - chunk - for chunk in materialized_chunks - if chunk["chunk_type"] in {"image", "table"} - ] - - assert retrieval_body["namespace"] == "contract-demo" - assert retrieval_results - assert retrieval_results[0]["source"]["document_id"] == document_id - assert retrieval_results[0]["source"]["section_path"] != "Root" - assert media_chunks - assert media_chunks[0]["file_path"] - uploaded_files = fake_result_storage.raw_files_by_job_id[str(job_row["job_id"])] - assert any(file_path.startswith("images/") for file_path in uploaded_files) - assert any(file_path.startswith("tables/") for file_path in uploaded_files) - assert any( - file_path.startswith("page_citation_assets/") for file_path in uploaded_files - ) - - -@pytest.mark.asyncio -async def test_should_release_materialization_claims_after_publication_failure( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], - monkeypatch: MonkeyPatch, -) -> None: - async def fail_materialization(*args: Any, **kwargs: Any) -> Any: - raise TimeoutError("publication timed out") - - async with developer_api_client_factory() as api_client: - import app.services.demo.source_materializer as source_materializer_module - - monkeypatch.setattr( - source_materializer_module.DemoSourceMaterializer, - "_materialize_source", - fail_materialization, - ) - - with pytest.raises(TimeoutError, match="publication timed out"): - await api_client.post( - "/api/v1/demo/materializations", - json={ - "namespace": "contract-demo-publication-failure", - "demo_source_ids": [DEMO_SOURCE_ID], - }, - ) - - materialization_rows = await ContractDatabase.fetch_all( - """ - SELECT demo_source_id, status, document_id - FROM demo_materializations - WHERE user_id = 'local-dev-user' - AND namespace = 'contract-demo-publication-failure' - """, - ) - - assert materialization_rows == [] - - -@pytest.mark.asyncio -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() - - 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_reject_sequential_duplicate_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 = await api_client.post( - "/api/v1/demo/materializations", - json={"namespace": "contract-demo-duplicate", "demo_source_ids": [DEMO_SOURCE_ID]}, - ) - second_response = await api_client.post( - "/api/v1/demo/materializations", - json={"namespace": "contract-demo-duplicate", "demo_source_ids": [DEMO_SOURCE_ID]}, - ) - - assert first_response.status_code == 200 - assert second_response.status_code == 409 - assert second_response.json()["error"]["details"]["reason"] == "ALREADY_EXISTS" - assert len(fake_result_storage.raw_files_by_job_id) == 1 - - -@pytest.mark.asyncio -async def test_should_reject_concurrent_duplicate_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", - json={ - "namespace": "contract-demo-race", - "demo_source_ids": [DEMO_SOURCE_ID], - }, - ), - api_client.post( - "/api/v1/demo/materializations", - json={ - "namespace": "contract-demo-race", - "demo_source_ids": [DEMO_SOURCE_ID], - }, - ), - ) - - statuses = {first_response.status_code, second_response.status_code} - assert statuses == {200, 409} - conflict_response = ( - first_response if first_response.status_code == 409 else second_response - ) - assert conflict_response.json()["error"]["details"]["reason"] == "ABORTED" - - 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-race' - AND demo_source_id = :demo_source_id - """, - {"demo_source_id": DEMO_SOURCE_ID}, - ) - job_rows = await ContractDatabase.fetch_all( - """ - SELECT job_id - FROM jobs - WHERE user_id = 'local-dev-user' - AND job_metadata ->> 'namespace' = 'contract-demo-race' - AND job_metadata ->> 'demo_source_id' = :demo_source_id - """, - {"demo_source_id": DEMO_SOURCE_ID}, - ) - - assert len(materialization_rows) == 1 - assert materialization_rows[0]["document_id"] is not None - assert len(job_rows) == 1 - - -@pytest.mark.asyncio -async def test_should_reject_blank_demo_materialization_selection( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], -) -> None: - async with developer_api_client_factory() as api_client: - response = await api_client.post( - "/api/v1/demo/materializations", - json={ - "namespace": "contract-demo-blank", - "demo_source_ids": [" ", "\t"], - }, - ) - - assert response.status_code == 400 - - -@pytest.mark.asyncio -async def test_should_reject_mixed_demo_materialization_selection_before_upload( - 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, - ) - response = await api_client.post( - "/api/v1/demo/materializations", - json={ - "namespace": "contract-demo-mixed-invalid", - "demo_source_ids": [DEMO_SOURCE_ID, "missing-demo-source"], - }, - ) - - job_rows = await ContractDatabase.fetch_all( - """ - SELECT job_id - FROM jobs - WHERE user_id = 'local-dev-user' - AND job_metadata ->> 'namespace' = 'contract-demo-mixed-invalid' - """, - ) - 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-mixed-invalid' - """, - ) - - assert response.status_code == 404 - assert fake_result_storage.raw_files_by_job_id == {} - assert job_rows == [] - assert materialization_rows == [] diff --git a/apps/api/tests/contract/test_demo_preparation_cache_contract.py b/apps/api/tests/contract/test_demo_preparation_cache_contract.py deleted file mode 100644 index 20005952a..000000000 --- a/apps/api/tests/contract/test_demo_preparation_cache_contract.py +++ /dev/null @@ -1,167 +0,0 @@ -"""HTTP and database contract for demo preparation-cache parity.""" - -from __future__ import annotations - -from collections.abc import Callable -from contextlib import AbstractAsyncContextManager -from typing import Any, cast - -import pytest -from httpx import AsyncClient -from pytest import MonkeyPatch - -from tests.support.contract_database import ContractDatabase - -_SOURCE_ID = "demo-tsla-q4-2025" - - -async def _publish( - client: AsyncClient, - *, - namespace: str, -) -> str: - response = await client.post( - "/api/v1/demo/materializations", - json={"namespace": namespace, "demo_source_ids": [_SOURCE_ID]}, - ) - assert response.status_code == 200 - return str(response.json()["sources"][0]["document_id"]) - - -async def _snapshot(document_id: str) -> dict[str, list[dict[str, Any]]]: - chunks = await ContractDatabase.fetch_all( - """ - SELECT chunk_id, chunk_type, content, content_search_text, - path_search_text, term_search_text, source_chunk_path, - file_path, sort_order - FROM document_chunks - WHERE document_id = :document_id - ORDER BY sort_order, chunk_id, id - """, - {"document_id": document_id}, - ) - units = await ContractDatabase.fetch_all( - """ - SELECT s.section_path, m.unit_kind, m.path_token_count, - m.content_token_count, m.has_image, m.has_table, m.sort_order - FROM document_map_units AS m - JOIN document_sections AS s ON s.section_id = m.section_id - WHERE m.document_id = :document_id - ORDER BY m.sort_order, s.section_path, m.unit_kind, m.id - """, - {"document_id": document_id}, - ) - tokens = await ContractDatabase.fetch_all( - """ - SELECT s.section_path, m.unit_kind, m.sort_order, - t.channel, t.token, t.token_hash, t.frequency - FROM document_map_unit_tokens AS t - JOIN document_map_units AS m ON m.id = t.map_unit_id - JOIN document_sections AS s ON s.section_id = m.section_id - WHERE m.document_id = :document_id - ORDER BY m.sort_order, s.section_path, m.unit_kind, t.channel, t.token - """, - {"document_id": document_id}, - ) - indexes = await ContractDatabase.fetch_all( - """ - SELECT format_version, unit_count, token_count, - average_idf_path, average_idf_content, - path_document_count, path_total_length, - content_document_count, content_total_length - FROM document_map_unit_indexes - WHERE document_id = :document_id - """, - {"document_id": document_id}, - ) - units = sorted(units, key=lambda row: tuple(str(value) for value in row.values())) - tokens = sorted( - tokens, key=lambda row: tuple(str(value) for value in row.values()) - ) - return {"chunks": chunks, "units": units, "tokens": tokens, "indexes": indexes} - - -@pytest.mark.asyncio -async def test_demo_preparation_cache_preserves_persisted_semantics_across_modes( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], - monkeypatch: MonkeyPatch, -) -> None: - monkeypatch.setenv("DEMO_CANONICAL_BUNDLE_ENABLED", "true") - monkeypatch.setenv("DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED", "false") - async with developer_api_client_factory() as client: - from shared.services.retrieval.publication_preparation_cache import ( - PublicationPreparationCache, - ) - - PublicationPreparationCache.clear() - cold_off_id = await _publish(client, namespace="prep-cache-off") - off_snapshot = await _snapshot(cold_off_id) - - import shared.core.config as config_module - - monkeypatch.setattr( - config_module.settings, - "DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED", - True, - ) - PublicationPreparationCache.clear() - cache_counts = {"hits": 0, "misses": 0} - original_get = PublicationPreparationCache._get - - def observe_get(key: object) -> object: - value = original_get(key) # type: ignore[arg-type] - cache_counts["hits" if value is not None else "misses"] += 1 - return value - - monkeypatch.setattr( - PublicationPreparationCache, - "_get", - classmethod(lambda _cls, key: observe_get(key)), - ) - cold_on_id = await _publish(client, namespace="prep-cache-cold") - cold_cache_hits = cache_counts["hits"] - cold_cache_misses = cache_counts["misses"] - warm_on_id = await _publish(client, namespace="prep-cache-warm") - assert cold_cache_misses > 0 - assert cache_counts["hits"] > cold_cache_hits - assert cache_counts["misses"] == cold_cache_misses - assert len(PublicationPreparationCache._entries) > 0 - cold_snapshot = await _snapshot(cold_on_id) - warm_snapshot = await _snapshot(warm_on_id) - - monkeypatch.setattr( - config_module.settings, - "DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED", - False, - ) - rollback_id = await _publish(client, namespace="prep-cache-rollback") - rollback_snapshot = await _snapshot(rollback_id) - - for namespace, expected_document_id in ( - ("prep-cache-off", cold_off_id), - ("prep-cache-cold", cold_on_id), - ("prep-cache-warm", warm_on_id), - ("prep-cache-rollback", rollback_id), - ): - retrieval = await client.post( - "/api/v1/retrieval/query", - json={ - "namespace": namespace, - "query": "Tesla investment", - "top_k": 3, - "use_agentic": False, - }, - ) - assert retrieval.status_code == 200 - result_documents = { - str(result["source"]["document_id"]) - for result in cast(list[dict[str, Any]], retrieval.json()["results"]) - } - assert result_documents - assert result_documents == {expected_document_id} - - assert cold_snapshot == off_snapshot - assert warm_snapshot == off_snapshot - assert rollback_snapshot == off_snapshot diff --git a/apps/api/tests/contract/test_demo_publication_trace_owner_contract.py b/apps/api/tests/contract/test_demo_publication_trace_owner_contract.py deleted file mode 100644 index e6d21e389..000000000 --- a/apps/api/tests/contract/test_demo_publication_trace_owner_contract.py +++ /dev/null @@ -1,244 +0,0 @@ -"""Contract for the demo adapter's publication transaction trace boundary.""" - -from __future__ import annotations - -from collections.abc import Callable -from pathlib import Path -from types import SimpleNamespace -from typing import cast - -import pytest -from loguru import logger -from pytest import MonkeyPatch -from sqlalchemy.ext.asyncio import AsyncSession - -from tests.support.import_environment import ( - configure_import_environment, - ensure_import_paths, -) - -configure_import_environment() -ensure_import_paths() - -from app.services.demo.source_catalog import DemoSourceCatalog, DemoSourceDefinition # noqa: E402 -from app.services.demo.canonical_bundle_result import CanonicalDemoBundle # noqa: E402 -from app.services.demo.source_materializer import DemoSourceMaterializer # noqa: E402 -from shared.models.database.demo_materialization import DemoMaterialization # noqa: E402 -from shared.models.database.job import Job # noqa: E402 -from shared.services.jobs.lifecycle.publication_trace import PublicationTrace # noqa: E402 -from shared.services.retrieval.publication_models import PublishedDocumentState # noqa: E402 -from shared.services.retrieval.publication_service import RetrievalPublicationService # noqa: E402 -from shared.services.redis import RedisPublicationSemaphore, RedisService # noqa: E402 - - -class FakeCatalog: - def source_directory(self, source: DemoSourceDefinition) -> Path: - return Path("/tmp") - - def publication_chunks( - self, source: DemoSourceDefinition - ) -> list[dict[str, object]]: - return [{"type": "text", "content": "contract content", "metadata": {}}] - - -class FakeConnection: - def __init__(self) -> None: - self.engine: object = object() - self.info: dict[str, object] = {} - - -class FakeDatabase: - def __init__(self, actions: list[str]) -> None: - self.actions = actions - self.bind: object = object() - self.connection_value = FakeConnection() - self.job: Job | None = None - - def add(self, value: object) -> None: - if isinstance(value, Job): - self.job = value - - async def connection(self) -> FakeConnection: - self.actions.append("connection") - return self.connection_value - - async def flush(self) -> None: - self.actions.append("flush") - - async def commit(self) -> None: - self.actions.append("commit") - - async def rollback(self) -> None: - self.actions.append("rollback") - - async def run_sync(self, operation: Callable[[object], object]) -> object: - return operation(self) - - -class FakeSemaphore: - def __init__(self, actions: list[str]) -> None: - self.actions = actions - - async def acquire(self) -> float: - self.actions.append("acquire") - return 0.0 - - async def release(self) -> None: - self.actions.append("release") - - -class FakePublicationService: - def __init__(self, database: FakeDatabase, *, should_fail: bool = False) -> None: - self.database = database - self.should_fail = should_fail - self.traces: list[PublicationTrace | None] = [] - - def publish_document_state( - self, database: object, **kwargs: object - ) -> PublishedDocumentState: - assert database is self.database - trace = cast(PublicationTrace | None, kwargs["trace"]) - self.traces.append(trace) - if self.should_fail: - raise RuntimeError("publication failed") - assert self.database.job is not None - metadata = cast(dict[str, str], self.database.job.job_metadata) - return PublishedDocumentState( - user_id=str(self.database.job.user_id), - namespace=metadata["namespace"], - document_id=metadata["document_id"], - manifest_payload={"version": "1"}, - ) - - def publish_document_graph(self, database: object, **kwargs: object) -> None: - assert database is self.database - self.traces.append(cast(PublicationTrace | None, kwargs["trace"])) - - def update_namespace_snapshot(self, database: object, **kwargs: object) -> None: - assert database is self.database - self.traces.append(cast(PublicationTrace | None, kwargs["trace"])) - - -def make_source() -> DemoSourceDefinition: - return DemoSourceDefinition( - demo_source_id="demo-contract", - canonical_document_id="demo-doc-contract", - title="Contract source.pdf", - mime_type="application/pdf", - size_bytes=10, - asset_directory="contract", - chunk_count=1, - examples=(), - ) - - -def make_materializer( - publication: FakePublicationService, - materializer_type: type[DemoSourceMaterializer], -) -> DemoSourceMaterializer: - materializer = object.__new__(materializer_type) - materializer._catalog = cast(DemoSourceCatalog, FakeCatalog()) - materializer._publication_service = cast(RetrievalPublicationService, publication) - materializer._redis_service = cast(RedisService, object()) - return materializer - - -@pytest.mark.asyncio -@pytest.mark.parametrize("should_fail", [False, True]) -async def test_demo_owner_finishes_trace_at_publication_commit_or_rollback( - monkeypatch: MonkeyPatch, - should_fail: bool, -) -> None: - import app.services.demo.source_materializer as materializer_module - - actions: list[str] = [] - database = FakeDatabase(actions) - publication = FakePublicationService(database, should_fail=should_fail) - materializer = make_materializer( - publication, - materializer_module.DemoSourceMaterializer, - ) - source = make_source() - claim = SimpleNamespace( - document_id=None, - status="materializing", - claimed_at=object(), - updated_at=None, - ) - - async def use_bundle(**kwargs: object) -> CanonicalDemoBundle: - assert kwargs["demo_source_id"] == source.demo_source_id - assert kwargs["source_directory"] == Path("/tmp") - return CanonicalDemoBundle( - zip_key="zip/123", - zip_size=123, - raw_prefix="results/demo-canonical/", - content_version="v1", - reused=False, - ) - - monkeypatch.setattr(materializer_module, "_upload_demo_result_bundle", use_bundle) - monkeypatch.setattr( - materializer_module, - "install_publication_trace_sql_instrumentation", - lambda engine: actions.append("install"), - ) - monkeypatch.setattr( - materializer_module, - "settings", - SimpleNamespace( - KNOWHERE_PUBLICATION_TRACE_ENABLED=True, - DEMO_CANONICAL_BUNDLE_ENABLED=True, - DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED=False, - ), - ) - terminal_payloads: list[dict[str, object]] = [] - - def capture_terminal(message: object) -> None: - record = message.record # type: ignore[attr-defined] - if record["extra"].get("event") == "publication.trace.terminal": - actions.append("terminal") - terminal_payloads.append( - cast(dict[str, object], record["extra"]["publication_trace"]) - ) - - sink_id = logger.add(capture_terminal, level="INFO") - try: - if should_fail: - with pytest.raises(RuntimeError, match="publication failed"): - await materializer._materialize_source( - cast(AsyncSession, database), - user_id="user_123", - namespace="namespace_123", - source=source, - claim=cast(DemoMaterialization, claim), - semaphore=cast(RedisPublicationSemaphore, FakeSemaphore(actions)), - ) - else: - await materializer._materialize_source( - cast(AsyncSession, database), - user_id="user_123", - namespace="namespace_123", - source=source, - claim=cast(DemoMaterialization, claim), - semaphore=cast(RedisPublicationSemaphore, FakeSemaphore(actions)), - ) - finally: - logger.remove(sink_id) - - assert len(terminal_payloads) == 1 - assert terminal_payloads[0]["owner"] == "async" - assert terminal_payloads[0]["outcome"] == ("rollback" if should_fail else "success") - assert terminal_payloads[0]["scope_fingerprint"] != "namespace_123" - assert database.connection_value.info == {} - assert actions.index("install") < actions.index("connection") - if should_fail: - assert actions[-3:] == ["rollback", "terminal", "release"] - assert len(publication.traces) == 1 - else: - assert actions[-4:] == ["flush", "commit", "terminal", "release"] - assert claim.status == "ready" - assert actions.index("commit") < actions.index("terminal") - assert actions.count("commit") == 1 - assert len(publication.traces) == 3 - assert all(trace is not None for trace in publication.traces) diff --git a/apps/api/tests/contract/test_job_rate_limit_contract.py b/apps/api/tests/contract/test_job_rate_limit_contract.py index 150d85fc5..79a4dd7e6 100644 --- a/apps/api/tests/contract/test_job_rate_limit_contract.py +++ b/apps/api/tests/contract/test_job_rate_limit_contract.py @@ -13,6 +13,7 @@ from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from shared.testing.contract_runtime import ( + configure_runtime_database_role, PostgreSQLProcess, clear_application_modules, cleanup_contract_runtime_async, @@ -130,6 +131,13 @@ async def _create_rate_limited_developer_api_client( monkeypatch.setenv("BILLING_ENABLED", "true" if billing_enabled else "false") monkeypatch.setenv("RATE_LIMIT_ENABLED", "true" if rate_limit_enabled else "false") await prepare_contract_storage() + engine = await _create_contract_engine() + try: + async with engine.connect() as connection: + previousTier = (await connection.execute(text("SELECT max_concurrent_jobs, rpm_limit, daily_quota FROM tier_limits WHERE tier_name='tier_5'"))).mappings().one() + previousSystem = (await connection.execute(text("SELECT rpm, period FROM system_limits WHERE method='*' AND api_pattern='*'"))).mappings().one() + finally: + await engine.dispose() await _set_local_developer_tier_limits( max_concurrent_jobs=max_concurrent_jobs, rpm_limit=rpm_limit, @@ -139,6 +147,7 @@ async def _create_rate_limited_developer_api_client( rpm=default_system_rpm, period=default_system_period, ) + await configure_runtime_database_role(monkeypatch, get_contract_database_url(postgresql_process)) _ensure_import_paths() clear_application_modules() @@ -158,6 +167,8 @@ async def _create_rate_limited_developer_api_client( ) yield client finally: + await _set_local_developer_tier_limits(**dict(previousTier)) + await _set_default_system_limit(**dict(previousSystem)) await cleanup_contract_runtime_async(remove_test_directories=True) diff --git a/apps/api/tests/contract/test_publication_benchmark_runner_contract.py b/apps/api/tests/contract/test_publication_benchmark_runner_contract.py index 818be5707..274b96c17 100644 --- a/apps/api/tests/contract/test_publication_benchmark_runner_contract.py +++ b/apps/api/tests/contract/test_publication_benchmark_runner_contract.py @@ -150,7 +150,7 @@ def test_run_publication_refuses_a_missing_frozen_input(tmp_path: Path) -> None: owner="sync", database_url=DATABASE_URL, input_directory=layout.frozen_input_directory(), - strategy="baseline", + strategy="candidate", mode="cold", run_id=RUN_ID, ) @@ -166,41 +166,22 @@ def test_run_publication_refuses_an_unlisted_clone(tmp_path: Path) -> None: owner="sync", database_url=DATABASE_URL, input_directory=layout.frozen_input_directory(), - strategy="baseline", + strategy="candidate", mode="cold", run_id=RUN_ID, expectations=expectations_for(payload), ) -def test_run_publication_refuses_a_strategy_that_differs_from_settings( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - import importlib - import sys - +def test_run_publication_refuses_the_removed_baseline(tmp_path: Path) -> None: layout = benchmark_layout(tmp_path) payload, _ = write_synthetic_frozen_input(layout.frozen_input_directory()) write_listed_clone(layout, run_id=RUN_ID, database_url=DATABASE_URL) - live_config = importlib.import_module("shared.core.config") - monkeypatch.setattr( - live_config.settings, - "KNOWHERE_PUBLICATION_STRATEGY", - "candidate", - ) - monkeypatch.setitem(sys.modules, "shared.core.config", live_config) - with pytest.raises(BenchmarkGuardError, match="does not match"): run_publication_point( - layout=layout, - owner="sync", - database_url=DATABASE_URL, - input_directory=layout.frozen_input_directory(), - strategy="baseline", - mode="cold", - run_id=RUN_ID, - expectations=expectations_for(payload), + layout=layout, owner="sync", database_url=DATABASE_URL, + input_directory=layout.frozen_input_directory(), strategy="baseline", + mode="cold", run_id=RUN_ID, expectations=expectations_for(payload), ) @@ -222,7 +203,7 @@ def test_run_publication_refuses_a_cold_sample_on_a_non_fresh_clone( owner="sync", database_url=DATABASE_URL, input_directory=layout.frozen_input_directory(), - strategy="baseline", + strategy="candidate", mode="cold", run_id=RUN_ID, expectations=expectations_for(payload), @@ -273,7 +254,7 @@ def test_run_publication_records_a_sample_with_the_owner_executor( owner="sync", database_url=DATABASE_URL, input_directory=layout.frozen_input_directory(), - strategy="baseline", + strategy="candidate", mode="cold", run_id=RUN_ID, expectations=expectations_for(payload), @@ -289,7 +270,7 @@ def test_run_publication_records_a_sample_with_the_owner_executor( assert result["record"]["outcome"] == "committed" assert result["record"]["publication_duration_ms"] == pytest.approx(1234.5) assert result["record"]["input_digest"] == manifest.digest - assert result["record"]["strategy"] == "baseline" + assert result["record"]["strategy"] == "candidate" assert result["record"]["source_content_digest"] == SOURCE_CONTENT_DIGEST assert result["record"]["publication_attempt_ref"].startswith("attempt-") assert read_clone_record(layout.clone_record_path(RUN_ID)).state == "sampled" @@ -333,7 +314,7 @@ def test_run_publication_rejects_source_drift_without_writing_a_sample( owner="sync", database_url=DATABASE_URL, input_directory=layout.frozen_input_directory(), - strategy="baseline", + strategy="candidate", mode="cold", run_id=RUN_ID, expectations=expectations_for(payload), @@ -384,7 +365,7 @@ def test_run_publication_uses_database_bounded_job_and_result_ids( owner="sync", database_url=DATABASE_URL, input_directory=layout.frozen_input_directory(), - strategy="baseline", + strategy="candidate", mode="cold", run_id=long_run_id, expectations=expectations_for(payload), @@ -410,7 +391,7 @@ def test_run_publication_explicit_replacement_uses_existing_document_and_new_ref owner="sync", database_url=DATABASE_URL, input_directory=layout.frozen_input_directory(), - strategy="baseline", + strategy="candidate", mode="warm", run_id=RUN_ID, replacement_document_id="doc_existing_123", @@ -439,7 +420,7 @@ def test_run_publication_persists_trace_for_enabled_sample( owner="sync", database_url=DATABASE_URL, input_directory=layout.frozen_input_directory(), - strategy="baseline", + strategy="candidate", mode="cold", run_id=RUN_ID, expectations=expectations_for(payload), @@ -484,7 +465,7 @@ def test_run_publication_accepts_a_freshly_reset_clone_for_cold_samples( owner="async", database_url=DATABASE_URL, input_directory=layout.frozen_input_directory(), - strategy="baseline", + strategy="candidate", mode="cold", run_id=RUN_ID, expectations=expectations_for(payload), @@ -514,7 +495,7 @@ def test_run_publication_refuses_a_clone_already_used_for_a_cold_sample( owner="sync", database_url=DATABASE_URL, input_directory=layout.frozen_input_directory(), - strategy="baseline", + strategy="candidate", mode="cold", run_id=RUN_ID, expectations=expectations_for(payload), @@ -539,7 +520,7 @@ def test_run_publication_executes_the_sync_owner_against_a_clone( owner="sync", database_url=database_url, input_directory=layout.frozen_input_directory(), - strategy="baseline", + strategy="candidate", mode="cold", run_id=RUN_ID, expectations=expectations_for(payload), @@ -594,7 +575,7 @@ def test_run_publication_executes_the_async_owner_against_a_clone( owner="async", database_url=database_url, input_directory=layout.frozen_input_directory(), - strategy="baseline", + strategy="candidate", mode="cold", run_id=RUN_ID, expectations=expectations_for(payload), @@ -642,12 +623,12 @@ def test_run_publication_rejects_shell_alternation_style_inputs() -> None: "--input", "inputs/spacex-s1-production", "--strategy", - "baseline", + "candidate", "--mode", "cold", ] ) - assert arguments.strategy == "baseline" + assert arguments.strategy == "candidate" assert arguments.mode == "cold" assert arguments.trace == "disabled" enabled = parser.parse_args( @@ -659,7 +640,7 @@ def test_run_publication_rejects_shell_alternation_style_inputs() -> None: "--input", "inputs/spacex-s1-production", "--strategy", - "baseline", + "candidate", "--mode", "cold", "--trace", @@ -694,7 +675,6 @@ def test_publication_environment_uses_only_benchmark_placeholders() -> None: assert environment["DATABASE_URL"] == ( "postgresql+psycopg2://postgres:secret@127.0.0.1:55450/db" ) - assert environment["KNOWHERE_PUBLICATION_STRATEGY"] == "candidate" assert environment["S3_BUCKET_NAME"] == "knowhere-publication-benchmark" assert environment["DS_URL"].endswith(".invalid/v1") assert environment["TELEMETRY_ENABLED"] == "false" @@ -732,7 +712,7 @@ def test_run_publication_applies_benchmark_environment( owner="async", database_url=DATABASE_URL, input_directory=layout.frozen_input_directory(), - strategy="baseline", + strategy="candidate", mode="cold", run_id=RUN_ID, expectations=expectations_for(payload), diff --git a/apps/api/tests/contract/test_publication_capacity_batch_contract.py b/apps/api/tests/contract/test_publication_capacity_batch_contract.py index 41110f6be..38874006d 100644 --- a/apps/api/tests/contract/test_publication_capacity_batch_contract.py +++ b/apps/api/tests/contract/test_publication_capacity_batch_contract.py @@ -16,7 +16,6 @@ get_contract_database_url, prepare_contract_storage, ) -from shared.core.config import settings from tests.support.publication_benchmark_support import ensure_benchmark_import_path ensure_benchmark_import_path() @@ -38,7 +37,7 @@ @pytest.mark.parametrize("owner", ("sync", "async")) -@pytest.mark.parametrize("strategy", ("baseline", "candidate")) +@pytest.mark.parametrize("strategy", ("candidate",)) def test_capacity_batch_converges_while_protected_retrieval_stays_stable( owner: str, strategy: str, @@ -47,8 +46,6 @@ def test_capacity_batch_converges_while_protected_retrieval_stays_stable( tmp_path: Path, ) -> None: configure_contract_environment(monkeypatch, postgresql_proc) - original_strategy = settings.KNOWHERE_PUBLICATION_STRATEGY - monkeypatch.setattr(settings, "KNOWHERE_PUBLICATION_STRATEGY", original_strategy) environment_keys = build_publication_environment( database_url="postgresql://benchmark.invalid/test", run_id="capacity-contract", diff --git a/apps/api/tests/contract/test_publication_gin_settings_contract.py b/apps/api/tests/contract/test_publication_gin_settings_contract.py index a6a987e70..5528fcb98 100644 --- a/apps/api/tests/contract/test_publication_gin_settings_contract.py +++ b/apps/api/tests/contract/test_publication_gin_settings_contract.py @@ -10,7 +10,6 @@ from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from sqlalchemy.orm import Session -from shared.core.config import settings from shared.services.retrieval.publication_service import RetrievalPublicationService from shared.testing.contract_runtime import ( PostgreSQLProcess, @@ -59,7 +58,7 @@ def _publish(db: Session, scope: PublicationScope) -> str: return str(db.execute(text("SHOW gin_pending_list_limit")).scalar_one()) -@pytest.mark.parametrize("strategy", ("baseline", "candidate")) +@pytest.mark.parametrize("strategy", ("candidate",)) def test_gin_limit_is_local_to_sync_publication_transaction( strategy: str, monkeypatch: pytest.MonkeyPatch, @@ -67,7 +66,6 @@ def test_gin_limit_is_local_to_sync_publication_transaction( ) -> None: configure_contract_environment(monkeypatch, postgresql_proc) asyncio.run(prepare_contract_storage()) - monkeypatch.setattr(settings, "KNOWHERE_PUBLICATION_STRATEGY", strategy) database_url = make_url(get_contract_database_url()).set( drivername="postgresql+psycopg2" ) @@ -79,14 +77,14 @@ def test_gin_limit_is_local_to_sync_publication_transaction( session.commit() default_limit = str(session.execute(text("SHOW gin_pending_list_limit")).scalar_one()) observed_limit = _publish(session, scope) - assert observed_limit == ("64MB" if strategy == "candidate" else default_limit) + assert observed_limit == "64MB" session.rollback() assert str(session.execute(text("SHOW gin_pending_list_limit")).scalar_one()) == default_limit finally: engine.dispose() -@pytest.mark.parametrize("strategy", ("baseline", "candidate")) +@pytest.mark.parametrize("strategy", ("candidate",)) async def test_gin_limit_is_local_to_async_publication_transaction( strategy: str, monkeypatch: pytest.MonkeyPatch, @@ -94,7 +92,6 @@ async def test_gin_limit_is_local_to_async_publication_transaction( ) -> None: configure_contract_environment(monkeypatch, postgresql_proc) await prepare_contract_storage() - monkeypatch.setattr(settings, "KNOWHERE_PUBLICATION_STRATEGY", strategy) database_url = make_url(get_contract_database_url()).set( drivername="postgresql+asyncpg" ) @@ -107,7 +104,7 @@ async def test_gin_limit_is_local_to_async_publication_transaction( await session.commit() default_limit = str((await session.execute(text("SHOW gin_pending_list_limit"))).scalar_one()) observed_limit = await session.run_sync(lambda db: _publish(db, scope)) - assert observed_limit == ("64MB" if strategy == "candidate" else default_limit) + assert observed_limit == "64MB" await session.rollback() assert str((await session.execute(text("SHOW gin_pending_list_limit"))).scalar_one()) == default_limit finally: diff --git a/apps/api/tests/contract/test_publication_namespace_move_contract.py b/apps/api/tests/contract/test_publication_namespace_move_contract.py index 49bc8d04e..70dc536ed 100644 --- a/apps/api/tests/contract/test_publication_namespace_move_contract.py +++ b/apps/api/tests/contract/test_publication_namespace_move_contract.py @@ -71,7 +71,7 @@ def _keyword_chunks(chunk_id: str) -> list[dict[str, object]]: ] -@pytest.mark.parametrize("strategy", ("baseline", "candidate")) +@pytest.mark.parametrize("strategy", ("candidate",)) def test_replacement_moves_graph_and_snapshot_in_one_transaction( strategy: str, monkeypatch: pytest.MonkeyPatch, @@ -79,7 +79,6 @@ def test_replacement_moves_graph_and_snapshot_in_one_transaction( ) -> None: configure_contract_environment(monkeypatch, postgresql_proc) asyncio.run(prepare_contract_storage()) - monkeypatch.setattr(settings, "KNOWHERE_PUBLICATION_STRATEGY", strategy) database_url = make_url(get_contract_database_url()).set( drivername="postgresql+psycopg2" ) @@ -176,7 +175,6 @@ def test_graph_publication_filters_peers_by_keyword_candidates( ) -> None: configure_contract_environment(monkeypatch, postgresql_proc) asyncio.run(prepare_contract_storage()) - monkeypatch.setattr(settings, "KNOWHERE_PUBLICATION_STRATEGY", "candidate") database_url = make_url(get_contract_database_url()).set( drivername="postgresql+psycopg2" ) @@ -241,7 +239,6 @@ def test_oversized_snapshot_uses_generation_mismatch_fallback( ) -> None: configure_contract_environment(monkeypatch, postgresql_proc) asyncio.run(prepare_contract_storage()) - monkeypatch.setattr(settings, "KNOWHERE_PUBLICATION_STRATEGY", "candidate") monkeypatch.setattr( settings, "KNOWHERE_PUBLICATION_NAMESPACE_SNAPSHOT_MAX_BYTES", 1 ) diff --git a/apps/api/tests/contract/test_publication_revision_replay_contract.py b/apps/api/tests/contract/test_publication_revision_replay_contract.py index 19f59c202..bc6a8faeb 100644 --- a/apps/api/tests/contract/test_publication_revision_replay_contract.py +++ b/apps/api/tests/contract/test_publication_revision_replay_contract.py @@ -9,7 +9,6 @@ from sqlalchemy.engine import make_url from sqlalchemy.orm import Session -from shared.core.config import settings from shared.models.database.document import Document from shared.models.database.job_result import JobResult from shared.services.retrieval.publication_service import RetrievalPublicationService @@ -52,7 +51,7 @@ def _chunks(scope: PublicationScope) -> list[dict[str, object]]: ] -@pytest.mark.parametrize("strategy", ("baseline", "candidate")) +@pytest.mark.parametrize("strategy", ("candidate",)) def test_completed_revision_replay_is_a_no_op( strategy: str, monkeypatch: pytest.MonkeyPatch, @@ -60,7 +59,6 @@ def test_completed_revision_replay_is_a_no_op( ) -> None: configure_contract_environment(monkeypatch, postgresql_proc) asyncio.run(prepare_contract_storage()) - monkeypatch.setattr(settings, "KNOWHERE_PUBLICATION_STRATEGY", strategy) database_url = make_url(get_contract_database_url()).set( drivername="postgresql+psycopg2" ) diff --git a/apps/api/tests/contract/test_publication_strategy_contract.py b/apps/api/tests/contract/test_publication_strategy_contract.py deleted file mode 100644 index d14eede08..000000000 --- a/apps/api/tests/contract/test_publication_strategy_contract.py +++ /dev/null @@ -1,115 +0,0 @@ -"""Contract tests for the temporary global Publication strategy setting.""" - -# ruff: noqa: E402 - -from __future__ import annotations - -import inspect -from pathlib import Path - -import pytest -from loguru import logger -from pydantic import ValidationError - -from tests.support.publication_benchmark_support import ( - REPOSITORY_ROOT, - ensure_benchmark_import_path, -) - -ensure_benchmark_import_path() - -from shared.core.config import settings # noqa: E402 -from shared.core.config.publication import ( # noqa: E402 - DEFAULT_PUBLICATION_NAMESPACE_SNAPSHOT_MAX_BYTES, - DEFAULT_PUBLICATION_STRATEGY, - PUBLICATION_STRATEGIES, - PUBLICATION_STRATEGY_ENV_VAR, - PublicationConfig, - PublicationStrategyError, - normalize_publication_strategy, -) -from shared.services.retrieval.publication_strategy import ( # noqa: E402 - PublicationStrategyAnnouncer, - announce_publication_strategy, - resolve_publication_strategy, -) - - -def test_publication_strategy_contract_defaults_to_baseline() -> None: - assert DEFAULT_PUBLICATION_STRATEGY == "baseline" - assert PUBLICATION_STRATEGIES == ("baseline", "candidate") - assert settings.KNOWHERE_PUBLICATION_STRATEGY in PUBLICATION_STRATEGIES - assert ( - settings.KNOWHERE_PUBLICATION_NAMESPACE_SNAPSHOT_MAX_BYTES - == DEFAULT_PUBLICATION_NAMESPACE_SNAPSHOT_MAX_BYTES - ) - - -def test_publication_snapshot_bound_accepts_zero_to_disable() -> None: - config = PublicationConfig( - KNOWHERE_PUBLICATION_STRATEGY="candidate", - KNOWHERE_PUBLICATION_NAMESPACE_SNAPSHOT_MAX_BYTES=0, - ) - assert config.KNOWHERE_PUBLICATION_NAMESPACE_SNAPSHOT_MAX_BYTES == 0 - - -def test_publication_config_rejects_unsupported_strategy_values() -> None: - config = PublicationConfig(KNOWHERE_PUBLICATION_STRATEGY="candidate") - assert config.KNOWHERE_PUBLICATION_STRATEGY == "candidate" - - with pytest.raises(ValidationError) as error: - PublicationConfig(KNOWHERE_PUBLICATION_STRATEGY="legacy") - - assert PUBLICATION_STRATEGY_ENV_VAR in str(error.value) - with pytest.raises(PublicationStrategyError): - normalize_publication_strategy("percent-50") - - -def test_strategy_resolver_accepts_no_override_argument() -> None: - assert inspect.signature(resolve_publication_strategy).parameters == {} - - -def test_strategy_resolver_reads_the_live_global_setting( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr(settings, "KNOWHERE_PUBLICATION_STRATEGY", "candidate") - assert resolve_publication_strategy() == "candidate" - - monkeypatch.setattr(settings, "KNOWHERE_PUBLICATION_STRATEGY", "legacy") - with pytest.raises(PublicationStrategyError): - resolve_publication_strategy() - - -def test_strategy_announcer_emits_one_event_per_process() -> None: - messages: list[str] = [] - sink_id = logger.add( - lambda message: messages.append(message.record["message"]), - level="INFO", - ) - try: - announcer = PublicationStrategyAnnouncer() - first = announcer.announce(service_name="knowhere-worker") - second = announcer.announce(service_name="knowhere-worker") - finally: - logger.remove(sink_id) - - assert first == second == settings.KNOWHERE_PUBLICATION_STRATEGY - assert len(messages) == 1 - assert "publication strategy" in messages[0] - assert PUBLICATION_STRATEGY_ENV_VAR in messages[0] - - -def test_strategy_announcement_singleton_is_exposed() -> None: - assert announce_publication_strategy is not None - assert resolve_publication_strategy() in PUBLICATION_STRATEGIES - - -def test_worker_and_api_announce_strategy_at_startup() -> None: - repository_root = Path(REPOSITORY_ROOT) - worker_source = ( - repository_root / "apps/worker/app/core/worker_bootstrap.py" - ).read_text(encoding="utf-8") - api_source = (repository_root / "apps/api/main.py").read_text(encoding="utf-8") - - assert "announce_publication_strategy" in worker_source - assert "announce_publication_strategy" in api_source diff --git a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py index 690221d1f..be4dc7046 100644 --- a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py +++ b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py @@ -89,6 +89,7 @@ async def test_classic_route_maps_winning_unit_to_one_chunk( assert results[0]["chunk_type"] == "text" assert results[0]["source"] == { "document_id": first["document_id"], + "job_result_id": first["job_result_id"], "source_file_name": "first.pdf", "section_path": "Root / Hit / body", } diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 708c13921..41b95e917 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -190,6 +190,7 @@ async def test_should_return_seeded_retrieval_results_for_the_authenticated_user assert results[0]["score"] == 1.0 assert results[0]["source"] == { "document_id": seeded_document["document_id"], + "job_result_id": seeded_document["job_result_id"], "source_file_name": "contract-retrieval.pdf", "section_path": "contract/intro", } @@ -492,11 +493,13 @@ async def test_agent_explore_retrieval_should_return_seeded_chunk_via_fake_episo assert any( ref.get("chunk_id") == target["chunk_id"] and ref.get("document_id") == target["document_id"] + and ref.get("job_result_id") == target["job_result_id"] for ref in referenced_chunks ) assert results[0]["content"] == "explore seeded EBITDA marker content" assert results[0]["source"] == { "document_id": target["document_id"], + "job_result_id": target["job_result_id"], "source_file_name": "target.pdf", "section_path": target["section_path"], } diff --git a/apps/api/tests/contract/test_shared_demo_corpus_contract.py b/apps/api/tests/contract/test_shared_demo_corpus_contract.py new file mode 100644 index 000000000..ef783f84e --- /dev/null +++ b/apps/api/tests/contract/test_shared_demo_corpus_contract.py @@ -0,0 +1,865 @@ +"""HTTP, PostgreSQL/RLS, and retrieval contracts for shared demo publication.""" + +from __future__ import annotations + +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager +from typing import Any +from uuid import uuid4 + +import pytest +from httpx import AsyncClient +from sqlalchemy import select, text +from sqlalchemy.exc import DBAPIError +from sqlalchemy.orm import Session + +from tests.support.contract_database import ContractDatabase + +ClientFactory = Callable[[], AbstractAsyncContextManager[AsyncClient]] + + +async def create_demo( + client: AsyncClient, + *, + source_id: str = "contract-demo", + file_name: str = "shared.docx", +) -> dict[str, Any]: + from shared.core.config import settings + + settings.DEMO_MAINTAINER_USER_IDS = "local-dev-user" + response = await client.post( + "/api/v2/jobs", + json={ + "namespace": "__knowhere_demo__", + "source_type": "file", + "file_name": file_name, + "data_id": source_id, + }, + ) + assert response.status_code == 200, response.text + return response.json() + + +async def publish_demo( + job: dict[str, Any], + *, + content: str = "Shared launch capabilities include reusable rockets.", + has_page: bool = False, + should_fail: bool = False, +) -> str: + from shared.core.database import get_db_context + from shared.models.database.job import Job + from shared.models.database.job_result import JobResult + from shared.services.retrieval.publication_service import ( + RetrievalPublicationService, + ) + + revisionId: str = str(uuid4()) + + def publish(db: Session) -> None: + record = db.execute(select(Job).where(Job.job_id == job["job_id"])).scalar_one() + result = JobResult( + id=revisionId, + job_id=record.job_id, + delivery_mode="url", + document_metadata={ + "source_size_bytes": 123, + "parse_track": "page_memory" if has_page else "chunk", + "demo_asset_manifest": ["images/rocket.png", "tables/launch.html"], + }, + ) + db.add(result) + db.flush() + chunks: list[dict[str, Any]] = [ + { + "chunk_id": "body-" + revisionId[:8], + "type": "page" if has_page else "text", + "content": content, + "path": "shared.docx/Capabilities", + "metadata": { + "page_nums": [1] if has_page else [], + "connect_to": [ + {"target": "image-" + revisionId[:8], "relation": "embeds"} + ], + }, + }, + { + "chunk_id": "image-" + revisionId[:8], + "type": "image", + "content": "Reusable rocket image", + "path": "images/rocket.png", + "metadata": {"file_path": "images/rocket.png"}, + }, + { + "chunk_id": "table-" + revisionId[:8], + "type": "table", + "content": "
Reusable launch
", + "path": "tables/launch.html", + "metadata": {"file_path": "tables/launch.html"}, + }, + ] + chunks.append( + { + "chunk_id": "filler-" + revisionId[:8], + "type": "text", + "content": "Independent astronomy paragraph about stars and planets.", + "path": "shared.docx/Astronomy", + "metadata": {}, + } + ) + service = RetrievalPublicationService() + outcome = service.publish_document_state( + db, job_id=record.job_id, job_result_id=revisionId, chunks=chunks + ) + assert outcome is not None and outcome.document_id == job["document_id"] + service.publish_document_graph( + db, job_id=record.job_id, job_result_id=revisionId + ) + if should_fail: + raise RuntimeError("Interrupted publication contract") + record.status = "done" + + async with get_db_context() as db: + await db.run_sync(publish) + await db.commit() + return revisionId + + +@pytest.mark.asyncio +async def test_preparing_readiness_and_removed_materialization( + developer_api_client_factory: ClientFactory, +) -> None: + async with developer_api_client_factory() as client: + catalog = await client.get("/api/v1/demo/catalog") + assert catalog.status_code == 200 and catalog.json()["sources"] == [] + assert catalog.headers["Cache-Control"] == "no-store" + assert all( + source["status"] == "planned" + for source in catalog.json()["official_library"]["sources"] + ) + query = await client.post( + "/api/v2/retrieval/query", + json={ + "namespace": "__knowhere_demo__", + "query": "launch", + "use_agentic": False, + }, + ) + assert ( + query.status_code == 503 + and query.json()["error"]["code"] == "DEMO_NOT_READY" + ) + original = await client.get("/api/v1/demo/sources/demo-spacex-s1/original") + assert ( + original.status_code == 409 + and original.json()["error"]["code"] == "DEMO_NOT_READY" + ) + removed = await client.post( + "/api/v1/demo/materializations", + json={"demo_source_ids": ["demo-spacex-s1"]}, + ) + assert ( + removed.status_code == 410 + and removed.json()["error"]["code"] == "DEMO_MATERIALIZATION_REMOVED" + ) + + +@pytest.mark.asyncio +async def test_demo_writes_require_maintainers( + developer_api_client_factory: ClientFactory, +) -> None: + async with developer_api_client_factory() as client: + from shared.core.config import settings + + settings.DEMO_MAINTAINER_USER_IDS = "" + for version in ("v1", "v2"): + response = await client.post( + f"/api/{version}/jobs", + json={ + "namespace": " __knowhere_demo__ ", + "source_type": "file", + "file_name": "shared.docx", + "data_id": "forbidden", + }, + ) + assert response.status_code == 403, response.text + archive = await client.post("/api/v1/documents/ddoc_forbidden/archive") + assert archive.status_code == 403, archive.text + from shared.core.database import get_db_context + + async with get_db_context() as db: + result = await db.execute( + text("UPDATE demo_documents SET title='Unauthorized'") + ) + assert result.rowcount == 0 + with pytest.raises(DBAPIError): + await db.execute( + text( + "INSERT INTO demo_documents (document_id, demo_source_id, user_id, namespace, status, title, category) VALUES ('ddoc_forbidden', 'forbidden', '__knowhere_demo__', '__knowhere_demo__', 'active', 'Unauthorized', 'other')" + ) + ) + await db.rollback() + + +@pytest.mark.asyncio +async def test_shared_lifecycle_reads_retrieval_history_and_global_cache( + developer_api_client_factory: ClientFactory, +) -> None: + async with developer_api_client_factory() as client: + job = await create_demo(client) + conflict = await client.post( + "/api/v2/jobs", + json={ + "namespace": "__knowhere_demo__", + "source_type": "file", + "file_name": "shared.docx", + "data_id": "contract-demo", + }, + ) + assert conflict.status_code == 409 + revisionId = await publish_demo(job) + catalog = (await client.get("/api/v1/demo/catalog")).json() + source = next( + source + for source in catalog["sources"] + if source["demo_source_id"] == "contract-demo" + ) + assert source["canonical_document_id"] == job["document_id"] + assert ( + source["job_result_id"] == revisionId + and source["chunk_count"] == 4 + and source["size_bytes"] == 123 + ) + readerKey: str = "sk_contract_reader_two" + await ContractDatabase.insert_authenticated_user( + user_id="reader-two", + api_key=readerKey, + credits_balance=10000000, + user_tier="tier_5", + ) + identities = [client.headers["Authorization"], "Bearer " + readerKey] + queryPayload = { + "namespace": "__knowhere_demo__", + "query": "capabilities", + "include_document_ids": [job["document_id"]], + "use_agentic": False, + "top_k": 1, + } + for identity in identities: + response = await client.post( + "/api/v2/retrieval/query", + json=queryPayload, + headers={"Authorization": identity}, + ) + assert response.status_code == 200, response.text + assert ( + response.json()["results"] + and response.json()["results"][0]["source"]["job_result_id"] + == revisionId + ) + document = await client.get( + f"/api/v2/documents/{job['document_id']}", + headers={"Authorization": identity}, + ) + assert ( + document.status_code == 200 + and document.json()["job_result_id"] == revisionId + ) + chunks = ( + await client.get(f"/api/v2/documents/{job['document_id']}/chunks") + ).json() + chunkId = chunks["chunks"][0]["id"] + update = await create_demo(client) + assert update["document_id"] == job["document_id"] + newRevision = await publish_demo( + update, content="Updated capabilities include orbital refueling." + ) + for identity in identities: + response = await client.post( + "/api/v2/retrieval/query", + json={**queryPayload, "query": "capabilities"}, + headers={"Authorization": identity}, + ) + assert response.status_code == 200, response.text + assert ( + response.json()["results"][0]["source"]["job_result_id"] == newRevision + ) + history = await client.get( + f"/api/v1/documents/{job['document_id']}/chunks/{chunkId}", + params={"job_result_id": revisionId}, + ) + assert ( + history.status_code == 200 + and "reusable rockets" in history.json()["chunk"]["content"] + ) + wrongRevision = await client.get( + f"/api/v1/documents/{job['document_id']}/chunks", + params={"job_result_id": str(uuid4())}, + ) + assert wrongRevision.status_code == 404 + failedUpdate = await create_demo(client) + with pytest.raises(RuntimeError, match="Interrupted publication"): + await publish_demo(failedUpdate, should_fail=True) + document = await client.get(f"/api/v1/documents/{job['document_id']}") + assert document.json()["job_result_id"] == newRevision + from shared.core.config import settings + + settings.DEMO_MAINTAINER_USER_IDS = "" + assert ( + await client.post(f"/api/v2/documents/{job['document_id']}/archive") + ).status_code == 403 + settings.DEMO_MAINTAINER_USER_IDS = "local-dev-user" + assert ( + await client.post(f"/api/v2/documents/{job['document_id']}/archive") + ).status_code == 200 + assert ( + await client.get( + f"/api/v1/documents/{job['document_id']}/chunks", + params={"job_result_id": revisionId}, + ) + ).status_code == 404 + + +@pytest.mark.asyncio +async def test_demo_agent_tools_pin_history_and_private_isolation( + developer_api_client_factory: ClientFactory, +) -> None: + async with developer_api_client_factory() as client: + job = await create_demo(client, source_id="agent-demo") + revision = await publish_demo(job, has_page=True) + historyChunks = ( + await client.get(f"/api/v1/documents/{job['document_id']}/chunks") + ).json()["chunks"] + sectionPath = next( + chunk["section_path"] + for chunk in historyChunks + if chunk["chunk_type"] == "page" + ) + from shared.core.database import get_db_context + from shared.services.retrieval.execution.revision_pins import ( + capture_revision_pins, + ) + from shared.services.retrieval.corpus_revision_context import ( + CorpusRevisionContext, + ) + from shared.services.retrieval.agent_explore.dispatch import dispatch_tool_call + from shared.services.retrieval.agent_explore.ref_resolution import ( + resolve_finish_refs, + ) + + async with get_db_context() as db: + pins = await capture_revision_pins( + db, user_id="reader-two", namespace="__knowhere_demo__" + ) + update = await create_demo(client, source_id="agent-demo") + await publish_demo( + update, content="Updated orbital refueling evidence", has_page=True + ) + with CorpusRevisionContext.bind(pins): + calls = [ + ("corpus.list_documents", {}), + ("corpus.outline", {"scope": [{"document_id": job["document_id"]}]}), + ( + "corpus.grep", + { + "pattern": "reusable rockets", + "scope": [{"document_id": job["document_id"]}], + }, + ), + ( + "corpus.recall", + { + "query": "rockets", + "scope": [{"document_id": job["document_id"]}], + }, + ), + ("corpus.assets", {"scope": [{"document_id": job["document_id"]}]}), + ("corpus.neighbors", {"document_id": job["document_id"]}), + ( + "corpus.read", + { + "refs": [ + { + "document_id": job["document_id"], + "chunk_id": "body-" + revision[:8], + } + ] + }, + ), + ] + for name, arguments in calls: + result = await dispatch_tool_call( + name, + arguments, + db_factory=get_db_context, + user_id="reader-two", + namespace="__knowhere_demo__", + ) + assert result.error is None, (name, result.error) + assert "Updated orbital" not in result.text, (name, result.text) + if name == "corpus.recall": + assert any( + ref.get("chunk_id") == "body-" + revision[:8] + for ref in result.refs + ), result.text + if name == "corpus.read": + assert "reusable rockets" in result.text + if name == "corpus.grep": + assert result.payload.get("total_matches", 0) > 0 or result.refs + async with get_db_context() as db: + refs = await resolve_finish_refs( + db, + user_id="reader-two", + namespace="__knowhere_demo__", + refs=[ + {"document_id": job["document_id"], "section_path": sectionPath} + ], + ) + assert ( + refs.resolved + and refs.resolved[0]["chunk_id"] == "body-" + revision[:8] + ) + private = await client.post( + "/api/v2/jobs", + json={ + "namespace": "personal", + "source_type": "file", + "file_name": "shared.docx", + "data_id": "private-demo-copy", + }, + ) + assert private.status_code == 200, private.text + await publish_demo(private.json()) + readerKey = "sk_private_isolation_reader" + await ContractDatabase.insert_authenticated_user( + user_id="isolated-reader", + api_key=readerKey, + credits_balance=10000000, + user_tier="tier_5", + ) + ownDocument = await client.get( + f"/api/v1/documents/{private.json()['document_id']}" + ) + assert ownDocument.status_code == 200 + denied = await client.get( + f"/api/v1/documents/{private.json()['document_id']}", + headers={"Authorization": "Bearer " + readerKey}, + ) + assert denied.status_code == 404 + shared = await client.post( + "/api/v2/retrieval/query", + json={ + "namespace": "__knowhere_demo__", + "query": "orbital", + "top_k": 100, + "use_agentic": False, + }, + headers={"Authorization": "Bearer " + readerKey}, + ) + assert shared.status_code == 200 and shared.json()["results"] + assert all( + row["source"]["document_id"] != private.json()["document_id"] + for row in shared.json()["results"] + ) + + +@pytest.mark.asyncio +async def test_demo_revision_media_validation_and_job_retention( + developer_api_client_factory: ClientFactory, monkeypatch: pytest.MonkeyPatch +) -> None: + async with developer_api_client_factory() as client: + job = await create_demo( + client, source_id="media-demo", file_name="original.pdf" + ) + revision = await publish_demo(job, has_page=True) + from shared.services.storage.job_file_storage import JobFileStorage + + calls: list[dict[str, Any]] = [] + + def sign(key: str, **arguments: Any) -> str: + calls.append({"key": key, **arguments}) + return "https://storage.example.test/signed-object" + + class SignedStorage: + generate_presigned_url = staticmethod(sign) + + monkeypatch.setattr( + JobFileStorage, "storage_adapter", property(lambda self: SignedStorage()) + ) + for path, mime in ( + ("original", "application/pdf"), + ("assets/images/rocket.png", "image/png"), + ("assets/tables/launch.html", "text/html"), + ): + response = await client.get( + f"/api/v1/demo/sources/media-demo/{path}", + params={"job_result_id": revision}, + headers={"Range": "bytes=0-10"}, + follow_redirects=False, + ) + assert response.status_code == 307, response.text + assert calls[-1]["headers"]["Content-Type"] == mime + assert calls[-1]["headers"]["Content-Disposition"].startswith("inline;") + for path in ( + "images/not-listed.png", + "images/%2e%2e/secret.png", + "%2fimages/rocket.png", + "images/rocket.png/extra", + ): + response = await client.get( + f"/api/v1/demo/sources/media-demo/assets/{path}", + params={"job_result_id": revision}, + ) + assert response.status_code == 404, response.text + chunks = await client.get( + "/api/v1/demo/sources/media-demo/chunks", + params={"page_size": 1, "job_result_id": revision}, + ) + assert ( + chunks.status_code == 200 + and len(chunks.json()["chunks"]) == 1 + and chunks.json()["pagination"]["total"] == 4 + ) + chunk = await client.get( + f"/api/v1/demo/sources/media-demo/chunks/{chunks.json()['chunks'][0]['id']}", + params={"job_result_id": revision}, + ) + assert chunk.status_code == 200 and chunk.json()["job_result_id"] == revision + for version in ("v1", "v2"): + deletion = await client.delete(f"/api/{version}/jobs/{job['job_id']}") + assert ( + deletion.status_code == 409 + and "retained shared demo revision" in deletion.text + ) + from shared.core.database import get_db_context + + async with get_db_context() as db: + with pytest.raises(DBAPIError): + await db.execute( + text("DELETE FROM jobs WHERE job_id=:identity"), + {"identity": job["job_id"]}, + ) + await db.rollback() + + +@pytest.mark.asyncio +async def test_demo_source_claims_worker_validation_and_interrupted_batch_roll_back( + developer_api_client_factory: ClientFactory, monkeypatch: pytest.MonkeyPatch +) -> None: + async with developer_api_client_factory() as client: + import asyncio + + # Authorize before concurrent HTTP calls so the same source lock is + # exercised by distinct real database transactions. + from shared.core.config import settings + + settings.DEMO_MAINTAINER_USER_IDS = "local-dev-user" + responses = await asyncio.gather( + *( + client.post( + "/api/v2/jobs", + json={ + "namespace": "__knowhere_demo__", + "source_type": "file", + "file_name": "shared.docx", + "data_id": "concurrent-demo", + }, + ) + for _ in range(2) + ) + ) + assert sorted(response.status_code for response in responses) == [200, 409] + job = next( + response.json() for response in responses if response.status_code == 200 + ) + renewed = await client.post(f"/api/v2/jobs/{job['job_id']}/upload-url") + assert renewed.status_code == 200 and renewed.json()["upload_url"] + settings.DEMO_MAINTAINER_USER_IDS = "" + for suffix in ("upload-url", "confirm-upload"): + denied = await client.post( + f"/api/v2/jobs/{job['job_id']}/{suffix}", json={} + ) + assert denied.status_code == 403, denied.text + settings.DEMO_MAINTAINER_USER_IDS = "local-dev-user" + from shared.core.database import get_db_context + from shared.models.database.job import Job + from shared.models.database.job_result import JobResult + from shared.models.database.demo_corpus import DemoDocument, DemoDocumentChunk + from shared.services.retrieval.demo_job_scope import validate_job_corpus + from shared.services.retrieval.publication_service import ( + RetrievalPublicationService, + ) + from sqlalchemy import event, func + + insertedBatches: list[int] = [] + + def interrupt_chunk_batch( + connection: object, + cursor: object, + statement: str, + parameters: object, + context: object, + executemany: bool, + ) -> None: + if statement.startswith("INSERT INTO demo_document_chunks"): + insertedBatches.append(1) + if len(insertedBatches) == 2: + raise RuntimeError("Second chunk batch interrupted") + + async with get_db_context() as db: + + def publish(session: Session) -> None: + record = session.execute( + select(Job).where(Job.job_id == job["job_id"]) + ).scalar_one() + malformed = dict(record.job_metadata or {}) + record.job_metadata = {**malformed, "corpus_target": "PRIVATE"} + from shared.core.exceptions.domain_exceptions import ( + PermissionDeniedException, + ) + + with pytest.raises(PermissionDeniedException): + validate_job_corpus(session, job=record) + record.job_metadata = malformed + session.add( + JobResult( + id=str(uuid4()), job_id=job["job_id"], delivery_mode="url" + ) + ) + session.flush() + revision = session.execute( + select(JobResult.id).where(JobResult.job_id == job["job_id"]) + ).scalar_one() + chunks: list[dict[str, Any]] = [ + { + "chunk_id": f"batch-{index}", + "type": "text", + "content": f"Unique publication evidence {index}", + "path": f"shared.docx/Section {index}", + "metadata": {}, + } + for index in range(1001) + ] + connection = session.connection() + event.listen(connection, "before_cursor_execute", interrupt_chunk_batch) + try: + RetrievalPublicationService().publish_document_state( + session, + job_id=job["job_id"], + job_result_id=revision, + chunks=chunks, + ) + finally: + event.remove( + connection, "before_cursor_execute", interrupt_chunk_batch + ) + + with pytest.raises(RuntimeError, match="Second chunk batch"): + await db.run_sync(publish) + await db.rollback() + assert len(insertedBatches) == 2 + async with get_db_context() as db: + count = ( + await db.execute( + select(func.count()) + .select_from(DemoDocumentChunk) + .where(DemoDocumentChunk.document_id == job["document_id"]) + ) + ).scalar_one() + document = ( + await db.execute( + select(DemoDocument).where( + DemoDocument.document_id == job["document_id"] + ) + ) + ).scalar_one() + assert count == 0 and document.current_job_result_id is None + missing = await client.post( + "/api/v2/retrieval/query", + json={ + "namespace": "__knowhere_demo__", + "query": "launch", + "include_document_ids": [job["document_id"]], + "use_agentic": False, + }, + ) + assert missing.status_code == 409 + + +@pytest.mark.asyncio +async def test_demo_filters_unpublished_revisions_examples_and_caller_stats( + developer_api_client_factory: ClientFactory, +) -> None: + async with developer_api_client_factory() as client: + job = await create_demo(client, source_id="filtered-demo") + from shared.core.database import get_db_context + from shared.models.database.demo_corpus import DemoDocument + from shared.models.database.job import Job + from shared.models.database.job_result import JobResult + from shared.services.retrieval.demo_authorization import ( + authorize_demo_transaction, + authorize_demo_caller, + ) + from shared.services.retrieval.stats.service import record_retrieval_hits + + async with get_db_context() as db: + await db.run_sync( + lambda session: authorize_demo_transaction( + session, user_id="local-dev-user" + ) + ) + document = ( + await db.execute( + select(DemoDocument).where( + DemoDocument.document_id == job["document_id"] + ) + ) + ).scalar_one() + document.catalog_metadata = { + "examples": [ + { + "id": "unique", + "question": "Which capability?", + "answer": "Reusable rockets.", + "citations": [ + { + "content": "Shared launch capabilities include reusable rockets." + } + ], + }, + { + "id": "missing", + "question": "Missing?", + "citations": [{"content": "Missing curated evidence"}], + }, + { + "id": "ambiguous", + "question": "Reusable?", + "citations": [{"content": "Reusable"}], + }, + ] + } + await db.commit() + revision = await publish_demo(job) + catalog = (await client.get("/api/v1/demo/catalog")).json() + source = next( + row + for row in catalog["sources"] + if row["demo_source_id"] == "filtered-demo" + ) + assert [example["id"] for example in source["examples"]] == ["unique"] + assert source["examples"][0]["citations"][0]["job_result_id"] == revision + chunks = ( + await client.get(f"/api/v2/documents/{job['document_id']}/chunks") + ).json()["chunks"] + section = next( + row["section_path"] + for row in chunks + if row["chunk_type"] == "text" and "capabilities" in row["content"] + ) + payload = { + "namespace": "__knowhere_demo__", + "query": "capabilities", + "use_agentic": False, + "top_k": 1, + "include_document_ids": [job["document_id"]], + } + for filters in ( + {"include_document_ids": []}, + {"exclude_document_ids": [job["document_id"]]}, + { + "exclude_sections": [ + {"document_id": job["document_id"], "section_path": section} + ] + }, + {"signal_paths": ["unmatched-section"], "filter_mode": "keep"}, + {"query": "nonexistentuniquephrase"}, + ): + result = await client.post( + "/api/v2/retrieval/query", json={**payload, **filters} + ) + assert result.status_code == 200 and result.json()["results"] == [], ( + result.text + ) + for chunk_type in ("text", "image", "table"): + result = await client.get( + f"/api/v1/documents/{job['document_id']}/chunks", + params={"chunk_type": chunk_type}, + ) + assert result.status_code == 200 + assert result.json()["chunks"] and all( + row["chunk_type"] == chunk_type for row in result.json()["chunks"] + ) + pending = await create_demo(client, source_id="unpublished-demo") + staged_revision = str(uuid4()) + async with get_db_context() as db: + record = ( + await db.execute(select(Job).where(Job.job_id == pending["job_id"])) + ).scalar_one() + record.status = "failed" + db.add( + JobResult( + id=staged_revision, + job_id=record.job_id, + demo_document_id=pending["document_id"], + delivery_mode="url", + ) + ) + await db.commit() + assert ( + await client.get( + f"/api/v2/documents/{pending['document_id']}/chunks", + params={"job_result_id": staged_revision}, + ) + ).status_code == 404 + query = await client.post( + "/api/v2/retrieval/query", + json={**payload, "include_document_ids": [pending["document_id"]]}, + ) + assert ( + query.status_code == 409 + and query.json()["error"]["code"] == "DEMO_NOT_READY" + ) + wrong_source = await client.get( + f"/api/v2/documents/{job['document_id']}/chunks", + params={"job_result_id": staged_revision}, + ) + assert wrong_source.status_code == 404 + retry = await create_demo(client, source_id="unpublished-demo") + assert retry["document_id"] == pending["document_id"] + async with get_db_context() as db: + await record_retrieval_hits( + db, + user_id="actual-reader", + namespace="__knowhere_demo__", + results=[ + { + "document_id": job["document_id"], + "chunk_id": "body-" + revision[:8], + } + ], + ) + await db.commit() + async with get_db_context() as db: + assert not ( + await db.execute(text("SELECT user_id FROM demo_retrieval_hit_stats")) + ).all() + await db.run_sync( + lambda session: authorize_demo_caller(session, user_id="actual-reader") + ) + assert { + row[0] + for row in ( + await db.execute( + text("SELECT user_id FROM demo_retrieval_hit_stats") + ) + ).all() + } == {"actual-reader"} + with pytest.raises(DBAPIError): + await db.execute( + text( + "INSERT INTO demo_retrieval_hit_stats (id, user_id, namespace, hit_kind, document_id, hit_count, created_at, updated_at) VALUES ('forged-hit', 'other-reader', '__knowhere_demo__', 'document', :document, 1, now(), now())" + ), + {"document": job["document_id"]}, + ) + await db.rollback() diff --git a/apps/api/tests/migrations/test_legacy_token_lookup_retirement_contract.py b/apps/api/tests/migrations/test_legacy_token_lookup_retirement_contract.py index f3ff96473..5295b5eed 100644 --- a/apps/api/tests/migrations/test_legacy_token_lookup_retirement_contract.py +++ b/apps/api/tests/migrations/test_legacy_token_lookup_retirement_contract.py @@ -50,6 +50,10 @@ def test_should_keep_default_indexes_then_apply_and_restore_at_same_head( ) -> None: config: Config = _build_config() command.upgrade(config, "heads") + with alembic_engine.connect() as connection: + initialRevision: str = str( + connection.scalar(text("SELECT version_num FROM alembic_version")) + ) initial: dict[str, dict[str, str | bool]] = _inspect_indexes(alembic_engine) assert len(initial) == 5 assert _LEGACY_INDEX in initial @@ -64,7 +68,7 @@ def test_should_keep_default_indexes_then_apply_and_restore_at_same_head( name: state for name, state in initial.items() if name != _LEGACY_INDEX } with alembic_engine.connect() as connection: - assert connection.scalar(text("SELECT version_num FROM alembic_version")) == _REVISION + assert connection.scalar(text("SELECT version_num FROM alembic_version")) == initialRevision for _ in range(2): result = _run_operator("--restore") assert result.returncode == 0, result.stderr diff --git a/apps/api/tests/migrations/test_schema_contract.py b/apps/api/tests/migrations/test_schema_contract.py index ff6ed485f..9c0438184 100644 --- a/apps/api/tests/migrations/test_schema_contract.py +++ b/apps/api/tests/migrations/test_schema_contract.py @@ -427,7 +427,10 @@ def test_should_restore_and_reapply_compact_token_indexes( external_connection: bool, ) -> None: config = _build_alembic_command_config(engine=alembic_engine) - _upgrade_to_heads(engine=alembic_engine) + # Exercise the reversible token-index chain independently of the retained + # demo schema, whose rollback deliberately preserves revisions and assets. + tokenIndexRevision: str = "3c4d5e6f7a8b" + command.upgrade(config, tokenIndexRevision) if external_connection: with alembic_engine.begin() as connection: @@ -458,9 +461,9 @@ def test_should_restore_and_reapply_compact_token_indexes( if external_connection: with alembic_engine.begin() as connection: config.attributes["connection"] = connection - command.upgrade(config, "heads") + command.upgrade(config, tokenIndexRevision) else: - command.upgrade(config, "heads") + command.upgrade(config, tokenIndexRevision) with alembic_engine.begin() as connection: compact_names = { str(row[0]) diff --git a/apps/worker/.env.example b/apps/worker/.env.example index 91020f97f..438c866ff 100644 --- a/apps/worker/.env.example +++ b/apps/worker/.env.example @@ -28,7 +28,7 @@ TMP_PATH=/tmp/knowhere LOGFIRE_TOKEN= # Required for local startup: database -DATABASE_URL=postgresql+asyncpg://root:root123@localhost:5432/Knowhere +DATABASE_URL=postgresql+asyncpg://knowhere_runtime:runtime123@localhost:5432/Knowhere DB_SSL_MODE=disable # DB_SSL_MODE options: disable, allow, prefer, require, verify-ca, verify-full # DB_SSL_CERT=/path/to/client-cert.pem @@ -148,3 +148,9 @@ ENTITY_TYPES=person,location,organization # Re-summarize detected page-track figures/charts/tables with numeric extraction # (ยง4.3). Gated separately from detection for safe rollout; defaults off. + +# Shared demo maintainers still require normal API write permission. +DEMO_MAINTAINER_USER_IDS= +# API/Worker DATABASE_URL must use NOSUPERUSER NOBYPASSRLS. +# Use a separate privileged role for migrations. +MIGRATION_DATABASE_URL= diff --git a/apps/worker/app/core/worker_bootstrap.py b/apps/worker/app/core/worker_bootstrap.py index 31778a80a..2d1f8b411 100644 --- a/apps/worker/app/core/worker_bootstrap.py +++ b/apps/worker/app/core/worker_bootstrap.py @@ -10,9 +10,6 @@ from shared.core.celery_app import celery_app from shared.core.logging import setup_logging -from shared.services.retrieval.publication_strategy import ( - announce_publication_strategy, -) from shared.services.worker_health import start_worker_heartbeat, stop_worker_heartbeat _CHILD_PROCESS_TERM_TIMEOUT_SECONDS: float = 5 @@ -47,7 +44,15 @@ def _stop_child_process( def init_worker(**kwargs: object) -> None: """Initialize structured logging and sync Redis when worker process starts.""" setup_logging(service_name="knowhere-worker") - announce_publication_strategy("knowhere-worker") + from shared.core.database_sync import get_sync_db_context + from shared.services.retrieval.demo_authorization import verify_runtime_database_role + with get_sync_db_context() as database: + try: + verify_runtime_database_role(database) + except RuntimeError as error: + # Celery logs and suppresses ordinary signal receiver exceptions. + # Exit the process so a bypassing role cannot consume any tasks. + raise SystemExit(str(error)) from error start_worker_heartbeat() # Do not cancel gevent tasks from Celery's reconnect or shutdown lifecycle. @@ -142,6 +147,11 @@ def run_worker() -> None: lock, while Kombu's broker mutex protects the restoration transaction. """ from shared.core.config import settings + from shared.core.database_sync import get_sync_db_context + from shared.services.retrieval.demo_authorization import verify_runtime_database_role + + with get_sync_db_context() as database: + verify_runtime_database_role(database) _register_task_modules() diff --git a/apps/worker/app/services/document_ingestion/processing_context.py b/apps/worker/app/services/document_ingestion/processing_context.py index 0d8e51728..f21940b0f 100644 --- a/apps/worker/app/services/document_ingestion/processing_context.py +++ b/apps/worker/app/services/document_ingestion/processing_context.py @@ -12,6 +12,7 @@ NotFoundException, ) from shared.models.database.job import Job +from shared.services.retrieval.demo_job_scope import validate_job_corpus from shared.services.redis.redis_sync_service import ( SyncJobInfoRedisService, SyncJobMetadataService, @@ -57,6 +58,12 @@ def load_parse_job_context( requested_user_id: str | None, redis_service: Any, ) -> ParseJobContext: + with get_sync_db_context() as authorizationDb: + durableJob: Job | None = authorizationDb.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() + if durableJob is None: + raise NotFoundException(resource="Job", resource_id=job_id) + validate_job_corpus(authorizationDb, job=durableJob) + durableScope: dict[str, object] = {key: value for key, value in (durableJob.job_metadata or {}).items() if key in ("namespace", "document_id", "demo_source_id", "corpus_target")} job_info_service = SyncJobInfoRedisService(redis_service) job_info = job_info_service.get_job_info(job_id) job_row: Job | None = None @@ -120,6 +127,7 @@ def load_parse_job_context( internal_message=f"Job metadata not found for job_id={job_id}", ) + raw_job_metadata.update(durableScope) return ParseJobContext( job_metadata=dict(raw_job_metadata), job_user_id=job_user_id, diff --git a/apps/worker/app/services/document_ingestion/source_preparation.py b/apps/worker/app/services/document_ingestion/source_preparation.py index 1aa6cb590..ca27344c8 100644 --- a/apps/worker/app/services/document_ingestion/source_preparation.py +++ b/apps/worker/app/services/document_ingestion/source_preparation.py @@ -56,6 +56,8 @@ def prepare_source_file( suffix=file_extension, temp_dir=input_dir, ) + if job_context.job_metadata.get("corpus_target") == "DEMO": + persist_job_metadata_updates(job_id=job_id, job_context=job_context, metadata_updates={"source_size_bytes": os.path.getsize(local_file_path)}) logger.info(f"File downloaded: job_id={job_id}, local_path={local_file_path}") prepared_parse_input = prepare_internal_parse_input( diff --git a/apps/worker/app/services/document_ingestion/success_finalization.py b/apps/worker/app/services/document_ingestion/success_finalization.py index 42330dbcc..0facf5a52 100644 --- a/apps/worker/app/services/document_ingestion/success_finalization.py +++ b/apps/worker/app/services/document_ingestion/success_finalization.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +import logfire from collections.abc import Callable from datetime import datetime, timezone from typing import Any @@ -26,7 +27,7 @@ from shared.models.schemas.job_metadata import JobMetadataHelper from shared.services.ai.token_tracking import get_current_token_tracker -from app.services.document_parser.support.stage_profiler import get_current_stage_tracker +from app.services.document_parser.support.stage_profiler import get_current_stage_tracker, stage_timer from shared.services.storage.result_storage import ResultStorage, get_result_storage from shared.services.storage.zip_result_service import ZipResultService @@ -82,7 +83,9 @@ def finalize_parse_success( generated_package=generated_package, job_id=job_id, result_storage_factory=result_storage_factory, + job_context=job_context, ) + _refresh_processing_stages(job_context) stored_count = 0 finalization_response = lifecycle_service.finalize_job_success( @@ -252,15 +255,19 @@ def _upload_result_package( generated_package: GeneratedResultPackage, job_id: str, result_storage_factory: ResultStorageFactory, + job_context: ParseJobContext | None = None, ) -> str: artifact_refs = collect_referenced_artifact_refs(result_package.chunks) add_dir = str(result_package.artifact.add_dir) if result_package.artifact.add_dir else "" if add_dir and os.path.isfile(os.path.join(add_dir, "source.pdf")): artifact_refs.add("source.pdf") - result_bundle = result_storage_factory().upload( - job_id=job_id, - result_dir=add_dir, - zip_file_path=generated_package.zip_file_path, - artifact_refs=artifact_refs, - ) + with logfire.span("Upload result assets for {job_id}", job_id=job_id), stage_timer("worker.result.assets_upload", job_id=job_id): + result_bundle = result_storage_factory().upload( + job_id=job_id, + result_dir=add_dir, + zip_file_path=generated_package.zip_file_path, + artifact_refs=artifact_refs, + ) + if job_context is not None and job_context.job_metadata.get("corpus_target") == "DEMO": + persist_job_metadata_updates(job_id=job_id, job_context=job_context, metadata_updates={"demo_asset_manifest": sorted(result_bundle.raw_files), "result_raw_prefix": result_bundle.raw_prefix}) return result_bundle.zip_key diff --git a/apps/worker/app/services/page_memory/fine_hierarchy.py b/apps/worker/app/services/page_memory/fine_hierarchy.py index 287e8d542..f34b9d8fb 100644 --- a/apps/worker/app/services/page_memory/fine_hierarchy.py +++ b/apps/worker/app/services/page_memory/fine_hierarchy.py @@ -14,7 +14,6 @@ from __future__ import annotations -import json import os import re from typing import Any @@ -23,11 +22,10 @@ from app.services.document_parser.structure.body_boundary import normalize_match_text from app.services.page_memory.page_tagger import PageTagResult +from app.services.page_memory.hierarchy_level_resolver import PageHierarchyLevelResolver from app.services.page_memory.skeleton_extractor import SectionSkeleton from app.services.page_memory._utils import page_scope_info, sort_skeletons from shared.services.ai.llm_overrides import get_text_client -from shared.services.ai.prompt_service import build_prompt -from shared.services.ai.response_process_service import eval_response from shared.services.chunks.path_segments import append_document_path @@ -91,6 +89,10 @@ def refine_fat_leaf_skeletons( ) if deeper: + # The first observed heading may start after the coarse section's + # opening body pages. Keep its parent so assembly owns that prefix. + if deeper[0].start_page > skeleton.start_page: + refined.append(skeleton) refined.extend(deeper) else: refined.append(skeleton) @@ -306,70 +308,20 @@ def _run_hierarchy_on_candidates( if not candidates: return None - input_json = json.dumps( - [ - { - "id": cand["id"], - "page": cand["page"], - "prominence": cand.get("prominence"), - "heading": cand["heading"], - } - for cand in candidates - ], - ensure_ascii=False, - indent=2, - ) - coarse_context = ( + coarse_context: str = ( f"title={skeleton.title}\n" f"path={skeleton.section_path}\n" f"pages={skeleton.start_page}-{skeleton.end_page}" ) - - try: - prompt, temperature, _top_p, prompt_max_tokens = build_prompt( - "page-memory-hierarchy", - input_json, - "", - paras={ - "max_depth": max_depth, - "max_tokens": max_tokens, - "coarse_context": coarse_context, - }, - ) - resolved_model = ( - model_name - or os.environ.get( - "HIERARCHY_LLM_MODEL", - os.environ.get("NORMOL_MODEL"), - ) - ) - client, resolved_model = get_text_client(requested_model=resolved_model) - answer = client.chat_completion( - messages=[ - {"role": "system", "content": "you are a document structure expert"}, - {"role": "user", "content": prompt}, - ], - model=resolved_model, - max_tokens=prompt_max_tokens, - temperature=temperature, - usage_task="page_memory.hierarchy", - ) - result = eval_response(answer) - except Exception as exc: - logger.warning( - "[page_memory.fine_hierarchy] LLM failed for skeleton {}: {}", - skeleton.section_path, - exc, - ) - return None - - levels_by_id = _parse_hierarchy_result(result, max_depth=max_depth) - if not levels_by_id: - logger.warning( - "[page_memory.fine_hierarchy] empty hierarchy result for skeleton {}", - skeleton.section_path, - ) - return None + resolved_model: str | None = model_name or os.environ.get( + "HIERARCHY_LLM_MODEL", os.environ.get("NORMOL_MODEL"), + ) + client, resolved_model = get_text_client(requested_model=resolved_model) + resolver: PageHierarchyLevelResolver = PageHierarchyLevelResolver( + client=client, model=resolved_model, max_tokens=max_tokens, + max_depth=max_depth, coarse_context=coarse_context, + ) + levels_by_id: dict[int, int] = resolver.resolve_levels(candidates) # Keep candidate order (== id order == the sequence fed to the LLM). ordered: list[dict[str, Any]] = [] @@ -476,28 +428,6 @@ def _title_key(title: str | None) -> str: return normalized -def _parse_hierarchy_result(result: Any, *, max_depth: int) -> dict[int, int]: - if not isinstance(result, list): - logger.warning( - "[page_memory.fine_hierarchy] unexpected hierarchy response type: {}", - type(result).__name__, - ) - return {} - parsed: dict[int, int] = {} - for item in result: - if not isinstance(item, dict): - continue - try: - row_id = int(item["id"]) - level = int(item["level"]) - except (KeyError, TypeError, ValueError): - continue - if level < 1: - continue - parsed[row_id] = min(level, max_depth) - return parsed - - def _record_trace_hierarchy( *, trace_recorder: Any | None, diff --git a/apps/worker/app/services/page_memory/hierarchy_level_resolver.py b/apps/worker/app/services/page_memory/hierarchy_level_resolver.py new file mode 100644 index 000000000..658c4cafb --- /dev/null +++ b/apps/worker/app/services/page_memory/hierarchy_level_resolver.py @@ -0,0 +1,204 @@ +"""Bound provider output while preserving hierarchy context between batches.""" + +from __future__ import annotations + +import json +from collections.abc import Mapping, Sequence +from typing import Protocol + +from loguru import logger +from openai.types.chat import ChatCompletionMessageParam + +from shared.services.ai.prompt_service import build_prompt + + +class _HierarchyClient(Protocol): + def chat_completion( + self, + *, + messages: list[ChatCompletionMessageParam], + model: str | None, + max_tokens: int, + temperature: float, + usage_task: str, + ) -> str: ... + + +class _HierarchyResponseError(ValueError): + """The provider did not classify exactly the requested candidate IDs.""" + + +class PageHierarchyLevelResolver: + """Classify every heading or noise row, with bounded recovery on bad output. + + Ancestors retain their relative levels across batches. Invalid responses + are retried in smaller batches, never interpreted as an empty hierarchy. + Exhausted recovery fails the import so a published revision stays intact. + """ + + def __init__( + self, + *, + client: _HierarchyClient, + model: str | None, + max_tokens: int, + max_depth: int, + coarse_context: str, + ) -> None: + self._client: _HierarchyClient = client + self._model: str | None = model + self._max_tokens: int = max_tokens + self._max_depth: int = max_depth + self._coarse_context: str = coarse_context + + def resolve_levels( + self, candidates: Sequence[Mapping[str, object]] + ) -> dict[int, int]: + # A response row needs ID, level and JSON punctuation. Leave ample + # completion headroom; cap input rows as well as the token estimate. + batch_size: int = max(1, min(64, self._max_tokens // 32)) + levels: dict[int, int] = {} + for start in range(0, len(candidates), batch_size): + batch: Sequence[Mapping[str, object]] = candidates[ + start : start + batch_size + ] + levels.update( + self._resolve_batch( + batch, + preceding_candidates=candidates[:start], + preceding_levels=levels, + ) + ) + return levels + + def _resolve_batch( + self, + candidates: Sequence[Mapping[str, object]], + *, + preceding_candidates: Sequence[Mapping[str, object]], + preceding_levels: Mapping[int, int], + can_retry: bool = True, + ) -> dict[int, int]: + prompt: str + temperature: float + completion_budget: int + context: str = self._build_context(preceding_candidates, preceding_levels) + prompt, temperature, _, completion_budget = build_prompt( + "page-memory-hierarchy", + json.dumps(list(candidates), ensure_ascii=False, indent=2), + "", + paras={ + "max_depth": self._max_depth, + "max_tokens": self._max_tokens, + "coarse_context": context, + }, + ) + answer: str = self._client.chat_completion( + messages=[ + {"role": "system", "content": "you are a document structure expert"}, + {"role": "user", "content": prompt}, + ], + model=self._model, + max_tokens=completion_budget, + temperature=temperature, + usage_task="page_memory.hierarchy", + ) + try: + return self._parse_levels(answer, candidates) + except _HierarchyResponseError as err: + logger.warning( + "[page_memory.fine_hierarchy] invalid response for {} candidates; {}", + len(candidates), + err, + ) + if len(candidates) == 1: + if can_retry: + return self._resolve_batch( + candidates, + preceding_candidates=preceding_candidates, + preceding_levels=preceding_levels, + can_retry=False, + ) + raise + midpoint: int = len(candidates) // 2 + left: Sequence[Mapping[str, object]] = candidates[:midpoint] + levels: dict[int, int] = self._resolve_batch( + left, + preceding_candidates=preceding_candidates, + preceding_levels=preceding_levels, + ) + levels.update( + self._resolve_batch( + candidates[midpoint:], + preceding_candidates=[*preceding_candidates, *left], + preceding_levels={**preceding_levels, **levels}, + ) + ) + return levels + + def _parse_levels( + self, + answer: str, + candidates: Sequence[Mapping[str, object]], + ) -> dict[int, int]: + content: str = answer.strip() + try: + if content.startswith("```") and content.endswith("```"): + content = content.split("\n", 1)[1].rsplit("```", 1)[0].strip() + response: object = json.loads(content) + except (json.JSONDecodeError, IndexError) as err: + raise _HierarchyResponseError( + "Hierarchy JSON is incomplete or malformed." + ) from err + if not isinstance(response, list): + raise _HierarchyResponseError("Hierarchy must be an array.") + expected_ids: set[int] = {int(str(candidate["id"])) for candidate in candidates} + levels: dict[int, int] = {} + for row in response: + if not isinstance(row, dict): + raise _HierarchyResponseError("Hierarchy row must be an object.") + identifier: object = row.get("id") + level: object = row.get("level") + if type(identifier) is not int or type(level) is not int: + raise _HierarchyResponseError( + "Hierarchy IDs and levels must be integers." + ) + if ( + identifier not in expected_ids + or identifier in levels + or not 0 <= level <= self._max_depth + ): + raise _HierarchyResponseError( + "Hierarchy contains duplicate, unknown or invalid rows." + ) + levels[identifier] = level + if set(levels) != expected_ids: + raise _HierarchyResponseError( + "Hierarchy did not classify every requested heading." + ) + return levels + + def _build_context( + self, + candidates: Sequence[Mapping[str, object]], + levels: Mapping[int, int], + ) -> str: + ancestors: list[str] = [] + ancestor_levels: list[int] = [] + for candidate in candidates: + identifier: int = int(str(candidate["id"])) + level: int = levels[identifier] + if level == 0: + continue + while ancestor_levels and ancestor_levels[-1] >= level: + ancestor_levels.pop() + ancestors.pop() + ancestor_levels.append(level) + ancestors.append(f"level={level}, heading={candidate['heading']}") + return self._coarse_context + ( + "\nPreviously classified ancestors in reading order (context only; do not output them):\n" + + "\n".join(ancestors) + + "\nContinue their relative levels; a batch boundary does not start a new subtree." + if ancestors + else "" + ) diff --git a/apps/worker/tests/contract/conftest.py b/apps/worker/tests/contract/conftest.py index 1a9ae74a9..3691a0443 100644 --- a/apps/worker/tests/contract/conftest.py +++ b/apps/worker/tests/contract/conftest.py @@ -110,6 +110,7 @@ def worker_contract_environment( ) -> Generator[None, None, None]: contract_runtime.configure_contract_environment(monkeypatch, postgresql_proc) asyncio.run(contract_runtime.prepare_contract_storage()) + asyncio.run(contract_runtime.configure_runtime_database_role(monkeypatch, contract_runtime.get_contract_database_url(postgresql_proc))) _ensure_worker_import_context() contract_runtime.clear_application_modules() diff --git a/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py b/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py index d2eaf126a..84e8488f1 100644 --- a/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py +++ b/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py @@ -2,6 +2,10 @@ import json import os +import re + +import pytest +from pytest import MonkeyPatch os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") @@ -23,6 +27,212 @@ def chat_completion(self, **_kwargs) -> str: return json.dumps(self.response) +def test_refine_large_leaf_preserves_headings_when_provider_output_is_truncated( + monkeypatch: MonkeyPatch, +) -> None: + """A bounded provider response must never collapse a whole PDF into Root.""" + + class BoundedClient: + def chat_completion(self, **kwargs: object) -> str: + messages = kwargs["messages"] + assert isinstance(messages, list) + prompt: str = str(messages[-1]["content"]) + identifiers: list[int] = [ + int(value) for value in re.findall(r'"id": (\d+)', prompt) + ] + response: str = json.dumps( + [{"id": identifier, "level": 1} for identifier in identifiers] + ) + # SpaceX's real response stopped mid-object at its completion budget. + budget: int = int(str(kwargs["max_tokens"])) + return response[: budget * 2] + + monkeypatch.setattr( + fine_hierarchy, + "get_text_client", + lambda requested_model=None: (BoundedClient(), "test-model"), + ) + skeleton: SectionSkeleton = SectionSkeleton( + section_path="spacex.pdf/Root", + level=1, + start_page=1, + end_page=407, + title="Root", + parent_path="spacex.pdf", + ) + tags: list[PageTagResult] = [ + PageTagResult( + page_index=page, + observed_titles=[{"text": f"Section {page}", "prominence": 1.0}], + ) + for page in range(1, 408) + ] + + refined: list[SectionSkeleton] = fine_hierarchy.refine_fat_leaf_skeletons( + coarse_skeletons=[skeleton], + tag_results=tags, + fat_leaf_pages=set(range(1, 408)), + model_name="test-model", + max_tokens=2000, + ) + + assert len(refined) == 407 + assert [item.start_page for item in refined] == list(range(1, 408)) + assert all(item.end_page == item.start_page for item in refined) + + +def test_refine_batches_preserve_parent_context_and_recover_incomplete_classification( + monkeypatch: MonkeyPatch, +) -> None: + prompts: list[str] = [] + + class IncompleteClient: + def chat_completion(self, **kwargs: object) -> str: + messages = kwargs["messages"] + assert isinstance(messages, list) + prompt: str = str(messages[-1]["content"]) + prompts.append(prompt) + identifiers: list[int] = [ + int(value) for value in re.findall(r'"id": (\d+)', prompt) + ] + # A well-formed but partial array must also trigger smaller requests. + if len(identifiers) > 2: + identifiers = identifiers[:-1] + return json.dumps( + [ + { + "id": identifier, + "level": 1 + if identifier == 1 + else (0 if identifier == 4 else 2), + } + for identifier in identifiers + ] + ) + + monkeypatch.setattr( + fine_hierarchy, + "get_text_client", + lambda requested_model=None: (IncompleteClient(), "test-model"), + ) + skeleton: SectionSkeleton = SectionSkeleton( + section_path="demo.pdf/Root", + level=1, + start_page=1, + end_page=8, + title="Root", + parent_path="demo.pdf", + ) + tags: list[PageTagResult] = [ + PageTagResult( + page_index=page, + observed_titles=[{"text": f"Heading {page}", "prominence": 1.0}], + ) + for page in range(1, 9) + ] + refined: list[SectionSkeleton] = fine_hierarchy.refine_fat_leaf_skeletons( + coarse_skeletons=[skeleton], + tag_results=tags, + fat_leaf_pages=set(range(1, 9)), + max_tokens=128, + model_name="test-model", + ) + + assert [item.title for item in refined] == [ + f"Heading {page}" for page in [1, 2, 3, 5, 6, 7, 8] + ] + assert all(item.parent_path == "demo.pdf/Root/Heading 1" for item in refined[1:]) + assert any( + "level=1, heading=Heading 1" in prompt and '"id": 5' in prompt + for prompt in prompts + ) + + +def test_refine_rejects_exhausted_invalid_output_instead_of_publishing_root( + monkeypatch: MonkeyPatch, +) -> None: + class InvalidClient: + def chat_completion(self, **kwargs: object) -> str: + return '[{"id": 1, "level":' + + monkeypatch.setattr( + fine_hierarchy, + "get_text_client", + lambda requested_model=None: (InvalidClient(), "test-model"), + ) + skeleton: SectionSkeleton = SectionSkeleton( + section_path="demo.pdf/Root", + level=1, + start_page=1, + end_page=407, + title="Root", + parent_path="demo.pdf", + ) + with pytest.raises(ValueError, match="incomplete or malformed"): + fine_hierarchy.refine_fat_leaf_skeletons( + coarse_skeletons=[skeleton], + tag_results=[ + PageTagResult( + page_index=1, + observed_titles=[ + {"text": "Launch capabilities", "prominence": 1.0} + ], + ) + ], + fat_leaf_pages={1}, + model_name="test-model", + ) + + +def test_refine_preserves_body_pages_before_first_observed_heading( + monkeypatch: MonkeyPatch, +) -> None: + from app.services.page_memory.node_assembler import identify_leaf_nodes + + monkeypatch.setattr( + fine_hierarchy, + "get_text_client", + lambda requested_model=None: ( + _FakeClient([{"id": 1, "level": 1}, {"id": 2, "level": 1}]), + "test-model", + ), + ) + skeleton: SectionSkeleton = SectionSkeleton( + section_path="demo.pdf/Root", + level=1, + start_page=1, + end_page=10, + title="Root", + parent_path="demo.pdf", + ) + refined: list[SectionSkeleton] = fine_hierarchy.refine_fat_leaf_skeletons( + coarse_skeletons=[skeleton], + tag_results=[ + PageTagResult( + page_index=3, + observed_titles=[{"text": "Launch capabilities", "prominence": 1.0}], + ), + PageTagResult( + page_index=7, + observed_titles=[{"text": "Financial statements", "prominence": 1.0}], + ), + ], + fat_leaf_pages=set(range(1, 11)), + model_name="test-model", + ) + leaves = identify_leaf_nodes(refined) + owned_pages: set[int] = { + page + for leaf in leaves + for page in ( + leaf.body_pages + if leaf.body_pages is not None + else range(leaf.start_page, leaf.end_page + 1) + ) + } + assert owned_pages == set(range(1, 11)) + + def test_compute_fat_leaf_pages_uses_exclusive_boundaries() -> None: skeletons = [ SectionSkeleton( @@ -140,12 +350,13 @@ def test_refine_fat_leaf_skeletons_excludes_next_section_start_when_unordered( ) assert [item.title for item in refined] == [ + "Section A", "A.1 Last Heading", "Section C", "B.1 Boundary Heading", ] - assert refined[0].parent_path == "demo.pdf/Section A" - assert refined[2].parent_path == "demo.pdf/Section B" + assert refined[1].parent_path == "demo.pdf/Section A" + assert refined[3].parent_path == "demo.pdf/Section B" def test_refine_fat_leaf_skeletons_uses_page_memory_prompt_without_demoting_siblings( @@ -370,6 +581,6 @@ def test_refine_fat_leaf_drops_boundary_page_on_tail_miss(monkeypatch) -> None: model_name="test-model", ) - assert [item.title for item in refined] == ["NCC 2022"] + assert [item.title for item in refined] == ["History", "NCC 2022"] assert all("amendment" not in item.title.casefold() for item in refined) assert all(item.start_page != 23 for item in refined) diff --git a/apps/worker/tests/contract/test_page_memory_retrieval_contract.py b/apps/worker/tests/contract/test_page_memory_retrieval_contract.py index f42f219b7..44ef55390 100644 --- a/apps/worker/tests/contract/test_page_memory_retrieval_contract.py +++ b/apps/worker/tests/contract/test_page_memory_retrieval_contract.py @@ -15,10 +15,6 @@ os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") os.environ.setdefault("S3_TEMP_PATH", "/tmp") -from shared.services.retrieval.hydration.assets import ( # noqa: E402 - build_retrieval_asset_url_map, - enrich_rows_with_retrieval_asset_url, -) from shared.services.retrieval.hydration.result_assembly import ( # noqa: E402 assemble_retrieval_results, ) @@ -28,9 +24,6 @@ build_term_search_text, ) from shared.services.retrieval.settings import normalize_chunk_types # noqa: E402 -from shared.services.retrieval.execution.reference_resolver import ( # noqa: E402 - resolve_workflow_references, -) from shared.core.exceptions.domain_exceptions import StorageServiceException # noqa: E402 from shared.services.storage.result_storage import JobResultStorage # noqa: E402 from shared.services.storage.page_pdf_crop import crop_source_pdf_pages # noqa: E402 @@ -155,6 +148,13 @@ async def test_table_result_assembly_uses_summary_not_html() -> None: @pytest.mark.asyncio async def test_page_asset_url_is_generated_from_page_nums(monkeypatch) -> None: + # Database contracts clear retrieval modules between environments. Resolve + # the live functions so this test patches the module it actually calls. + from shared.services.retrieval.hydration.assets import ( + build_retrieval_asset_url_map, + enrich_rows_with_retrieval_asset_url, + ) + monkeypatch.setattr( "shared.services.retrieval.hydration.assets.crop_source_pdf_pages", lambda *, job_id, pages: ( @@ -188,6 +188,11 @@ async def test_page_asset_url_is_generated_from_page_nums(monkeypatch) -> None: @pytest.mark.asyncio async def test_page_citation_asset_precedes_lazy_page_pdf_fallback(monkeypatch) -> None: + from shared.services.retrieval.hydration.assets import ( + build_retrieval_asset_url_map, + enrich_rows_with_retrieval_asset_url, + ) + page_pdf_calls: list[tuple[str, list[int]]] = [] def fake_crop_source_pdf_pages(*, job_id, pages): @@ -630,6 +635,10 @@ def generate_artifact_url(self, **_kwargs) -> str: async def test_referenced_chunks_get_page_asset_url_from_hydrated_rows( monkeypatch, ) -> None: + from shared.services.retrieval.execution.reference_resolver import ( + resolve_workflow_references, + ) + async def fake_hydrate_referenced_chunk_rows(**_kwargs): return [ { diff --git a/apps/worker/tests/contract/test_worker_bootstrap_contract.py b/apps/worker/tests/contract/test_worker_bootstrap_contract.py index b997796a9..f521954d9 100644 --- a/apps/worker/tests/contract/test_worker_bootstrap_contract.py +++ b/apps/worker/tests/contract/test_worker_bootstrap_contract.py @@ -5,6 +5,7 @@ import pytest from pytest import MonkeyPatch +from shared.testing.contract_runtime import PostgreSQLProcess def test_should_register_worker_task_modules_for_celery_consumers( @@ -187,3 +188,36 @@ def fail_watchdog_spawn(command: list[str]) -> FakeBeatProcess: assert beat_process.was_terminated is True assert beat_process.wait_timeouts == [5] + + +def test_should_reject_bypassing_database_role_before_starting_consumers( + worker_contract_environment: None, + monkeypatch: MonkeyPatch, + postgresql_proc: PostgreSQLProcess, +) -> None: + from contextlib import contextmanager + from collections.abc import Iterator + from sqlalchemy import create_engine + from sqlalchemy.engine import make_url + from sqlalchemy.orm import Session + from app.core import worker_bootstrap + from shared.core import database_sync + from shared.testing.contract_runtime import get_contract_database_url + + engine = create_engine(make_url(get_contract_database_url(postgresql_proc)).set(drivername="postgresql+psycopg2")) + + @contextmanager + def use_admin_database() -> Iterator[Session]: + with Session(engine) as database: + yield database + + def reject_process_spawn(*args: object, **kwargs: object) -> None: + raise AssertionError("A bypassing runtime role must not start consumers or sidecars") + + monkeypatch.setattr(database_sync, "get_sync_db_context", use_admin_database) + monkeypatch.setattr(worker_bootstrap.subprocess, "Popen", reject_process_spawn) + try: + with pytest.raises(RuntimeError, match="non-superuser"): + worker_bootstrap.run_worker() + finally: + engine.dispose() diff --git a/deploy/ecs/task-definition-api.staging.json b/deploy/ecs/task-definition-api.staging.json index dde385e26..0e6cf9a97 100644 --- a/deploy/ecs/task-definition-api.staging.json +++ b/deploy/ecs/task-definition-api.staging.json @@ -59,6 +59,7 @@ ], "secrets": [ {"name": "DATABASE_URL", "valueFrom": "${SECRETS_ARN}:DATABASE_URL::"}, + {"name": "DEMO_MAINTAINER_USER_IDS", "valueFrom": "${SECRETS_ARN}:DEMO_MAINTAINER_USER_IDS::"}, {"name": "REDIS_HOST", "valueFrom": "${SECRETS_ARN}:REDIS_HOST::"}, {"name": "REDIS_PORT", "valueFrom": "${SECRETS_ARN}:REDIS_PORT::"}, {"name": "REDIS_PASSWORD", "valueFrom": "${SECRETS_ARN}:REDIS_PASSWORD::"}, diff --git a/deploy/ecs/task-definition-worker.staging.json b/deploy/ecs/task-definition-worker.staging.json index accb00367..04dbdf17a 100644 --- a/deploy/ecs/task-definition-worker.staging.json +++ b/deploy/ecs/task-definition-worker.staging.json @@ -64,6 +64,7 @@ ], "secrets": [ {"name": "DATABASE_URL", "valueFrom": "${SECRETS_ARN}:DATABASE_URL::"}, + {"name": "DEMO_MAINTAINER_USER_IDS", "valueFrom": "${SECRETS_ARN}:DEMO_MAINTAINER_USER_IDS::"}, {"name": "REDIS_HOST", "valueFrom": "${SECRETS_ARN}:REDIS_HOST::"}, {"name": "REDIS_PORT", "valueFrom": "${SECRETS_ARN}:REDIS_PORT::"}, {"name": "REDIS_PASSWORD", "valueFrom": "${SECRETS_ARN}:REDIS_PASSWORD::"}, diff --git a/deploy/local-dev/postgres/init/02-runtime-role.sql b/deploy/local-dev/postgres/init/02-runtime-role.sql new file mode 100644 index 000000000..c97b26046 --- /dev/null +++ b/deploy/local-dev/postgres/init/02-runtime-role.sql @@ -0,0 +1,20 @@ +-- Local development only. Production passwords belong in a secret manager. +-- Root owns migrations; runtime processes receive DML without bypassing RLS. +DO $bootstrap$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'knowhere_runtime') THEN + CREATE ROLE knowhere_runtime LOGIN PASSWORD 'runtime123' + NOSUPERUSER NOCREATEDB NOCREATEROLE NOINHERIT NOBYPASSRLS; + END IF; +END +$bootstrap$; + +ALTER ROLE knowhere_runtime NOSUPERUSER NOCREATEDB NOCREATEROLE NOINHERIT NOBYPASSRLS; +GRANT CONNECT ON DATABASE "Knowhere" TO knowhere_runtime; +GRANT USAGE ON SCHEMA public TO knowhere_runtime; +GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public TO knowhere_runtime; +GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO knowhere_runtime; +ALTER DEFAULT PRIVILEGES FOR ROLE root IN SCHEMA public + GRANT SELECT, INSERT, UPDATE, DELETE ON TABLES TO knowhere_runtime; +ALTER DEFAULT PRIVILEGES FOR ROLE root IN SCHEMA public + GRANT USAGE, SELECT ON SEQUENCES TO knowhere_runtime; diff --git a/deploy/local-dev/start-dev.sh b/deploy/local-dev/start-dev.sh index 6224235fe..2756d214c 100755 --- a/deploy/local-dev/start-dev.sh +++ b/deploy/local-dev/start-dev.sh @@ -100,12 +100,14 @@ Local development services are ready. Service endpoints: - LocalStack: http://localhost:4566 - - PostgreSQL: localhost:5432 (root/root123) + - PostgreSQL: localhost:5432 (runtime: knowhere_runtime/runtime123; migrations: root/root123) - Redis: localhost:6379 Next steps: - 1. Start the API: cd apps/api && uv run uvicorn main:app --host 0.0.0.0 --port 5005 --reload - 2. Start the worker: cd apps/worker && uv run python worker.py + 1. Copy the API/Worker .env.example files to .env and configure your provider keys. + 2. Migrate: cd apps/api && MIGRATION_DATABASE_URL=postgresql+asyncpg://root:root123@localhost:5432/Knowhere uv run alembic upgrade heads + 3. Start the API: cd apps/api && uv run uvicorn main:app --host 0.0.0.0 --port 5005 --reload + 4. Start the worker: cd apps/worker && uv run python worker.py Stop services: ${SCRIPT_DIR}/stop-dev.sh @@ -120,6 +122,9 @@ main() { run_compose -f "${COMPOSE_FILE}" up -d wait_for_postgres + # Init scripts only run on empty volumes. Provision existing volumes as well. + docker exec -i knowhere_postgres psql -U root -d Knowhere -v ON_ERROR_STOP=1 \ + < "${SCRIPT_DIR}/postgres/init/02-runtime-role.sql" wait_for_redis wait_for_localstack print_summary diff --git a/docs/shared-demo-pr-description.md b/docs/shared-demo-pr-description.md new file mode 100644 index 000000000..1ab5d82d4 --- /dev/null +++ b/docs/shared-demo-pr-description.md @@ -0,0 +1,49 @@ +## Summary + +Demo selection previously created a private copy for each user. We publish each source once into twelve dedicated PostgreSQL tables in `__knowhere_demo__`, with stable document IDs, retained revisions/assets, shared retrieval and caller-specific statistics. Optimized batched publication is now the only implementation. + +Maintainers use the existing V1/V2 upload, confirmation and Worker parser flow. API and Worker validate the server-owned corpus target; FORCE RLS rejects unapproved writes and both services reject database roles that bypass RLS. Publication atomically switches the revision, manifests and global generation. Document, chunk, media, agent tools and citation reads honor captured revision pins. + +Real SpaceX validation exposed two failures addressed here: hierarchy JSON truncation silently produced a single Root chunk, and agent context compaction removed useful evidence and caused repeated reads. Bounded hierarchy batches validate every heading, retry malformed output and fail the import if recovery is exhausted. The OpenAI harness retains successful evidence while compacting old discovery results. + +### Client migration + +- `POST /api/v1/demo/materializations` returns **410 `DEMO_MATERIALIZATION_REMOVED`**. Existing private copies remain accessible through ordinary APIs. +- Save catalog `canonical_document_id` and `job_result_id`; use the shared namespace directly. Hide a demo source locally instead of archiving the shared document. Mixed private/demo retrieval needs separate requests. +- Pass `job_result_id` when opening historical citations, chunks, originals and assets. Old static chunk IDs are not stable after reparsing. +- Catalog is `no-store`. Explicit preparing sources return 409; an empty demo corpus returns 503. The initial migration seeds directory metadata without READY content. + +```json +{ + "namespace": "__knowhere_demo__", + "query": "What are SpaceX's launch capabilities?", + "include_document_ids": [""], + "use_agentic": false +} +``` + +## Verification + +- 434 API contracts, 30 migration contracts and 342 Worker contracts passed. The full API run found three stale migration-contract assumptions; their revision scope was corrected and the complete migration suite rerun successfully. Runtime code was unchanged after that full run. Existing API unit tests also passed; no new unit tests were added. +- Ruff, Pyright, staged diff checks and a redacted Gitleaks tree scan passed in an isolated checkout of the submitted code. +- Real local HTTP create โ†’ PUT โ†’ confirm โ†’ poll โ†’ retrieval/media checks used the production Celery/gevent Worker, enforced RLS, Redis, LocalStack and real providers for mixed DOCX, Tesla and the 153,808,384-byte SpaceX original. +- SpaceX now has 502 page chunks covering all 407 pages. The originally failing agent query finishes for both readers in 7โ€“8 seconds with citations to pages 31โ€“32, 167โ€“168 and 208โ€“209. Historical reads and PDF 206 Range responses passed; readers create no jobs, chunks, tokens or stored assets. +- All 27 originals were backed up and SHA-256/format validated. Full SpaceX processing took about 43 minutes. Measured publication took 18.081 seconds; timed runs included two pre-existing local projection optimizations excluded from this PR, so this is not a benchmark of the exact submitted tree. +- Notebook migration and production deployment are separate; the complete end-user flow awaits that client upgrade. + +## Deployment Notes + +- Remove `KNOWHERE_PUBLICATION_STRATEGY`. Set `DEMO_MAINTAINER_USER_IDS` and use a NOSUPERUSER/NOBYPASSRLS runtime `DATABASE_URL`; the default maintainer set is empty. +- Run additive migration `4d5e6f7a8b9c` with separate management credentials (`MIGRATION_DATABASE_URL`). It stops on private documents already using the reserved namespace. Grant runtime DML permissions after migration. +- Verify the original-file backup before cutover. Stop old materialization admission and drain tasks before updating API/Worker together. Restore Tesla first, then SpaceX, then other sources using the resumable helper; service startup does not trigger parsing. +- Roll back API/Worker images together after stopping imports and draining/cancelling new tasks. Retain additive tables, originals and revision assets; do not run a destructive down migration. + +See `docs/shared-demo-rollout.md` for backup/runtime grants, upload examples, recovery commands, client migration and rollback, and `docs/shared-demo-validation-20261003.md` for measured results and limitations. + +## Checklist + +- [x] Contract tests added or updated for changed behavior +- [x] Public documentation, examples and migration contracts updated +- [x] Additive migrations and runtime-role provisioning verified locally +- [x] Logs, errors and validation paths avoid leaking credentials +- [x] Breaking changes and release ordering documented diff --git a/docs/shared-demo-rollout.md b/docs/shared-demo-rollout.md new file mode 100644 index 000000000..6985ddca9 --- /dev/null +++ b/docs/shared-demo-rollout.md @@ -0,0 +1,270 @@ +# Shared demo corpus rollout and client migration + +This backend replaces personal demo materialization with one shared corpus in +`__knowhere_demo__`. Publication always uses the optimized implementation. +Remove `KNOWHERE_PUBLICATION_STRATEGY` from deployment settings; no switch enables +the new demo path. Historical benchmark baseline/candidate labels are retained. +This document defines the rollout contract for the shared demo backend. + +## Storage and access + +The existing PostgreSQL database gains twelve dedicated `demo_*` tables. +Private documents remain in their existing tables. A source's `data_id` is its +unique `demo_source_id` (up to 128 characters); `ddoc_...` document IDs remain +stable across reparses. Source directory metadata is seeded, but no old static +chunks are imported as READY content. Publication atomically writes sections, +chunks, map units, tokens, completeness markers, manifests, graph, generation +and the current revision. Demo INSERT batches contain at most 1,000 rows. +PostgreSQL prohibits COPY FROM into an RLS table, so private COPY and Worker +psycogreen/execute_values paths remain separate from demo INSERTs. + +Only user IDs in `DEMO_MAINTAINER_USER_IDS` may create, replace or archive demos; +the default is empty. They must also hold the normal API write permission. +The Worker revalidates durable job target/source/document metadata and identity. +Moves between private and demo scopes are rejected. Other authenticated users +can query the same shared documents. Caller identity remains the basis for +authentication, billing, limits, cache keys and hit statistics. + +All dedicated tables use FORCE RLS. Request processes must use a non-superuser +role without BYPASSRLS. API and Worker refuse to start with a bypassing role. +The server sets transaction-local maintenance/caller context after authorization. +Database credentials are internal service credentials, never end-user access. +Public content policies require an active source and a complete bound revision; +directory reads can include preparing metadata. RLS UPDATE of invisible rows +affects zero rows; unauthorized INSERT fails. Hit-stat policies bind the actual +caller. Retained revisions restrict deletion of their jobs/results even after +archive. API job deletion requests for these revisions return 409. This checkout +has no general job deletion/cancellation API; ordinary jobs remain retained +processing records, and DELETE remains unsupported for them. + +## Backup gate + +Use the authoritative originals in `knowhere-storage-staging`, not public URLs. +`scripts/demo-original-sources.json` freezes all 27 source mappings from the +existing catalog and each `app/data/demo_documents//manifest.json`. +Each mapping uses `uploads/`. + +```bash +uv run --all-packages python scripts/backup-demo-originals.py \ + --profile knowhere --bucket knowhere-storage-staging \ + --directory "$HOME/knowhere-demo-originals/20261003" +``` + +The inventory records source ID, original job ID, S3 key, filename, actual size, +SHA-256 and format validation. The command checks every PDF page, DOCX archive +CRC and document structure, and fails the gate if any object is missing or +invalid. Resume reruns reuse only objects whose size and source ETag still match. +Do not switch with an incomplete inventory. These S3 originals have no versioned +recovery copy, so preserve an independent persistent backup. + +The 2026-10-03 backup contains 27 verified originals totaling 225,749,274 bytes. +The prior backup run copied them to +`/Users/ontosai/knowhere-demo-originals/20261003` and independently verified remote +SHA-256 values. The local validation copy is +`/tmp/knowhere-demo-originals/20261003`; this temporary copy is not the persistent +backup. The current local run also copied and independently rehashed all originals +into `.demo-originals/20261003` in this checkout. That directory is Git-ignored +and survives removal of temporary validation files. Preserve it before deleting +the checkout. Mac availability was not reverified after the prior backup run. +SpaceX is 153,808,384 bytes / 407 pages. Ensure **both** API and Worker +configure `MAX_FILE_SIZE=314572800`; older local configurations may still have +a 100 MiB limit. + +## Migration and runtime roles + +Stop if this preflight returns rows: + +```sql +SELECT document_id, user_id FROM documents +WHERE namespace = '__knowhere_demo__'; +SELECT current_user, rolsuper, rolbypassrls FROM pg_roles +WHERE rolname = current_user; +``` + +Migration `4d5e6f7a8b9c` repeats the reserved-namespace conflict check. It is +additive and freezes its DDL/catalog independently of future ORM changes. +Run Alembic using a separate schema-management connection. Supply +`MIGRATION_DATABASE_URL` only to the migration process; use the nonprivileged +`DATABASE_URL` for API and Worker. Provision existing auth tables/users with +management credentials before starting runtime processes. + +An administrator can provision a runtime role and grant existing service DML: + +```sql +CREATE ROLE knowhere_runtime LOGIN NOSUPERUSER NOBYPASSRLS; +-- Set its password through your secret manager, not a checked-in SQL file. +GRANT CONNECT ON DATABASE "Knowhere" TO knowhere_runtime; +GRANT USAGE ON SCHEMA public TO knowhere_runtime; +GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public TO knowhere_runtime; +GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO knowhere_runtime; +-- Execute for the actual migration table owner: +ALTER DEFAULT PRIVILEGES IN SCHEMA public +GRANT SELECT, INSERT, UPDATE, DELETE ON TABLES TO knowhere_runtime; +ALTER DEFAULT PRIVILEGES IN SCHEMA public +GRANT USAGE, SELECT ON SEQUENCES TO knowhere_runtime; +``` + +The runtime role does not need schema ownership, superuser or BYPASSRLS. Preserve +existing deployment-specific function permissions and extensions. Migration +credentials must retain DDL permissions; runtime credentials must not replace +them. Verify both processes' effective roles before accepting traffic. + +For the repository's local Docker stack, `deploy/local-dev/start-dev.sh` +provisions `knowhere_runtime` on both new and existing volumes. The API/Worker +examples use that role; the local password is development-only. Run migrations +separately before starting either service: + +```bash +cd apps/api +MIGRATION_DATABASE_URL=postgresql+asyncpg://root:root123@localhost:5432/Knowhere \ + uv run alembic upgrade heads +``` + +Do not put the root connection in the API/Worker `DATABASE_URL`. Local role +provisioning preserves an existing runtime password; adapt copied examples if +that password has already been changed. + +## Coordinated cutover + +1. Complete backup, contracts and real HTTP/Worker validation. +2. Check for reserved private-namespace conflicts and role bypass. +3. Run additive migrations; confirm the directory is planned with no READY data. +4. Configure maintainer IDs and runtime credentials; remove strategy settings. + Before deploying ECS, add `DEMO_MAINTAINER_USER_IDS` to the environment's + runtime secret. Both task definitions require that key. Use a comma-separated + list of authorized user IDs, or an empty string to disable demo writes until + a maintainer is selected. Adding an API key is not required. +5. Stop old materialization admission and drain its tasks before updating API and + Worker together. Old tasks must not execute new publication code. +6. Verify private upload/retrieval, catalog preparing state and materialization + 410 before manually restoring sources. Demo downtime is expected. +7. Restore Tesla first as a smoke test, then SpaceX, then remaining inventory. + +No service startup automatically parses demos. A failed source does not block +other READY sources. A failed update leaves its earlier complete revision active. +Initial failures remain unavailable. Keep historical revisions and all assets; +there is no automatic garbage collection. + +## Manual HTTP restoration + +Provide a maintainer API key through `KNOWHERE_DEMO_API_KEY` in your environment. +Do not put it in command arguments or commit it. The helper uses v2 jobs, +uploads the verified file through the returned PUT capability, confirms upload, +polls with increasing intervals, and checks the READY catalog. + +```bash +uv run --all-packages python scripts/restore-demo-originals.py \ + --api-url https://your-knowhere-host \ + --inventory "$HOME/knowhere-demo-originals/20261003/inventory.json" \ + --source demo-tsla-q4-2025 --state-directory ./demo-restoration + +# Repeat with --source demo-spacex-s1, then omit --source for all inventory. +# Default concurrency is 2. Recorded jobs are always resumed. +# --retry-failed creates a new job after failure; --reparse-ready creates +# a new revision for already completed recorded jobs. +``` + +State files persist each job/document/revision and a sanitized result report. +They omit provider secrets and signed URLs. Waiting-file jobs resume through +`POST /api/v2/jobs/{job_id}/upload-url`, which renews the upload capability only +for the owner with write permission and a valid corpus target. Interrupting the +helper stops polling; it does not cancel Celery processing. Resume follows the +same recorded job. To stop an import before rollback, drain the Worker or use +the existing operator failure/stale-job lifecycle; do not forcibly kill gevent +greenlets. Retry failed imports with a new job, retaining the old ledger entry. +GET jobs now returns the immutable `job_result_id` when available. The helper +validates the recorded source and original hash, and verifies this exact revision +through document reads; it does not attribute a later concurrent publication to +the earlier job. + +The equivalent creation request is: + +```http +POST /api/v2/jobs +Authorization: Bearer +Content-Type: application/json + +{"namespace":"__knowhere_demo__","source_type":"file","file_name":"spacex-s1.pdf","data_id":"demo-spacex-s1"} +``` + +PUT the original using `upload_url` and `upload_headers`, POST `{}` to +`/api/v2/jobs/{job_id}/confirm-upload`, then GET the job until done/failed. +The source claim serializes updates across API processes; another active import +returns 409. After completion inspect catalog, retrieval, chunks, original, +image/table/page assets and v2 `files/page-citation-source`. + +## Client breaking changes + +`POST /api/v1/demo/materializations` now returns 410 +`DEMO_MATERIALIZATION_REMOVED` with shared-namespace migration guidance. Existing +personal copies remain private documents accessible through ordinary APIs. +Old static canonical/chunk IDs are not IDs for newly parsed shared content. + +Catalog uses `Cache-Control: no-store`. `sources` contains READY sources only; +`import_status` and official-library planned entries describe processing/failure. +Catalog returns stable `canonical_document_id`, source ID, `job_result_id`, +namespace and actual parsed counts/size. Unmatched curated example citations +are omitted until uniquely matched or manually corrected. + +Query one namespace at a time: + +```json +{ + "namespace": "__knowhere_demo__", + "query": "What are SpaceX's launch capabilities?", + "include_document_ids": [""], + "use_agentic": false +} +``` + +Mixing private sources and demos requires separate namespace requests and +client-side result merging. Retrieval source/evidence/referenced chunks include +their `job_result_id`. Document/chunk/demo original/assets and v2 page-citation +source reads accept an optional `job_result_id`: omitted means current complete +revision; specified means a complete revision belonging to that source. Demo +chunk `id` is the dedicated row ID; `chunk_id` is the content ID. Preserve both. +Original/assets issue short-lived signed 307 redirects with MIME/disposition; +object storage serves Range. Only exact paths in that revision's asset manifest +are accepted, with traversal rejection. +Chunk and retrieval asset projections use the same manifest validation and +one-hour MIME-aware signatures. Demo page previews reuse published page images +or the retained normalized PDF; reader requests never create cropped PDFs. + +| State | Response | +| --- | --- | +| No READY demos in catalog | 200, `sources=[]`, official library planned | +| Query whole corpus without READY sources | 503 `DEMO_NOT_READY` | +| Explicitly request a known preparing source | 409 `DEMO_NOT_READY` | +| Unknown, archived, or wrong-source revision | 404 | +| Some READY sources | Query READY subset; other sources remain preparing | + +Notebook needs a separate change before the complete user workflow can be +accepted: remove materialization; save shared source/document/revision; split +private/demo requests with include IDs; hide local demo selections instead of +archiving shared documents; pass revision into citations and media reads; key +chunk caches by revision; make catalog no-store; update folder/exclusion maps; +remove personal/static citation remapping; render 410/409/503 migration/preparing +states. Other clients using static IDs or materialization require the same migration. + +## Verification and rollback + +Contract tests cover private publication's COPY/gevent compatibility and shared +HTTP/RLS publication/reads, source claims, history, cache generations and archive. +Real validation runs API, production gevent Celery Worker, PostgreSQL, Redis and +S3-compatible object storage through HTTP job creation/upload/confirm/poll. +Provider calls must be real; payload replay is only a publication performance +diagnostic. Record API/upload, queue, parse, asset upload and publication timings +separately. Compare the same captured payload once between private and demo +storage. Multiple readers must not create jobs/chunks/tokens or reupload assets. +Full SpaceX parsing time is separate from user selection/read time; no ten-second +full-parsing promise is made. + +Before downgrade, stop demo import admission and drain or fail active new jobs. +Restore previous API and Worker images together. Preserve all additive tables, +originals and revision assets; do not run a destructive down migration. Demo may +remain unavailable during rollback. Re-enable imports only after the compatible +pair is running. Production deployment and Notebook changes are separate from +this backend implementation. + +See [the local verification report](shared-demo-validation-20261003.md) for +measured results, the SpaceX hierarchy/agent fixes and validation scope. diff --git a/docs/shared-demo-validation-20261003.md b/docs/shared-demo-validation-20261003.md new file mode 100644 index 000000000..ec8bce9ed --- /dev/null +++ b/docs/shared-demo-validation-20261003.md @@ -0,0 +1,244 @@ +# Shared demo backend validation โ€” 2026-10-03 + +## Result and boundary + +The backend was implemented and validated locally on branch +`feat/wangbinqi/shared-demo-corpus`. API and production Celery/gevent Worker ran +against isolated PostgreSQL 16, Redis, and LocalStack S3. Both service processes +used `demo_runtime`, a non-superuser role without BYPASSRLS. DeepSeek and Ali +provider calls were real; LLM mocking was disabled. The agent harness was the +OpenAI-compatible harness with DeepSeek. The real-provider validation run made +no production deployment, cloud mutation or Notebook change. + +Timed runs used the local working tree, including pre-existing chunk-read +projection optimizations in `map_unit_index.py` and `serving_manifest.py`. Those +projection hunks and the earlier `prepare_demo_bundles.py` script edit are outside +this PR and remain local. Publication times below characterize that measured +tree, rather than a benchmark of the exact submitted tree. The submitted code +is separately checked in an isolated checkout. API and Worker contract suites +run sequentially because their fixtures share `/tmp/knowhere-api-tests`. +Three page-retrieval contracts now resolve live imports after runtime module +reset, so their provider/storage patches affect the functions they execute. + +The shared upload/publication/read flow works. Follow-up fixes resolved the two +SpaceX failures found during initial validation: truncated hierarchy responses +silently published one Root chunk, and agent context compaction removed useful +read evidence, causing repeated reads until the token budget ran out. + +The real SpaceX update now publishes 502 page chunks covering all 407 pages; +the largest section spans ten pages. Both readers completed classic and real +agent queries with citations to specific page ranges and the same new revision. +Historical reads and PDF byte ranges remain available, and reading creates no +jobs, chunks, token rows or stored assets. These checks cover the launch queries +described below, rather than a comprehensive retrieval quality benchmark. + +## Automated verification + +| Check | Result | +| --- | --- | +| Full API contract suite, including evidence-retention regression | 434 passed | +| API migration contracts after historical revision-scope corrections | 30 passed | +| Existing API unit tests in the full API run | 12 passed | +| Full Worker contract suite in the submitted checkout | 342 passed | +| Relevant Worker contracts: parsing, processing, bootstrap, publication trace | 29 passed | +| New shared demo contract file | 7 passed | +| Full page-memory contracts after hierarchy fix | 83 passed | +| Agent evidence-retention and shared demo contracts | 8 passed | +| `make check` | Ruff and Pyright passed | +| `git diff --check` | Passed | + +The isolated full API command initially produced 473 passes and three migration +contract failures. One assumed the old revision was still head; two attempted +to downgrade through the new retained demo schema. The maintenance contract now +checks its captured initial head, and reversible token-index tests explicitly +target their historical migration chain. All 30 migration contracts then passed. +Only those test files changed after the full API run; runtime code was unchanged. +The isolated Worker command passed all 342 contracts after correcting the stale +test imports described above. Gitleaks found no leaks in the submitted tree. + +Commands for the submitted checkout were the CI equivalents of +`uv run pytest apps/api/tests -q`, `uv run pytest apps/api/tests/migrations -q`, +`uv run pytest apps/worker/tests/contract -q`, Ruff and Pyright. The isolated +checkout reused the existing virtual environment with its own shared/API/Worker +paths; package dependencies were unchanged. + +The local-development role bootstrap was also checked in an isolated +PostgreSQL 15 container using the actual init scripts. Repeating provisioning +succeeded; the runtime role had no superuser, BYPASSRLS, database-creation or +role-creation privileges. A table created afterward inherited DML grants, while +an INSERT without an RLS policy was rejected. The startup shell passed `bash -n`. + +The new contracts exercise preparing/410 responses, maintainer authorization, +FORCE RLS, same shared data for two users, private isolation, concurrent source +claims, second-batch rollback, Worker target validation, update failures, retained +revisions, global cache generation, agent tools and final references, archive, +revision-aware media, traversal rejection, deletion protection, include/exclude, +section/type filters, no matches, unpublished/wrong-source revisions, retry after +failure, curated citation rebinding, and caller-specific hit statistics. The +Worker bootstrap contract verifies a bypassing role fails before consumers or +sidecars start. Tests use real PostgreSQL; provider fakes in contract tests are +separate from the real provider verification below. + +This checkout has no general job cancellation/deletion API. Ordinary DELETE +remains unsupported; retained shared revision DELETE returns 409. No new cancel +endpoint or destructive job cleanup was added. + +## Real HTTP and Worker verification + +Every parsed input below used HTTP create job, PUT original, confirm upload, +GET polling, and HTTP reads. Worker was started through `worker.py` with the +production gevent pool and psycogreen. An earlier SpaceX job correctly failed the +old local 100 MiB limit; validation then configured the planned 300 MiB limit on +both processes. The initial successful parse exposed the hierarchy bug. A real +HTTP update after fixing it preserved the document ID and retained that earlier +revision. + +| Input | Track | Published chunks | Parser time | Full Worker task | +| --- | --- | --- | --- | --- | +| Synthetic DOCX with text, image and table | chunk | 5: 3 text, 1 image, 1 table | 6.643 s in the final timed update | 8.106 s | +| Tesla Q4 2025 authoritative original | page_memory | 73: 37 page, 36 image | 348.778 s | 357.56 s | +| SpaceX initial import before hierarchy fix | page_memory | 1 page chunk after hierarchy fallback | 1,217.407 s | 1,243.392 s | +| SpaceX authoritative original update, 153,808,384 bytes / 407 pages | page_memory | 502 page chunks | 2,546.914 s | 2,585.419 s | +| Same synthetic DOCX uploaded privately | chunk | 5 | 6.851 s | HTTP completion observed after 10.107 s | + +The final demo DOCX request measured create 29 ms, original PUT 25 ms, +confirm 55 ms, retained result ZIP/assets upload 46 ms, and database publication +including commit 267.242 ms. Polling observed READY after 9.917 s. Queue delivery +was within the same one-second Worker log timestamp bucket; that is a bound, +not a precise queue latency measurement. Private DOCX measured create 60 ms, +PUT 21 ms, confirm 128 ms, and publication 303.254 ms. + +The large PDF jobs ran before the new upload timer and publication tracing were +enabled. Their original-upload figures (Tesla 0.246 s, SpaceX 1.351 s) include +the then-combined helper interval, rather than separately measured PUT/confirm. +ZIP/assets upload and database publication are only coarse log bounds for these +initial PDF jobs. The corrected SpaceX update has independent timings below. +The helper now records create, PUT and confirm separately, and the Worker records +`worker.result.assets_upload` independently of the publication trace. + +HTTP validation checked catalog no-store, document lists in both API versions, +chunk reads, originals, DOCX image/table assets, Tesla page images, and normalized +PDF page-citation-source. Storage responded 206 to byte ranges, with correct +PDF, DOCX, PNG and HTML MIME types. Invalid manifest paths returned 404; ordinary +demo upload returned 403; materialization returned 410. Tesla, DOCX and corrected +SpaceX agent episodes completed with revision-bearing sources and references. + +Two users executed classic/shared reads and agent retrieval. During this read +interval, before/after counts were identical: + +| Resource | Before | After | +| --- | ---: | ---: | +| Jobs | 5 | 5 | +| Demo chunks, including retained revisions | 83 | 83 | +| Demo token rows | 13,530 | 13,530 | +| Private chunks/token rows | 0 / 0 | 0 / 0 | +| S3 objects | 490 | 490 | + +The private upload, publication diagnostic, and later manual DOCX updates took +place after that interval and intentionally changed counts. A real DOCX update +retained `ddoc_706bbc5ce1f9`; both readers subsequently observed revision +`5d0eb1ca-5d35-42bb-b04b-a66488dcf26d` while the old +`a56acbeb-0ad3-4eee-adad-8eacefcfc7cb` remained readable through pinned chunk and +original requests. The recovery helper resumed completed Tesla/SpaceX jobs +without creating new parsing jobs and recorded their own immutable revisions. + +## SpaceX follow-up acceptance + +The updated source retains `ddoc_ceb0d9f0c884` and publishes revision +`c7e46329-9185-4d4d-a459-f2cf6e303735` through job `job_a0a5f79cfb1f`. +The old `73863515-a0d4-420b-adb0-d8efbaca6eaa` revision remains readable. +Fine hierarchy classified 842 observed titles into 562 skeleton nodes; one real +malformed response batch was detected and recovered by smaller retries. +Publication committed 553 sections, 502 chunks, 477 map units and 81,833 tokens. + +| Phase | Measured duration | +| --- | ---: | +| Profiling | 822.010 s | +| Title rendering / detection | 65.890 s / 546.147 s | +| Fine hierarchy | 26.217 s | +| Page rendering / tagging | 65.658 s / 420.940 s | +| Node assembly, including summaries | 599.273 s | +| Complete parser | 2,546.914 s | +| Result/assets upload | 7.169 s | +| Database publication, including commit | 18.081 s | +| Complete Worker task | 2,585.419 s (43 min 5 s) | + +The HTTP helper separately measured job creation at 0.102 s, original PUT at +0.992 s and upload confirmation at 0.128 s. Polling observed completion after +2,588.824 s. Queue delivery was not independently measured with subsecond +precision. + +All 407 pages occur in the published page metadata. Maximum raw chunk content +is 42,365 characters, compared with 1,467,478 in the initial Root chunk. This +revision has page images and its normalized source PDF, but no standalone image +or table chunks; DOCX and Tesla separately exercised those asset paths. + +Classic query `Falcon 9 launch capabilities payload capacity` returned three +results with page ranges 31โ€“32, 167โ€“168 and 107โ€“108. Reader request times were +2.458 s and 1.930 s. Real agent query `Compare the launch capabilities and payload +capacities of Falcon 9 and Falcon Heavy.` finished in 9.439 s and 9.831 s, +returning three and four references. Both sets included specific Falcon launch +sections and new-revision pins. Citation pages were checked against text extracted +from the backed-up original PDF. The citation PDF served 206 with PDF MIME for +byte range 0โ€“63. + +The original query that exhausted its token budget, `What are the launch +capabilities of Falcon 9 and Falcon Heavy?`, was also repeated after bypassing +each reader's local query cache. Both episodes finished in 7.017 s and 8.075 s +with three references to page ranges 31โ€“32, 167โ€“168 and 208โ€“209. API restart alone +was not used as proof of cache invalidation. The traces match the original +recall/grep and first read, then perform one complementary read and finish. +They consume 43,185 and 43,242 tokens, compared with 111,394 tokens in the +original budget-exhausted episode. + +During those two-reader checks, jobs remained 11, demo chunks including history +remained 668, token rows remained 99,883, and S3 objects remained 912. No parser +job or materialization was created. Agent context retention was changed only +after this new revision existed, so agent verification did not reparse the file. + +The evidence-retention contract reproduces a multi-section read whose second +fact is truncated, followed by a complementary read. Before the fix, evidence +vanished after two turns and reads alternated until the episode budget expired. +After the fix, the episode finishes with both references while stale discovery +content is compacted. Successful evidence messages with valid refs stay visible; +per-result text caps and episode budgets remain enforced. + +See [the hierarchy fix](spacex-hierarchy-fix.md) for batching, strict response +validation, cross-batch nesting and preservation of opening body pages. + +## Publication diagnostic + +The same captured Tesla ZIP's 73 chunks were replayed once into a private scope +and once into a demo scope using gevent/psycogreen and the real runtime role. +Both committed 45 sections, 73 chunks, 36 map units and 4,345 token rows. + +| Storage | Publication including commit | SQL statements | SQL time | +| --- | ---: | ---: | ---: | +| Private execute_values compatibility path | 1,366.279 ms | 51 | 461.289 ms | +| Demo INSERT under FORCE RLS | 1,087.970 ms | 58 | 540.103 ms | + +These are single diagnostic samples, with different existing namespace snapshot +and graph sizes and private cold-start effects. They show the storage adapter +can publish the same payload; they do not prove demo is faster. Replay is not +parsing E2E validation. The diagnostic demo was archived through the normal API; +its history is retained in the isolated validation database. + +## Backup and delivery artifacts + +All 27 authoritative originals were format-validated and copied into the local +persistent `.demo-originals/20261003` directory; every size and SHA-256 was checked +again after copying. Total: 225,749,274 bytes. The inventory and backup directory +are Git-ignored. A prior run independently verified the Mac backup as well; +this run did not reverify Mac availability after SSH became unreachable. + +Sanitized JSON evidence is preserved under +`.demo-originals/validation-20261003/`. Temporary service credentials and raw +provider logs are excluded. Local API uses port 55005; isolated PostgreSQL, +Redis and S3 use 55543, 56379 and 54566. Services may be restarted from the local +validation environment, but credentials must never be copied into delivery +documents. + +See [the rollout guide](shared-demo-rollout.md) for migration credentials, +maintainer/runtime configuration, backup and resumable recovery commands, +client breaking changes, query examples, coordinated cutover, and rollback. +Notebook must still migrate before accepting the complete end-user demo flow. diff --git a/docs/spacex-hierarchy-fix.md b/docs/spacex-hierarchy-fix.md new file mode 100644 index 000000000..62eca2ad0 --- /dev/null +++ b/docs/spacex-hierarchy-fix.md @@ -0,0 +1,52 @@ +# SpaceX hierarchy response recovery + +## Failure and reproduction + +The first real shared SpaceX import produced one page chunk spanning all 407 +pages. Page-title detection found 836 headings, but fine hierarchy sent its +deduplicated candidates in one request with a fixed 2,000-token completion +budget. The returned JSON ended midway through candidate 154. The generic +response parser returned raw text; hierarchy parsing interpreted that as no +levels and retained Root. The import was then marked successful. + +The regression contract feeds a 407-page coarse leaf through real candidate +collection, prompt construction and skeleton refinement, using a provider +boundary that truncates its response at the supplied budget. Before the fix, +`test_refine_large_leaf_preserves_headings_when_provider_output_is_truncated` +failed with `assert 1 == 407`. After the fix it preserves all 407 sections. + +## Change + +`PageHierarchyLevelResolver` bounds candidate batches by the completion budget, +with at most 64 rows. It supplies the preceding ancestor chain with its assigned +relative levels so a batch boundary does not reset nesting. Every requested ID +must have an integer level; level 0 explicitly marks filtered noise. Duplicate, +unknown, missing or invalid rows and malformed JSON are rejected. + +An invalid batch retries in smaller batches with the same hierarchy context. +An invalid singleton retries once, then fails the import. The current published +revision remains available; a malformed hierarchy is never treated as successful +Root fallback. Valid all-noise output can still retain the coarse section. + +Refinement also retains the coarse parent when body pages precede the first +observed heading. Assembly can then preserve those pages alongside the refined +children. A separate contract reproduced missing opening pages before this fix. + +## Verification + +All 83 page-memory contracts passed, including truncation, incomplete-array +recovery, cross-batch nesting, explicit noise, exhausted recovery and opening +body-page preservation. Ruff and Pyright passed. The real HTTP/Worker update +published 502 page chunks covering all 407 pages, including opening body pages, +with no section spanning more than ten pages. The source document ID remained +stable and its old revision remains readable. Full Worker time was 43 minutes +5 seconds; hierarchy refinement was 26.217 seconds. + +Follow-up agent validation exposed a separate context issue: successful reads +were compacted after two turns, causing repeated reads of complementary +sections. The OpenAI harness now retains successful evidence with refs while +compacting stale discovery/error messages. A full-episode provider contract +reproduced the loop before this change and finishes in four turns afterward. +Both real reader launch-comparison queries finished with precise page ranges. +See [the validation report](shared-demo-validation-20261003.md) for timings and +the exact verification boundary. diff --git a/packages/shared-python/shared/core/config/__init__.py b/packages/shared-python/shared/core/config/__init__.py index 25f9e2ca1..55c43ce06 100644 --- a/packages/shared-python/shared/core/config/__init__.py +++ b/packages/shared-python/shared/core/config/__init__.py @@ -19,13 +19,7 @@ from .mineru import MineruConfig from .publication import ( DEFAULT_PUBLICATION_NAMESPACE_SNAPSHOT_MAX_BYTES, - DEFAULT_PUBLICATION_STRATEGY, - PUBLICATION_STRATEGIES, - PUBLICATION_STRATEGY_ENV_VAR, PublicationConfig, - PublicationStrategy, - PublicationStrategyError, - normalize_publication_strategy, ) from .qstash import QStashConfig from .redis import RedisConfig, RedisConfigManager, RedisPoolManager @@ -45,13 +39,7 @@ "AIConfig", "MineruConfig", "PublicationConfig", - "PublicationStrategy", - "PublicationStrategyError", - "DEFAULT_PUBLICATION_STRATEGY", "DEFAULT_PUBLICATION_NAMESPACE_SNAPSHOT_MAX_BYTES", - "PUBLICATION_STRATEGIES", - "PUBLICATION_STRATEGY_ENV_VAR", - "normalize_publication_strategy", "RetrievalConfig", "AppConfig", "app_config", diff --git a/packages/shared-python/shared/core/config/app.py b/packages/shared-python/shared/core/config/app.py index e73264c7c..31ce98205 100644 --- a/packages/shared-python/shared/core/config/app.py +++ b/packages/shared-python/shared/core/config/app.py @@ -10,6 +10,7 @@ from .billing import BillingConfig from .celery import CeleryConfig from .database import DatabaseConfig +from .demo import DemoConfig from .job import JobConfig from .mineru import MineruConfig from .publication import PublicationConfig @@ -32,6 +33,7 @@ class AppConfig( JobConfig, RetrievalConfig, PublicationConfig, + DemoConfig, ): """Application configuration โ€” all config components merged.""" diff --git a/packages/shared-python/shared/core/config/database.py b/packages/shared-python/shared/core/config/database.py index 53d1b97d1..6c382c8e6 100644 --- a/packages/shared-python/shared/core/config/database.py +++ b/packages/shared-python/shared/core/config/database.py @@ -11,6 +11,8 @@ class DatabaseConfig(BaseModel): """Database configuration.""" # Core database settings. + MIGRATION_DATABASE_URL: str | None = Field(default=None, description="Privileged migration connection; never used by API/Worker request sessions.") + DATABASE_URL: str = Field(..., description="Database connection URL") # SSL configuration. diff --git a/packages/shared-python/shared/core/config/demo.py b/packages/shared-python/shared/core/config/demo.py new file mode 100644 index 000000000..1369f364a --- /dev/null +++ b/packages/shared-python/shared/core/config/demo.py @@ -0,0 +1,17 @@ +"""Authorization settings for the reserved shared demo corpus.""" + +from pydantic import Field +from pydantic_settings import BaseSettings + + +class DemoConfig(BaseSettings): + DEMO_MAINTAINER_USER_IDS: str = Field( + default="", description="Comma-separated user IDs authorized to maintain demos." + ) + + def get_demo_maintainer_ids(self) -> frozenset[str]: + return frozenset( + value.strip() + for value in self.DEMO_MAINTAINER_USER_IDS.split(",") + if value.strip() + ) diff --git a/packages/shared-python/shared/core/config/publication.py b/packages/shared-python/shared/core/config/publication.py index f03cd00dd..92b52d3b2 100644 --- a/packages/shared-python/shared/core/config/publication.py +++ b/packages/shared-python/shared/core/config/publication.py @@ -1,55 +1,19 @@ -"""Temporary global Publication strategy configuration. +"""Publication measurement and snapshot size settings.""" -The publication performance plan uses exactly one process-wide strategy -switch. It has two values and no per-request, per-namespace, or percentage -routing is permitted, so the value is validated as a closed enum at settings -construction time. -""" - -from typing import Literal, cast - -from pydantic import Field, field_validator +from pydantic import Field from pydantic_settings import BaseSettings -PublicationStrategy = Literal["baseline", "candidate"] - -PUBLICATION_STRATEGY_ENV_VAR: str = "KNOWHERE_PUBLICATION_STRATEGY" -PUBLICATION_STRATEGIES: tuple[PublicationStrategy, ...] = ("baseline", "candidate") -DEFAULT_PUBLICATION_STRATEGY: PublicationStrategy = "baseline" DEFAULT_PUBLICATION_NAMESPACE_SNAPSHOT_MAX_BYTES: int = 8 * 1024 * 1024 -class PublicationStrategyError(ValueError): - """Raised when a Publication strategy value is not supported.""" - - -def normalize_publication_strategy(raw_value: object) -> PublicationStrategy: - """Validate one raw strategy value and return its canonical form.""" - value = str(raw_value if raw_value is not None else "").strip().lower() - if value not in PUBLICATION_STRATEGIES: - raise PublicationStrategyError( - f"{PUBLICATION_STRATEGY_ENV_VAR} must be one of " - f"{', '.join(PUBLICATION_STRATEGIES)}; received {raw_value!r}" - ) - return cast(PublicationStrategy, value) - - class PublicationConfig(BaseSettings): - """Publication strategy settings.""" + """Measurement-only and bounded snapshot settings.""" KNOWHERE_PUBLICATION_TRACE_ENABLED: bool = Field( default=False, description="Enable measurement-only publication traces for every attempt.", ) - KNOWHERE_PUBLICATION_STRATEGY: str = Field( - default=DEFAULT_PUBLICATION_STRATEGY, - description=( - "Temporary global Publication strategy. Only 'baseline' and " - "'candidate' are accepted; no request or namespace may override it." - ), - ) - KNOWHERE_PUBLICATION_NAMESPACE_SNAPSHOT_MAX_BYTES: int = Field( default=DEFAULT_PUBLICATION_NAMESPACE_SNAPSHOT_MAX_BYTES, ge=0, @@ -60,9 +24,3 @@ class PublicationConfig(BaseSettings): "bound." ), ) - - @field_validator("KNOWHERE_PUBLICATION_STRATEGY") - @classmethod - def validate_publication_strategy(cls, value: str) -> str: - """Reject unsupported strategy values during settings construction.""" - return normalize_publication_strategy(value) diff --git a/packages/shared-python/shared/core/config/storage.py b/packages/shared-python/shared/core/config/storage.py index a51193377..930ae06cd 100644 --- a/packages/shared-python/shared/core/config/storage.py +++ b/packages/shared-python/shared/core/config/storage.py @@ -111,20 +111,6 @@ class StorageConfig(BaseModel): "ZIP bundles remain single-object uploads." ), ) - DEMO_CANONICAL_BUNDLE_ENABLED: bool = Field( - default=False, - description=( - "Reuse immutable demo ZIP and raw artifacts by content version. " - "Disabling only changes new uploads; existing canonical results remain readable." - ), - ) - DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED: bool = Field( - default=False, - description=( - "Reuse pure tokenization and token frequencies for versioned canonical " - "demo publication, in a bounded process-local cache. Requires canonical bundles." - ), - ) SUPPORTED_EXTENSIONS: str = Field( default=".doc,.docx,.pdf,.txt,.xls,.xlsx,.pptx,.jpg,.jpeg,.png,.md,.html,.htm", description="Supported file extensions", diff --git a/packages/shared-python/shared/core/exceptions/demo_materialization_removed.py b/packages/shared-python/shared/core/exceptions/demo_materialization_removed.py new file mode 100644 index 000000000..f0aeccfd2 --- /dev/null +++ b/packages/shared-python/shared/core/exceptions/demo_materialization_removed.py @@ -0,0 +1,9 @@ +"""Migration response for the retired personal demo copy operation.""" + +from shared.core.exceptions.knowhere_exception import KnowhereException +from shared.core.response.ErrorCode import ErrorCode + + +class DemoMaterializationRemovedException(KnowhereException): + def __init__(self) -> None: + super().__init__(code=ErrorCode.DEMO_MATERIALIZATION_REMOVED, internal_message="Personal demo materialization has been retired.", user_message="Use catalog canonical_document_id values to query namespace __knowhere_demo__. Existing personal copies remain accessible through the document API.", details={"namespace": "__knowhere_demo__"}, http_status_code=410) diff --git a/packages/shared-python/shared/core/exceptions/demo_not_ready.py b/packages/shared-python/shared/core/exceptions/demo_not_ready.py new file mode 100644 index 000000000..8950ca6a2 --- /dev/null +++ b/packages/shared-python/shared/core/exceptions/demo_not_ready.py @@ -0,0 +1,9 @@ +"""Public readiness failure for the shared demo corpus.""" + +from shared.core.exceptions.knowhere_exception import KnowhereException +from shared.core.response.ErrorCode import ErrorCode + + +class DemoNotReadyException(KnowhereException): + def __init__(self, *, is_explicit_source: bool) -> None: + super().__init__(code=ErrorCode.DEMO_NOT_READY, internal_message="Demo source has no complete published revision.", user_message="Demo documents are preparing. Retry after publication completes.", http_status_code=409 if is_explicit_source else 503, details={"retry_after": 30}) diff --git a/packages/shared-python/shared/core/response/ErrorCode.py b/packages/shared-python/shared/core/response/ErrorCode.py index cf5dd2f8f..7c7ef1d25 100644 --- a/packages/shared-python/shared/core/response/ErrorCode.py +++ b/packages/shared-python/shared/core/response/ErrorCode.py @@ -63,6 +63,9 @@ class ErrorCode(str, Enum): ) CANCELLED = "CANCELLED" # 499 - Client cancelled request + DEMO_NOT_READY = "DEMO_NOT_READY" + DEMO_MATERIALIZATION_REMOVED = "DEMO_MATERIALIZATION_REMOVED" + # Server Errors (5xx) UNKNOWN = "UNKNOWN" # 500 - Unknown error (fallback) INTERNAL_ERROR = "INTERNAL_ERROR" # 500 - Internal invariants broken @@ -83,6 +86,8 @@ class ErrorCodeMapper: _error_code_to_http_status = { ErrorCode.OK: 200, + ErrorCode.DEMO_NOT_READY: 503, + ErrorCode.DEMO_MATERIALIZATION_REMOVED: 410, ErrorCode.INVALID_ARGUMENT: 400, ErrorCode.FAILED_PRECONDITION: 400, ErrorCode.OUT_OF_RANGE: 400, diff --git a/packages/shared-python/shared/models/database/__init__.py b/packages/shared-python/shared/models/database/__init__.py index 5535cd9bf..7426b8262 100644 --- a/packages/shared-python/shared/models/database/__init__.py +++ b/packages/shared-python/shared/models/database/__init__.py @@ -83,3 +83,5 @@ "SystemLimit", # "OAuthProvider" ] + +from . import demo_corpus # noqa: F401 diff --git a/packages/shared-python/shared/models/database/demo_corpus.py b/packages/shared-python/shared/models/database/demo_corpus.py new file mode 100644 index 000000000..8e9231083 --- /dev/null +++ b/packages/shared-python/shared/models/database/demo_corpus.py @@ -0,0 +1,716 @@ +"""Dedicated shared demo persistence records; private tables remain isolated.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any, Dict, Optional +from uuid import uuid4 + +from sqlalchemy import ( + JSON, + CheckConstraint, + ForeignKeyConstraint, + DateTime, + Float, + ForeignKey, + BigInteger, + Index, + Integer, + LargeBinary, + String, + Text, + text, + UniqueConstraint, +) +from sqlalchemy.orm import Mapped, mapped_column + +from shared.core.database import Base +from shared.utils.utc_now import utc_now_naive + + +class DemoDocument(Base): + """Shared source directory and its current complete publication revision.""" + + __tablename__ = "demo_documents" + + document_id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: f"ddoc_{uuid4().hex[:12]}" + ) + demo_source_id: Mapped[str] = mapped_column( + String(128), nullable=False, unique=True + ) + title: Mapped[str] = mapped_column(Text, nullable=False) + category_id: Mapped[str] = mapped_column( + String(128), nullable=False, default="other" + ) + catalog_metadata: Mapped[dict[str, Any]] = mapped_column( + JSON, nullable=False, default=dict + ) + + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column( + String(255), nullable=False, default="__knowhere_demo__" + ) + status: Mapped[str] = mapped_column(String(32), nullable=False, default="active") + current_job_result_id: Mapped[Optional[str]] = mapped_column( + String(36), + ForeignKey("job_results.id", ondelete="RESTRICT", use_alter=True), + nullable=True, + ) + source_file_name: Mapped[Optional[str]] = mapped_column(Text, nullable=True) + document_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column( + JSON, nullable=True + ) + parse_track: Mapped[str] = mapped_column( + String(32), nullable=False, default="chunk" + ) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime, + default=utc_now_naive, + onupdate=utc_now_naive, + nullable=False, + ) + archived_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True) + + __table_args__ = ( + ForeignKeyConstraint( + ["current_job_result_id", "document_id"], + ["job_results.id", "job_results.demo_document_id"], + name="fk_demo_documents_current_revision", + deferrable=True, + initially="DEFERRED", + use_alter=True, + ), + CheckConstraint( + "namespace = '__knowhere_demo__'", name="ck_demo_documents_namespace" + ), + Index( + "idx_demo_documents_user_namespace_status", "user_id", "namespace", "status" + ), + Index("idx_demo_documents_current_job_result", "current_job_result_id"), + ) + + +class DemoDocumentSection(Base): + """Canonical hierarchy node for one published document revision.""" + + __tablename__ = "demo_document_sections" + + section_id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: f"sec_{uuid4().hex[:12]}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column( + String(255), nullable=False, default="default" + ) + document_id: Mapped[str] = mapped_column( + String(36), + ForeignKey("demo_documents.document_id", ondelete="CASCADE"), + nullable=False, + ) + job_result_id: Mapped[str] = mapped_column( + String(36), ForeignKey("job_results.id", ondelete="RESTRICT"), nullable=False + ) + parent_section_id: Mapped[Optional[str]] = mapped_column( + String(36), + ForeignKey("demo_document_sections.section_id", ondelete="SET NULL"), + nullable=True, + ) + section_path: Mapped[str] = mapped_column(Text, nullable=False) + section_title: Mapped[Optional[str]] = mapped_column(Text, nullable=True) + section_level: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + summary: Mapped[Optional[str]] = mapped_column(Text, nullable=True) + section_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column( + JSON, nullable=True + ) + sort_order: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + + __table_args__ = ( + UniqueConstraint( + "section_id", + "document_id", + "job_result_id", + name="uq_demo_sections_identity", + ), + ForeignKeyConstraint( + ["parent_section_id", "document_id", "job_result_id"], + [ + "demo_document_sections.section_id", + "demo_document_sections.document_id", + "demo_document_sections.job_result_id", + ], + deferrable=True, + initially="DEFERRED", + name="fk_demo_sections_parent_revision", + ), + ForeignKeyConstraint( + ["job_result_id", "document_id"], + ["job_results.id", "job_results.demo_document_id"], + ondelete="RESTRICT", + name="fk_demo_documentsection_revision", + ), + UniqueConstraint( + "document_id", + "job_result_id", + "section_path", + name="uq_demo_document_sections_revision_path", + ), + Index("idx_demo_document_sections_scope", "user_id", "namespace"), + Index( + "idx_demo_document_sections_doc_revision", "document_id", "job_result_id" + ), + Index( + "idx_demo_document_sections_revision_snapshot_order", + "document_id", + "job_result_id", + "sort_order", + "section_id", + ), + ) + + +class DemoDocumentChunk(Base): + """Canonical retrieval payload row for one published document revision.""" + + __tablename__ = "demo_document_chunks" + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: f"dchk_{uuid4().hex[:12]}" + ) + + chunk_id: Mapped[str] = mapped_column(String(64), nullable=False) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column( + String(255), nullable=False, default="default" + ) + document_id: Mapped[str] = mapped_column( + String(36), + ForeignKey("demo_documents.document_id", ondelete="CASCADE"), + nullable=False, + ) + job_result_id: Mapped[str] = mapped_column( + String(36), ForeignKey("job_results.id", ondelete="RESTRICT"), nullable=False + ) + section_id: Mapped[Optional[str]] = mapped_column( + String(36), + ForeignKey("demo_document_sections.section_id", ondelete="SET NULL"), + nullable=True, + ) + chunk_type: Mapped[str] = mapped_column(String(64), nullable=False, default="text") + content: Mapped[Optional[str]] = mapped_column(Text, nullable=True) + content_lexical_text: Mapped[Optional[str]] = mapped_column(Text, nullable=True) + path_lexical_text: Mapped[Optional[str]] = mapped_column(Text, nullable=True) + content_search_text: Mapped[Optional[str]] = mapped_column(Text, nullable=True) + path_search_text: Mapped[Optional[str]] = mapped_column(Text, nullable=True) + term_search_text: Mapped[Optional[str]] = mapped_column(Text, nullable=True) + source_chunk_path: Mapped[Optional[str]] = mapped_column(Text, nullable=True) + file_path: Mapped[Optional[str]] = mapped_column(Text, nullable=True) + chunk_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column( + JSON, nullable=True + ) + sort_order: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + + __table_args__ = ( + ForeignKeyConstraint( + ["section_id", "document_id", "job_result_id"], + [ + "demo_document_sections.section_id", + "demo_document_sections.document_id", + "demo_document_sections.job_result_id", + ], + deferrable=True, + initially="DEFERRED", + name="fk_demo_documentchunk_section_revision", + ), + ForeignKeyConstraint( + ["job_result_id", "document_id"], + ["job_results.id", "job_results.demo_document_id"], + ondelete="RESTRICT", + name="fk_demo_documentchunk_revision", + ), + UniqueConstraint( + "document_id", + "job_result_id", + "source_chunk_path", + name="uq_demo_document_chunks_revision_path", + ), + Index("idx_demo_document_chunks_scope", "user_id", "namespace"), + Index("idx_demo_document_chunks_chunk_id", "chunk_id"), + Index("idx_demo_document_chunks_doc_revision", "document_id", "job_result_id"), + Index( + "idx_demo_document_chunks_revision_snapshot_order", + "document_id", + "job_result_id", + "sort_order", + "chunk_id", + "id", + ), + Index( + "idx_demo_document_chunks_revision_section_order", + "document_id", + "job_result_id", + "section_id", + "sort_order", + "chunk_id", + "id", + ), + Index("idx_demo_document_chunks_section", "section_id"), + ) + + +class DemoDocumentMapUnit(Base): + """Persisted lexical map unit for one document revision. + + These rows are a derived index of the exact leaf and interstitial units + used by map-nav. Full chunk payloads remain in ``demo_document_chunks`` and are + loaded separately for evidence hydration. + """ + + __tablename__ = "demo_document_map_units" + + id: Mapped[str] = mapped_column(String(160), primary_key=True) + document_id: Mapped[str] = mapped_column( + String(36), + ForeignKey("demo_documents.document_id", ondelete="CASCADE"), + nullable=False, + ) + job_result_id: Mapped[str] = mapped_column( + String(36), ForeignKey("job_results.id", ondelete="RESTRICT"), nullable=False + ) + unit_id: Mapped[str] = mapped_column(String(128), nullable=False) + section_id: Mapped[str] = mapped_column(String(36), nullable=False) + unit_kind: Mapped[str] = mapped_column(String(32), nullable=False) + path_token_count: Mapped[int] = mapped_column(Integer, nullable=False) + content_token_count: Mapped[int] = mapped_column(Integer, nullable=False) + # Asset presence under this unit's section, after root-asset remount + # (``KnowhereProvider._remount_root_assets``). Lets type-scoped queries + # (e.g. chunk_types=["image"]) narrow map-unit candidates *before* + # scoring instead of scoring everything and discarding after the fact. + has_image: Mapped[bool] = mapped_column(nullable=False, default=False) + has_table: Mapped[bool] = mapped_column(nullable=False, default=False) + sort_order: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + + __table_args__ = ( + ForeignKeyConstraint( + ["section_id", "document_id", "job_result_id"], + [ + "demo_document_sections.section_id", + "demo_document_sections.document_id", + "demo_document_sections.job_result_id", + ], + deferrable=True, + initially="DEFERRED", + name="fk_demo_documentmapunit_section_revision", + ), + ForeignKeyConstraint( + ["job_result_id", "document_id"], + ["job_results.id", "job_results.demo_document_id"], + ondelete="RESTRICT", + name="fk_demo_documentmapunit_revision", + ), + Index( + "idx_demo_document_map_units_revision_order", + "document_id", + "job_result_id", + "sort_order", + "unit_id", + ), + Index("idx_demo_document_map_units_section", "section_id"), + Index( + "idx_demo_document_map_units_has_image", + "document_id", + "job_result_id", + postgresql_where=has_image.is_(True), + ), + Index( + "idx_demo_document_map_units_has_table", + "document_id", + "job_result_id", + postgresql_where=has_table.is_(True), + ), + ) + + +class DemoDocumentMapUnitToken(Base): + """One exact token frequency in a persisted map unit channel.""" + + __tablename__ = "demo_document_map_unit_tokens" + + id: Mapped[str] = mapped_column(String(36), primary_key=True) + map_unit_id: Mapped[str] = mapped_column( + String(160), + ForeignKey( + "demo_document_map_units.id", + ondelete="CASCADE", + deferrable=True, + initially="DEFERRED", + ), + nullable=False, + ) + channel: Mapped[str] = mapped_column(String(16), nullable=False) + token: Mapped[str] = mapped_column(Text, nullable=False) + token_hash: Mapped[str] = mapped_column(String(64), nullable=False) + frequency: Mapped[int] = mapped_column(Integer, nullable=False) + + __table_args__ = ( + Index( + "idx_demo_document_map_unit_tokens_token_lookup_binary", + "channel", + text("decode(token_hash, 'hex'::text)"), + postgresql_include=["map_unit_id", "token", "frequency"], + ), + Index( + "idx_demo_document_map_unit_tokens_unit", + "map_unit_id", + "channel", + postgresql_include=["token", "frequency"], + ), + ) + + +class DemoDocumentMapUnitIndex(Base): + """Completeness marker and corpus statistics for a materialized index.""" + + __tablename__ = "demo_document_map_unit_indexes" + + id: Mapped[str] = mapped_column(String(100), primary_key=True) + document_id: Mapped[str] = mapped_column( + String(36), + ForeignKey("demo_documents.document_id", ondelete="CASCADE"), + nullable=False, + ) + job_result_id: Mapped[str] = mapped_column( + String(36), ForeignKey("job_results.id", ondelete="RESTRICT"), nullable=False + ) + format_version: Mapped[int] = mapped_column(Integer, nullable=False) + unit_count: Mapped[int] = mapped_column(Integer, nullable=False) + token_count: Mapped[int] = mapped_column(Integer, nullable=False) + # rank_bm25 Okapi average IDF for the revision's units (path/content). + # Written at index time so query scoring never rescans all tokens. + average_idf_path: Mapped[float] = mapped_column(Float, nullable=False, default=0.0) + average_idf_content: Mapped[float] = mapped_column( + Float, nullable=False, default=0.0 + ) + path_document_count: Mapped[Optional[int]] = mapped_column(Integer, nullable=True) + path_total_length: Mapped[Optional[int]] = mapped_column(Integer, nullable=True) + content_document_count: Mapped[Optional[int]] = mapped_column( + Integer, nullable=True + ) + content_total_length: Mapped[Optional[int]] = mapped_column(Integer, nullable=True) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + + __table_args__ = ( + ForeignKeyConstraint( + ["job_result_id", "document_id"], + ["job_results.id", "job_results.demo_document_id"], + ondelete="RESTRICT", + name="fk_demo_documentmapunitindex_revision", + ), + UniqueConstraint( + "document_id", + "job_result_id", + name="uq_demo_document_map_unit_indexes_revision", + ), + Index( + "idx_demo_document_map_unit_indexes_revision", + "document_id", + "job_result_id", + ), + ) + + +class DemoRetrievalNamespaceGeneration(Base): + """Monotonic serving generation for one user-owned namespace.""" + + __tablename__ = "demo_retrieval_namespace_generations" + + id: Mapped[str] = mapped_column( + String(100), primary_key=True, default=lambda: f"rng_{uuid4().hex}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column(String(255), nullable=False) + generation: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0) + updated_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, onupdate=utc_now_naive, nullable=False + ) + + __table_args__ = ( + UniqueConstraint( + "user_id", + "namespace", + name="uq_demo_retrieval_namespace_generations_scope", + ), + ) + + +class DemoRetrievalServingRevisionManifest(Base): + """Compressed ordered metadata for one document revision.""" + + __tablename__ = "demo_retrieval_serving_revision_manifests" + + id: Mapped[str] = mapped_column( + String(100), primary_key=True, default=lambda: f"rsm_{uuid4().hex}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column(String(255), nullable=False) + document_id: Mapped[str] = mapped_column( + String(36), + ForeignKey("demo_documents.document_id", ondelete="CASCADE"), + nullable=False, + ) + job_result_id: Mapped[str] = mapped_column( + String(36), ForeignKey("job_results.id", ondelete="RESTRICT"), nullable=False + ) + format_version: Mapped[int] = mapped_column(Integer, nullable=False) + payload_zlib: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + checksum: Mapped[str] = mapped_column(String(64), nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + + __table_args__ = ( + ForeignKeyConstraint( + ["job_result_id", "document_id"], + ["job_results.id", "job_results.demo_document_id"], + ondelete="RESTRICT", + name="fk_demo_retrievalservingrevisionmanifest_revision", + ), + UniqueConstraint( + "document_id", + "job_result_id", + name="uq_demo_retrieval_serving_revision_manifests_revision", + ), + Index( + "idx_demo_retrieval_serving_revision_manifests_scope", + "user_id", + "namespace", + "document_id", + "job_result_id", + ), + ) + + +class DemoRetrievalNamespaceMapSnapshot(Base): + """Persisted namespace-level MAP (sections + chunk index + map units). + + Incrementally patched at publish/archive time (one document's subtree at + a time); query time reads this row directly instead of merging per-file + manifests. Overwritten in place -- no generation history is retained. + """ + + __tablename__ = "demo_retrieval_namespace_map_snapshots" + + id: Mapped[str] = mapped_column( + String(100), primary_key=True, default=lambda: f"rnmap_{uuid4().hex}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column(String(255), nullable=False) + generation: Mapped[int] = mapped_column(BigInteger, nullable=False) + format_version: Mapped[int] = mapped_column(Integer, nullable=False) + payload_zlib: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + checksum: Mapped[str] = mapped_column(String(64), nullable=False) + updated_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, onupdate=utc_now_naive, nullable=False + ) + + __table_args__ = ( + UniqueConstraint( + "user_id", + "namespace", + name="uq_demo_retrieval_namespace_map_snapshots_scope", + ), + ) + + +class DemoGraphNode(Base): + """Persisted derived graph node used for routing and expansion.""" + + __tablename__ = "demo_graph_nodes" + + node_id: Mapped[str] = mapped_column(String(128), primary_key=True) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column( + String(255), nullable=False, default="default" + ) + node_kind: Mapped[str] = mapped_column(String(32), nullable=False) + owner_document_id: Mapped[str] = mapped_column( + String(36), + ForeignKey("demo_documents.document_id", ondelete="CASCADE"), + nullable=False, + ) + job_result_id: Mapped[str] = mapped_column( + String(36), ForeignKey("job_results.id", ondelete="RESTRICT"), nullable=False + ) + ref_document_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True) + ref_section_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True) + properties: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime, + default=utc_now_naive, + onupdate=utc_now_naive, + nullable=False, + ) + + __table_args__ = ( + ForeignKeyConstraint( + ["job_result_id", "owner_document_id"], + ["job_results.id", "job_results.demo_document_id"], + ondelete="RESTRICT", + name="fk_demo_graphnode_revision", + ), + Index("idx_demo_graph_nodes_scope", "user_id", "namespace", "node_kind"), + Index( + "idx_demo_graph_nodes_owner_revision", "owner_document_id", "job_result_id" + ), + Index("idx_demo_graph_nodes_ref_document", "ref_document_id"), + Index("idx_demo_graph_nodes_ref_section", "ref_section_id"), + Index( + "idx_demo_graph_nodes_top_keywords_gin", + text("(properties::jsonb -> 'top_keywords')"), + postgresql_using="gin", + ), + Index( + "idx_demo_graph_nodes_top_entities_gin", + text("(properties::jsonb -> 'top_entities')"), + postgresql_using="gin", + ), + ) + + +class DemoGraphEdge(Base): + """Persisted derived graph edge used for routing and expansion.""" + + __tablename__ = "demo_graph_edges" + + edge_id: Mapped[str] = mapped_column(String(160), primary_key=True) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column( + String(255), nullable=False, default="default" + ) + edge_kind: Mapped[str] = mapped_column(String(32), nullable=False) + source_node_id: Mapped[str] = mapped_column( + String(128), + ForeignKey("demo_graph_nodes.node_id", ondelete="CASCADE"), + nullable=False, + ) + target_node_id: Mapped[str] = mapped_column( + String(128), + ForeignKey("demo_graph_nodes.node_id", ondelete="CASCADE"), + nullable=False, + ) + owner_document_id: Mapped[str] = mapped_column( + String(36), + ForeignKey("demo_documents.document_id", ondelete="CASCADE"), + nullable=False, + ) + job_result_id: Mapped[str] = mapped_column( + String(36), ForeignKey("job_results.id", ondelete="RESTRICT"), nullable=False + ) + is_directed: Mapped[bool] = mapped_column(nullable=False, default=True) + weight: Mapped[Optional[float]] = mapped_column(Float, nullable=True) + properties: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime, + default=utc_now_naive, + onupdate=utc_now_naive, + nullable=False, + ) + + __table_args__ = ( + ForeignKeyConstraint( + ["job_result_id", "owner_document_id"], + ["job_results.id", "job_results.demo_document_id"], + ondelete="RESTRICT", + name="fk_demo_graphedge_revision", + ), + Index("idx_demo_graph_edges_scope", "user_id", "namespace", "edge_kind"), + Index( + "idx_demo_graph_edges_owner_revision", "owner_document_id", "job_result_id" + ), + Index("idx_demo_graph_edges_source", "source_node_id"), + Index("idx_demo_graph_edges_target", "target_node_id"), + ) + + +class DemoRetrievalHitStat(Base): + """Append-only retrieval usage analytics row.""" + + __tablename__ = "demo_retrieval_hit_stats" + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: f"rhs_{uuid4().hex[:12]}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column( + String(255), nullable=False, default="default" + ) + hit_kind: Mapped[str] = mapped_column(String(32), nullable=False) + document_id: Mapped[str] = mapped_column( + String(36), + ForeignKey("demo_documents.document_id", ondelete="CASCADE"), + nullable=False, + ) + chunk_id: Mapped[Optional[str]] = mapped_column(String(64), nullable=True) + hit_count: Mapped[int] = mapped_column(Integer, nullable=False, default=1) + last_hit_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime, + default=utc_now_naive, + onupdate=utc_now_naive, + nullable=False, + ) + + __table_args__ = ( + Index( + "uq_demo_retrieval_hit_stats_document_key", + "user_id", + "namespace", + "hit_kind", + "document_id", + unique=True, + postgresql_where=chunk_id.is_(None), + ), + Index( + "uq_demo_retrieval_hit_stats_chunk_key", + "user_id", + "namespace", + "hit_kind", + "document_id", + "chunk_id", + unique=True, + postgresql_where=chunk_id.is_not(None), + ), + Index( + "idx_demo_retrieval_hit_stats_scope_kind", + "user_id", + "namespace", + "hit_kind", + ), + Index("idx_demo_retrieval_hit_stats_document", "document_id"), + Index("idx_demo_retrieval_hit_stats_chunk", "chunk_id"), + ) diff --git a/packages/shared-python/shared/models/database/job_result.py b/packages/shared-python/shared/models/database/job_result.py index f88f3b98b..7e0bd9604 100644 --- a/packages/shared-python/shared/models/database/job_result.py +++ b/packages/shared-python/shared/models/database/job_result.py @@ -6,7 +6,7 @@ from typing import Any, Dict, List, Optional from uuid import uuid4 -from sqlalchemy import JSON, DateTime, ForeignKey, Index, Integer, String, Text +from sqlalchemy import CheckConstraint, UniqueConstraint, JSON, DateTime, ForeignKey, Index, Integer, String, Text from sqlalchemy.orm import Mapped, mapped_column, relationship from shared.core.database import Base @@ -34,6 +34,15 @@ class JobResult(Base): index=True, ) + demo_document_id: Mapped[str | None] = mapped_column( + String(36), ForeignKey("demo_documents.document_id", ondelete="RESTRICT"), + nullable=True, index=True, + ) + __table_args__ = ( + CheckConstraint("document_id IS NULL OR demo_document_id IS NULL", name="ck_job_results_corpus_target"), + UniqueConstraint("id", "demo_document_id", name="uq_job_results_demo_revision"), + ) + delivery_mode: Mapped[str] = mapped_column(String(20), nullable=False) # inline/url document_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column( JSON, nullable=True diff --git a/packages/shared-python/shared/models/schemas/job.py b/packages/shared-python/shared/models/schemas/job.py index 1ea775fde..e16bb65ff 100644 --- a/packages/shared-python/shared/models/schemas/job.py +++ b/packages/shared-python/shared/models/schemas/job.py @@ -150,6 +150,8 @@ class StandardErrorObject(BaseModel): class JobResultResponse(BaseModel): """Job status query response (for GET /jobs/{job_id}/result)""" + job_result_id: Optional[str] = Field(None, description="Immutable completed publication revision, when available") + job_id: str = Field(..., description="Job ID") namespace: Optional[str] = Field(None, description="Effective retrieval namespace") document_id: Optional[str] = Field(None, description="Linked document ID") diff --git a/packages/shared-python/shared/services/ai/prompt_service.py b/packages/shared-python/shared/services/ai/prompt_service.py index 808b816c9..b850b7c87 100755 --- a/packages/shared-python/shared/services/ai/prompt_service.py +++ b/packages/shared-python/shared/services/ai/prompt_service.py @@ -585,14 +585,14 @@ def build_prompt(task, texts, query, **kwargs): - TOC/index headings such as "Contents", "็›ฎๅฝ•", "็›ฎๆฌก"; - front matter such as "ๅ‰่จ€" when it is outside the segment's body outline. 4. When two rows are near-duplicates, keep the clearer/more complete one and - omit the duplicate from output. + mark the duplicate with level 0. 5. Do not invent headings. Only return ids that exist in the input. Output requirements: - Output ONLY a valid JSON array. No markdown fences, no explanations. -- Include each retained heading as: - {{"id": , "level": }} -- Omitted ids are treated as filtered noise. +- Classify EVERY input id exactly once, including noise: + {{"id": , "level": }} +- Use level 0 only for filtered noise. Never omit an id. """ elif task == "page-memory-node-summary": diff --git a/packages/shared-python/shared/services/jobs/lifecycle/result_writer.py b/packages/shared-python/shared/services/jobs/lifecycle/result_writer.py index d1edc43b1..af373d64a 100644 --- a/packages/shared-python/shared/services/jobs/lifecycle/result_writer.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/result_writer.py @@ -6,6 +6,7 @@ from sqlalchemy import delete, insert, select from sqlalchemy.orm import Session +from shared.models.database.job import Job from shared.models.database.job_result import JobChunk, JobResult from shared.utils.json_utils import remove_nul_characters @@ -34,10 +35,13 @@ def upsert_job_result( db.flush() return existing + job = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() + metadata = job.job_metadata or {} if job is not None else {} + demoMetadata = {key: metadata[key] for key in ("source_size_bytes", "demo_asset_manifest", "result_raw_prefix", "parse_track") if key in metadata} if metadata.get("corpus_target") == "DEMO" else {} job_result = JobResult( job_id=job_id, delivery_mode=delivery_mode, - document_metadata={}, + document_metadata=demoMetadata, inline_payload=inline_payload, result_s3_key=result_s3_key, result_size=result_size, diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/dispatch.py b/packages/shared-python/shared/services/retrieval/agent_explore/dispatch.py index ffceac8f3..f002472bd 100644 --- a/packages/shared-python/shared/services/retrieval/agent_explore/dispatch.py +++ b/packages/shared-python/shared/services/retrieval/agent_explore/dispatch.py @@ -25,6 +25,7 @@ from typing import Any from sqlalchemy.ext.asyncio import AsyncSession +from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext from shared.services.retrieval.agent_tools import REGISTRY, ToolBudget, ToolContext, ToolResult @@ -55,13 +56,14 @@ async def dispatch_tool_call( """ try: async with db_factory() as db: + pins = CorpusRevisionContext.get_pins() tool_ctx = ToolContext( db=db, user_id=user_id, namespace=namespace, db_factory=db_factory, budget=budget or ToolBudget(), - document_scope=document_scope, + document_scope=document_scope.narrow(list(pins)) if pins is not None else document_scope, query=query, ) return await REGISTRY.dispatch(name, tool_ctx, args) diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py b/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py index 57124ca85..1d84fdabd 100644 --- a/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py +++ b/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py @@ -27,7 +27,7 @@ ``messages: list[dict]`` history, a mechanism the Cursor SDK harness has no equivalent hook for: -1. **Stale tool-message collapsing** (``_TOOL_MESSAGE_FRESH_TURNS``, +1. **Stale discovery-message collapsing** (``_TOOL_MESSAGE_FRESH_TURNS``, ``_collapse_stale_tool_messages``): ``messages`` only ever appended, so a single ``corpus.outline``/``corpus.node_filter`` call (each capped at ``ToolBudget.max_chars`` โ€” currently ``EVIDENCE_TEXT_CHAR_BUDGET=12_000``, @@ -36,6 +36,10 @@ ``RETRIEVAL_NAV_TOKEN_LIMIT`` (100k default) within 7-8 LLM turns from this resend alone, not from query difficulty โ€” per-turn token cost grew monotonically (q04: 4.4k -> 4.8k -> 19.8k -> 21.5k -> 24.2k -> 29.6k). + Successful evidence reads remain visible until finish. Removing them after + two turns caused a real SpaceX query to alternate between complementary + sections until its token budget ran out. Per-result text caps and the + episode budget still apply. 2. **Trajectory refs fallback** (``shared.select_episode_refs``): verified live that ``finish`` can be called with no ``refs`` key at all (raw ``function.arguments`` was literally ``'{}'``) even after the model had @@ -85,7 +89,7 @@ load_corpus_schema_text, ) -# A tool-role message is kept in full for the turn it was produced plus this +# A discovery or failed tool message is kept in full for the turn it was produced plus this # many additional turns, then collapsed to a placeholder โ€” see module # docstring point 1. Not tuned against a real recall-vs-token tradeoff yet; # 2 was chosen so a result stays fully visible for one full turn after the @@ -163,14 +167,15 @@ def _collapse_stale_tool_messages( current_turn: int, fresh_turns: int, ) -> None: - """Replace tool messages older than ``fresh_turns`` with a placeholder. + """Replace stale discovery/error messages while preserving read evidence. ``messages`` only ever grows within one episode (see module docstring point 1); this is what keeps that growth bounded instead of resending - every past tool result on every later turn. + every past discovery result on every later turn. Successful evidence must + remain visible so complementary reads can be synthesized without rereading. """ for entry in tool_message_log: - if entry["collapsed"]: + if entry["collapsed"] or entry.get("has_evidence", False): continue if current_turn - entry["turn_index"] < fresh_turns: continue @@ -213,7 +218,7 @@ async def run_episode( # module docstring point 2). trajectory_refs: list[dict[str, Any]] = [] # One entry per appended tool-role message: {message_index, turn_index, - # tool_name, original_chars, collapsed} โ€” see _collapse_stale_tool_messages. + # tool_name, original_chars, has_evidence, collapsed} โ€” see _collapse_stale_tool_messages. tool_message_log: list[dict[str, Any]] = [] turn_index = 0 @@ -381,6 +386,11 @@ async def run_episode( "turn_index": turn_index, "tool_name": canonical_name, "original_chars": len(content), + "has_evidence": ( + canonical_name in EVIDENCE_TOOL_NAMES + and not tool_result.error + and bool(tool_result.refs) + ), "collapsed": False, } ) diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/ref_resolution.py b/packages/shared-python/shared/services/retrieval/agent_explore/ref_resolution.py index ded7aac8f..d244e41ea 100644 --- a/packages/shared-python/shared/services/retrieval/agent_explore/ref_resolution.py +++ b/packages/shared-python/shared/services/retrieval/agent_explore/ref_resolution.py @@ -19,6 +19,10 @@ from __future__ import annotations +from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext + +from shared.services.retrieval.corpus_storage import CorpusStorage + from shared.services.retrieval.document_scope import DocumentScope from dataclasses import dataclass, field @@ -27,7 +31,6 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.services.retrieval.agent_tools.section_path_lookup import ( resolve_section_path_anchor, ) @@ -50,6 +53,7 @@ async def resolve_finish_refs( document_scope: DocumentScope = DocumentScope(), ) -> FinishRefResolution: """Return refs with ``chunk_id`` populated; dropped refs keep a reason.""" + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) resolution = FinishRefResolution() scoped: list[dict[str, Any]] = [] for ref in refs: @@ -73,19 +77,19 @@ async def resolve_finish_refs( documents = ( ( await db.execute( - select(Document) - .where(Document.document_id.in_(document_ids)) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == "active") - .where(document_scope.predicate(Document.document_id)) + select(corpusStorage.Document) + .where(corpusStorage.Document.document_id.in_(document_ids)) + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.Document.namespace == namespace) + .where(corpusStorage.Document.status == "active") + .where(document_scope.predicate(corpusStorage.Document.document_id)) ) ) .scalars() .all() ) revision_by_doc = { - d.document_id: d.current_job_result_id for d in documents if d.current_job_result_id + d.document_id: revision for d in documents if (revision := CorpusRevisionContext.resolve_revision(d.document_id, d.current_job_result_id)) } for ref in scoped: @@ -125,16 +129,16 @@ async def resolve_finish_refs( row = ( await db.execute( - select(DocumentChunk.chunk_id) - .select_from(DocumentChunk) + select(corpusStorage.DocumentChunk.chunk_id) + .select_from(corpusStorage.DocumentChunk) .join( - DocumentSection, - DocumentSection.section_id == DocumentChunk.section_id, + corpusStorage.DocumentSection, + corpusStorage.DocumentSection.section_id == corpusStorage.DocumentChunk.section_id, ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentSection.section_path == resolved_path) - .where(DocumentChunk.chunk_type.in_(_BODY_CHUNK_TYPES)) + .where(corpusStorage.DocumentChunk.document_id == document_id) + .where(corpusStorage.DocumentChunk.job_result_id == job_result_id) + .where(corpusStorage.DocumentSection.section_path == resolved_path) + .where(corpusStorage.DocumentChunk.chunk_type.in_(_BODY_CHUNK_TYPES)) ) ).first() if row is None: diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/asset_hosts.py b/packages/shared-python/shared/services/retrieval/agent_tools/asset_hosts.py index 8358cdf65..337793b97 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/asset_hosts.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/asset_hosts.py @@ -11,13 +11,16 @@ from __future__ import annotations +from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext + +from shared.services.retrieval.corpus_storage import CorpusStorage + from collections.abc import Iterable from dataclasses import dataclass from typing import Any from sqlalchemy import select -from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.services.retrieval.agent_tools.registry import ToolContext from shared.services.retrieval.agent_tools.scope import ScopeTarget from shared.services.retrieval.hydration.row_utils import iter_connected_target_ids @@ -66,30 +69,31 @@ async def load_asset_hosts( ``document_ids=None`` scans every document visible to the caller; ``asset_ids`` limits the result to those asset chunk ids. """ + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(ctx.namespace) wanted = set(asset_ids) if asset_ids is not None else None stmt = ( select( - DocumentChunk.document_id, - DocumentChunk.chunk_id, - DocumentChunk.chunk_metadata, - DocumentSection.section_path, - Document.source_file_name, + corpusStorage.DocumentChunk.document_id, + corpusStorage.DocumentChunk.chunk_id, + corpusStorage.DocumentChunk.chunk_metadata, + corpusStorage.DocumentSection.section_path, + corpusStorage.Document.source_file_name, ) - .select_from(DocumentChunk) - .join(Document, Document.document_id == DocumentChunk.document_id) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .select_from(corpusStorage.DocumentChunk) + .join(corpusStorage.Document, corpusStorage.Document.document_id == corpusStorage.DocumentChunk.document_id) + .outerjoin(corpusStorage.DocumentSection, corpusStorage.DocumentSection.section_id == corpusStorage.DocumentChunk.section_id) .where( - Document.user_id == ctx.user_id, - Document.namespace == ctx.namespace, - Document.status == "active", - Document.current_job_result_id == DocumentChunk.job_result_id, - ctx.document_scope.predicate(Document.document_id), - DocumentChunk.chunk_type.in_(_BODY_CHUNK_TYPES), + corpusStorage.Document.user_id == corpusStorage.resolve_owner(ctx.user_id), + corpusStorage.Document.namespace == ctx.namespace, + corpusStorage.Document.status == "active", + CorpusRevisionContext.build_revision_column(corpusStorage.Document) == corpusStorage.DocumentChunk.job_result_id, + ctx.document_scope.predicate(corpusStorage.Document.document_id), + corpusStorage.DocumentChunk.chunk_type.in_(_BODY_CHUNK_TYPES), ) - .order_by(DocumentChunk.document_id, DocumentChunk.sort_order, DocumentChunk.chunk_id) + .order_by(corpusStorage.DocumentChunk.document_id, corpusStorage.DocumentChunk.sort_order, corpusStorage.DocumentChunk.chunk_id) ) if document_ids is not None: - stmt = stmt.where(Document.document_id.in_(sorted(set(document_ids)))) + stmt = stmt.where(corpusStorage.Document.document_id.in_(sorted(set(document_ids)))) hosts: dict[AssetKey, list[AssetHost]] = {} for document_id, chunk_id, metadata, section_path, source_file_name in ( diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/explore_mount.py b/packages/shared-python/shared/services/retrieval/agent_tools/explore_mount.py index dc7a7f5ec..c4cbaa050 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/explore_mount.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/explore_mount.py @@ -10,6 +10,7 @@ from typing import Any from shared.services.retrieval.agent_tools.registry import ToolContext +from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext from shared.services.retrieval.hydration.assets import ( enrich_rows_with_retrieval_asset_url, ) @@ -71,6 +72,7 @@ async def mount_explore_hits( exclude_document_ids=[], document_scope=ctx.document_scope, exclude_sections=[], + revision_pins=CorpusRevisionContext.get_pins(), ) render_rows = [ diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/registry.py b/packages/shared-python/shared/services/retrieval/agent_tools/registry.py index 862531daf..282302f60 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/registry.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/registry.py @@ -16,7 +16,7 @@ import logging from collections.abc import Awaitable, Callable from contextlib import AbstractAsyncContextManager -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import Any import jsonschema @@ -24,6 +24,8 @@ from sqlalchemy.ext.asyncio import AsyncSession from shared.services.retrieval.document_scope import DocumentScope +from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext +from shared.services.retrieval.execution.revision_pins import capture_revision_pins from shared.services.retrieval.settings import EVIDENCE_TEXT_CHAR_BUDGET DbFactory = Callable[[], AbstractAsyncContextManager[AsyncSession]] @@ -174,7 +176,11 @@ async def dispatch( if schema_error is not None: return ToolResult(text="", error=schema_error) try: - return await spec.run(ctx, args) + pins = CorpusRevisionContext.get_pins() + if pins is None: + pins = await capture_revision_pins(ctx.db, user_id=ctx.user_id, namespace=ctx.namespace) + with CorpusRevisionContext.bind(pins): + return await spec.run(replace(ctx, document_scope=ctx.document_scope.narrow(list(pins))), args) except Exception as exc: # noqa: BLE001 - one broken call must not end the caller's loop logger.exception("corpus tool %s raised", name) return ToolResult( diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/scope.py b/packages/shared-python/shared/services/retrieval/agent_tools/scope.py index d0ebc4bf1..9bb853471 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/scope.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/scope.py @@ -20,13 +20,16 @@ from __future__ import annotations +from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext + +from shared.services.retrieval.corpus_storage import CorpusStorage + from dataclasses import dataclass from typing import Any from sqlalchemy import and_, or_, select from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import Document, DocumentSection from shared.services.retrieval.document_scope import DocumentScope from shared.services.retrieval.search.lexical_text import normalize_section_path @@ -88,6 +91,7 @@ async def resolve_scope( exact ``section_path`` on that document's current revision. Any single failure fails the whole call. """ + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) items = raw_scope if isinstance(raw_scope, list) else [] if not items: return [], "scope requires at least one {document_id, section_path?} item" @@ -104,19 +108,19 @@ async def resolve_scope( documents = ( ( await db.execute( - select(Document) - .where(Document.document_id.in_(document_ids)) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == "active") - .where(document_scope.predicate(Document.document_id)) + select(corpusStorage.Document) + .where(corpusStorage.Document.document_id.in_(document_ids)) + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.Document.namespace == namespace) + .where(corpusStorage.Document.status == "active") + .where(document_scope.predicate(corpusStorage.Document.document_id)) ) ) .scalars() .all() ) revision_by_doc = { - d.document_id: d.current_job_result_id for d in documents if d.current_job_result_id + d.document_id: revision for d in documents if (revision := CorpusRevisionContext.resolve_revision(d.document_id, d.current_job_result_id)) } targets: list[ScopeTarget] = [] @@ -132,10 +136,10 @@ async def resolve_scope( normalized = normalize_section_path(raw_path) exists = ( await db.execute( - select(DocumentSection.section_id) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - .where(DocumentSection.section_path == normalized) + select(corpusStorage.DocumentSection.section_id) + .where(corpusStorage.DocumentSection.document_id == document_id) + .where(corpusStorage.DocumentSection.job_result_id == job_result_id) + .where(corpusStorage.DocumentSection.section_path == normalized) ) ).first() if exists is None: diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/section_path_lookup.py b/packages/shared-python/shared/services/retrieval/agent_tools/section_path_lookup.py index 3df35992c..f15735181 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/section_path_lookup.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/section_path_lookup.py @@ -10,10 +10,11 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage + from sqlalchemy import or_, select from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import DocumentSection from shared.services.retrieval.search.lexical_text import normalize_section_path @@ -38,16 +39,18 @@ def format_ambiguous_section_path_error(normalized: str, matches: list[str]) -> ) -def section_path_anchor_filter(resolved_path: str): +def section_path_anchor_filter(resolved_path: str, document_id: str = ""): + corpusStorage: CorpusStorage = CorpusStorage.resolve_document(document_id) """SQL filter for one section (``mode=self`` anchor).""" - return DocumentSection.section_path == resolved_path + return corpusStorage.DocumentSection.section_path == resolved_path -def section_path_subtree_filter(resolved_path: str): +def section_path_subtree_filter(resolved_path: str, document_id: str = ""): + corpusStorage: CorpusStorage = CorpusStorage.resolve_document(document_id) """SQL filter for a section and all descendants (``mode=descendants``).""" return or_( - DocumentSection.section_path == resolved_path, - DocumentSection.section_path.like(f"{resolved_path} / %"), + corpusStorage.DocumentSection.section_path == resolved_path, + corpusStorage.DocumentSection.section_path.like(f"{resolved_path} / %"), ) @@ -59,23 +62,24 @@ async def resolve_section_path_anchor( section_path: str, ) -> tuple[str | None, str | None]: """Resolve one canonical ``section_path`` or return ``(None, error)``.""" + corpusStorage: CorpusStorage = CorpusStorage.resolve_document(document_id) normalized = normalize_section_path(section_path) base = ( - select(DocumentSection.section_path) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) + select(corpusStorage.DocumentSection.section_path) + .where(corpusStorage.DocumentSection.document_id == document_id) + .where(corpusStorage.DocumentSection.job_result_id == job_result_id) ) - exact_row = (await db.execute(base.where(DocumentSection.section_path == normalized))).first() + exact_row = (await db.execute(base.where(corpusStorage.DocumentSection.section_path == normalized))).first() if exact_row is not None and exact_row[0]: return str(exact_row[0]), None suffix_filter = or_( - DocumentSection.section_path == normalized, - DocumentSection.section_path.like(f"% / {normalized}"), + corpusStorage.DocumentSection.section_path == normalized, + corpusStorage.DocumentSection.section_path.like(f"% / {normalized}"), ) suffix_rows = ( - await db.execute(base.where(suffix_filter).order_by(DocumentSection.sort_order)) + await db.execute(base.where(suffix_filter).order_by(corpusStorage.DocumentSection.sort_order)) ).all() matches = paths_matching_section_ref( normalized, [str(row[0]) for row in suffix_rows if row and row[0]] diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py index 90a8621b2..d931b9389 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py @@ -23,11 +23,14 @@ from __future__ import annotations +from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext + +from shared.services.retrieval.corpus_storage import CorpusStorage + from typing import Any from sqlalchemy import select -from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.services.retrieval.agent_tools.asset_hosts import ( hosted_section_path, in_scope_hosts, @@ -118,24 +121,25 @@ async def _forward_search( asset_type: str, query: str, ) -> ToolResult: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(ctx.namespace) types = {asset_type} if asset_type in ASSET_CHUNK_TYPES else set(ASSET_CHUNK_TYPES) stmt = ( - select(DocumentChunk, DocumentSection.section_path, Document.source_file_name) - .select_from(DocumentChunk) - .join(Document, Document.document_id == DocumentChunk.document_id) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + select(corpusStorage.DocumentChunk, corpusStorage.DocumentSection.section_path, corpusStorage.Document.source_file_name) + .select_from(corpusStorage.DocumentChunk) + .join(corpusStorage.Document, corpusStorage.Document.document_id == corpusStorage.DocumentChunk.document_id) + .outerjoin(corpusStorage.DocumentSection, corpusStorage.DocumentSection.section_id == corpusStorage.DocumentChunk.section_id) .where( - ctx.document_scope.predicate(Document.document_id), - Document.user_id == ctx.user_id, - Document.namespace == ctx.namespace, - Document.status == "active", - Document.current_job_result_id == DocumentChunk.job_result_id, - DocumentChunk.chunk_type.in_(sorted(types)), + ctx.document_scope.predicate(corpusStorage.Document.document_id), + corpusStorage.Document.user_id == corpusStorage.resolve_owner(ctx.user_id), + corpusStorage.Document.namespace == ctx.namespace, + corpusStorage.Document.status == "active", + CorpusRevisionContext.build_revision_column(corpusStorage.Document) == corpusStorage.DocumentChunk.job_result_id, + corpusStorage.DocumentChunk.chunk_type.in_(sorted(types)), ) - .order_by(DocumentChunk.document_id, DocumentChunk.sort_order, DocumentChunk.chunk_id) + .order_by(corpusStorage.DocumentChunk.document_id, corpusStorage.DocumentChunk.sort_order, corpusStorage.DocumentChunk.chunk_id) ) if scope: - stmt = stmt.where(Document.document_id.in_(scope_document_ids(scope))) + stmt = stmt.where(corpusStorage.Document.document_id.in_(scope_document_ids(scope))) asset_rows = (await ctx.db.execute(stmt)).all() hosts = await load_asset_hosts( @@ -188,19 +192,20 @@ async def _reverse_lookup( scope: list[ScopeTarget], target_ids: list[str], ) -> ToolResult: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(ctx.namespace) hosts = await load_asset_hosts( ctx, document_ids=scope_document_ids(scope) if scope else None, asset_ids=target_ids, ) asset_stmt = ( - select(DocumentChunk.document_id, DocumentChunk.chunk_id, DocumentChunk.chunk_type) - .select_from(DocumentChunk) - .join(Document, Document.document_id == DocumentChunk.document_id) + select(corpusStorage.DocumentChunk.document_id, corpusStorage.DocumentChunk.chunk_id, corpusStorage.DocumentChunk.chunk_type) + .select_from(corpusStorage.DocumentChunk) + .join(corpusStorage.Document, corpusStorage.Document.document_id == corpusStorage.DocumentChunk.document_id) .where( - Document.current_job_result_id == DocumentChunk.job_result_id, - DocumentChunk.document_id.in_(sorted({key[0] for key in hosts})), - DocumentChunk.chunk_id.in_(target_ids), + CorpusRevisionContext.build_revision_column(corpusStorage.Document) == corpusStorage.DocumentChunk.job_result_id, + corpusStorage.DocumentChunk.document_id.in_(sorted({key[0] for key in hosts})), + corpusStorage.DocumentChunk.chunk_id.in_(target_ids), ) ) asset_type_by_key = { diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py index 87c40fe49..8789ddde9 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py @@ -20,12 +20,15 @@ from __future__ import annotations +from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext + +from shared.services.retrieval.corpus_storage import CorpusStorage + import re from typing import Any from sqlalchemy import func, or_, select -from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job_result import JobResult from shared.services.retrieval.agent_tools.asset_hosts import host_paths_for_hits from shared.services.retrieval.agent_tools.explore_mount import mount_explore_hits @@ -65,12 +68,12 @@ def _terms_from_args(args: dict[str, Any]) -> list[str]: return terms -def _indexed_term_search_text() -> Any: +def _indexed_term_search_text(namespace: str = "default") -> Any: """Haystack expression for ``idx_document_chunks_term_trgm``.""" - return func.lower(func.coalesce(DocumentChunk.term_search_text, "")) + return func.lower(func.coalesce(CorpusStorage.resolve_namespace(namespace).DocumentChunk.term_search_text, "")) -def _term_search(terms: list[str]) -> tuple[re.Pattern[str], Any]: +def _term_search(terms: list[str], namespace: str = "default") -> tuple[re.Pattern[str], Any]: """Build the Python matcher and the SQL term predicate. Every term uses ``LIKE`` on the indexed lowercased haystack. Several @@ -82,7 +85,7 @@ def _term_search(terms: list[str]) -> tuple[re.Pattern[str], Any]: ) return compiled, or_( *( - _indexed_term_search_text().like(f"%{term.lower()}%") + _indexed_term_search_text(namespace).like(f"%{term.lower()}%") for term in terms ) ) @@ -147,6 +150,7 @@ def _term_search(terms: list[str]) -> tuple[re.Pattern[str], Any]: }, ) async def grep(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(ctx.namespace) terms = _terms_from_args(args) if not terms: return ToolResult(text="", error="grep requires pattern or patterns") @@ -171,44 +175,44 @@ async def grep(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: return ToolResult(text="", error=f"grep: {scope_error}") scope_filter = scope_orm_clause( scope, - document_id_col=DocumentChunk.document_id, - section_path_col=DocumentSection.section_path, + document_id_col=corpusStorage.DocumentChunk.document_id, + section_path_col=corpusStorage.DocumentSection.section_path, ) - compiled, term_filter = _term_search(terms) + compiled, term_filter = _term_search(terms, ctx.namespace) # Match the indexed haystack first so PostgreSQL can use # idx_document_chunks_term_trgm. Section is joined here (not later) so a # scope's section-subtree condition can apply inside this same CTE. matched = ( select( - DocumentChunk.id, - DocumentChunk.chunk_id, - DocumentChunk.document_id, - DocumentChunk.job_result_id, - DocumentChunk.chunk_type, - DocumentChunk.term_search_text, - DocumentChunk.content, - DocumentChunk.file_path, - DocumentChunk.chunk_metadata, - DocumentChunk.section_id, - DocumentChunk.sort_order, - DocumentSection.section_path, + corpusStorage.DocumentChunk.id, + corpusStorage.DocumentChunk.chunk_id, + corpusStorage.DocumentChunk.document_id, + corpusStorage.DocumentChunk.job_result_id, + corpusStorage.DocumentChunk.chunk_type, + corpusStorage.DocumentChunk.term_search_text, + corpusStorage.DocumentChunk.content, + corpusStorage.DocumentChunk.file_path, + corpusStorage.DocumentChunk.chunk_metadata, + corpusStorage.DocumentChunk.section_id, + corpusStorage.DocumentChunk.sort_order, + corpusStorage.DocumentSection.section_path, ) - .select_from(DocumentChunk) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .select_from(corpusStorage.DocumentChunk) + .outerjoin(corpusStorage.DocumentSection, corpusStorage.DocumentSection.section_id == corpusStorage.DocumentChunk.section_id) .where( - DocumentChunk.term_search_text.is_not(None), + corpusStorage.DocumentChunk.term_search_text.is_not(None), term_filter, - DocumentChunk.user_id == ctx.user_id, - DocumentChunk.namespace == ctx.namespace, - ctx.document_scope.predicate(DocumentChunk.document_id), + corpusStorage.DocumentChunk.user_id == corpusStorage.resolve_owner(ctx.user_id), + corpusStorage.DocumentChunk.namespace == ctx.namespace, + ctx.document_scope.predicate(corpusStorage.DocumentChunk.document_id), ) ) if scope_filter is not None: matched = matched.where(scope_filter) if chunk_types: matched = matched.where( - func.lower(DocumentChunk.chunk_type).in_(sorted(chunk_types)) + func.lower(corpusStorage.DocumentChunk.chunk_type).in_(sorted(chunk_types)) ) matched = matched.cte("matched").prefix_with("MATERIALIZED") @@ -225,17 +229,17 @@ async def grep(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: JobResult.job_id, JobResult.document_metadata["result_raw_prefix"].as_string(), matched.c.section_path, - Document.source_file_name, + corpusStorage.Document.source_file_name, func.count().over().label("total_matches"), ) .select_from(matched) - .join(Document, Document.document_id == matched.c.document_id) - .outerjoin(JobResult, JobResult.id == Document.current_job_result_id) + .join(corpusStorage.Document, corpusStorage.Document.document_id == matched.c.document_id) + .outerjoin(JobResult, JobResult.id == CorpusRevisionContext.build_revision_column(corpusStorage.Document)) .where( - Document.user_id == ctx.user_id, - Document.namespace == ctx.namespace, - Document.status == "active", - Document.current_job_result_id == matched.c.job_result_id, + corpusStorage.Document.user_id == corpusStorage.resolve_owner(ctx.user_id), + corpusStorage.Document.namespace == ctx.namespace, + corpusStorage.Document.status == "active", + CorpusRevisionContext.build_revision_column(corpusStorage.Document) == matched.c.job_result_id, ) .order_by(matched.c.document_id, matched.c.sort_order) .limit(limit) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/list_documents.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/list_documents.py index 330ccb871..a742ee18f 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/list_documents.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/list_documents.py @@ -7,11 +7,13 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage +from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext + from typing import Any from sqlalchemy import select -from shared.models.database.document import Document, GraphNode from shared.services.retrieval.agent_tools.registry import ( ToolContext, ToolResult, @@ -36,18 +38,20 @@ }, ) async def list_documents(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(ctx.namespace) stmt = ( - select(Document, GraphNode.properties) + select(corpusStorage.Document, corpusStorage.GraphNode.properties) .outerjoin( - GraphNode, - (GraphNode.owner_document_id == Document.document_id) - & (GraphNode.node_kind == "document"), + corpusStorage.GraphNode, + (corpusStorage.GraphNode.owner_document_id == corpusStorage.Document.document_id) + & (corpusStorage.GraphNode.node_kind == "document") + & (corpusStorage.GraphNode.job_result_id == CorpusRevisionContext.build_revision_column(corpusStorage.Document)), ) - .where(Document.user_id == ctx.user_id) - .where(Document.namespace == ctx.namespace) - .where(Document.status == "active") - .where(ctx.document_scope.predicate(Document.document_id)) - .order_by(Document.source_file_name) + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(ctx.user_id)) + .where(corpusStorage.Document.namespace == ctx.namespace) + .where(corpusStorage.Document.status == "active") + .where(ctx.document_scope.predicate(corpusStorage.Document.document_id)) + .order_by(corpusStorage.Document.source_file_name) ) rows = (await ctx.db.execute(stmt)).all() @@ -57,6 +61,7 @@ async def list_documents(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: props = properties if isinstance(properties, dict) else {} entry = { "document_id": document.document_id, + "job_result_id": CorpusRevisionContext.resolve_revision(document.document_id, document.current_job_result_id), "source_file_name": document.source_file_name, "parse_track": document.parse_track, "top_keywords": props.get("top_keywords") or [], diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/neighbors.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/neighbors.py index e5498ee6e..e70ce9a5e 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/neighbors.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/neighbors.py @@ -8,11 +8,13 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage +from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext + from typing import Any from sqlalchemy import or_, select -from shared.models.database.document import Document, GraphEdge, GraphNode from shared.services.retrieval.agent_tools.registry import ( ToolContext, ToolResult, @@ -42,40 +44,49 @@ }, ) async def neighbors(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(ctx.namespace) document_id = str(args.get("document_id") or "").strip() if not document_id: return ToolResult(text="", error="neighbors requires document_id") document = ( await ctx.db.execute( - select(Document) - .where(Document.document_id == document_id) - .where(Document.user_id == ctx.user_id) - .where(Document.namespace == ctx.namespace) - .where(Document.status == "active") - .where(ctx.document_scope.predicate(Document.document_id)) + select(corpusStorage.Document) + .where(corpusStorage.Document.document_id == document_id) + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(ctx.user_id)) + .where(corpusStorage.Document.namespace == ctx.namespace) + .where(corpusStorage.Document.status == "active") + .where(ctx.document_scope.predicate(corpusStorage.Document.document_id)) ) ).scalar_one_or_none() if document is None: return ToolResult(text="", error=f"unknown document_id: {document_id}") - allowed_nodes = select(GraphNode.node_id).where( - ctx.document_scope.predicate(GraphNode.owner_document_id) + revision: str | None = CorpusRevisionContext.resolve_revision(document_id, document.current_job_result_id) + node_id: str = f"doc:{document_id}:{revision}" if corpusStorage.is_demo else f"doc:{document_id}" + allowed_nodes = select(corpusStorage.GraphNode.node_id).join( + corpusStorage.Document, + (corpusStorage.Document.document_id == corpusStorage.GraphNode.owner_document_id) + & (corpusStorage.GraphNode.job_result_id == CorpusRevisionContext.build_revision_column(corpusStorage.Document)), + ).where( + ctx.document_scope.predicate(corpusStorage.GraphNode.owner_document_id), + corpusStorage.Document.user_id == corpusStorage.resolve_owner(ctx.user_id), + corpusStorage.Document.namespace == ctx.namespace, + corpusStorage.Document.status == "active", ) - node_id = f"doc:{document_id}" edges = ( ( await ctx.db.execute( - select(GraphEdge) - .where(GraphEdge.user_id == ctx.user_id) - .where(GraphEdge.namespace == ctx.namespace) - .where(GraphEdge.edge_kind == "related") - .where(GraphEdge.source_node_id.in_(allowed_nodes)) - .where(GraphEdge.target_node_id.in_(allowed_nodes)) + select(corpusStorage.GraphEdge) + .where(corpusStorage.GraphEdge.user_id == corpusStorage.resolve_owner(ctx.user_id)) + .where(corpusStorage.GraphEdge.namespace == ctx.namespace) + .where(corpusStorage.GraphEdge.edge_kind == "related") + .where(corpusStorage.GraphEdge.source_node_id.in_(allowed_nodes)) + .where(corpusStorage.GraphEdge.target_node_id.in_(allowed_nodes)) .where( or_( - GraphEdge.source_node_id == node_id, - GraphEdge.target_node_id == node_id, + corpusStorage.GraphEdge.source_node_id == node_id, + corpusStorage.GraphEdge.target_node_id == node_id, ) ) ) @@ -93,9 +104,9 @@ async def neighbors(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: peer_nodes = ( ( await ctx.db.execute( - select(GraphNode).where( - GraphNode.node_id.in_(peer_node_ids), - ctx.document_scope.predicate(GraphNode.owner_document_id), + select(corpusStorage.GraphNode).where( + corpusStorage.GraphNode.node_id.in_(peer_node_ids), + ctx.document_scope.predicate(corpusStorage.GraphNode.owner_document_id), ) ) ) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py index 771dd261e..7df3d3459 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py @@ -20,11 +20,12 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage + from typing import Any from sqlalchemy import select -from shared.models.database.document import DocumentChunk, DocumentSection from shared.services.retrieval.agent_tools.map_render import render_map from shared.services.retrieval.agent_tools.registry import ( ToolContext, @@ -114,6 +115,7 @@ }, ) async def node_filter(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(ctx.namespace) targets, scope_error = await resolve_scope( ctx.db, user_id=ctx.user_id, @@ -161,9 +163,9 @@ async def node_filter(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: allowed_section_ids: set[str] | None = None if chunk_types: chunk_rows = await ctx.db.execute( - select(DocumentChunk.section_id, DocumentChunk.chunk_type).where( - DocumentChunk.document_id.in_(document_ids), - DocumentChunk.job_result_id.in_(list(revision_by_doc.values())), + select(corpusStorage.DocumentChunk.section_id, corpusStorage.DocumentChunk.chunk_type).where( + corpusStorage.DocumentChunk.document_id.in_(document_ids), + corpusStorage.DocumentChunk.job_result_id.in_(list(revision_by_doc.values())), ) ) allowed_section_ids = { @@ -179,10 +181,10 @@ async def node_filter(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: for target in targets: section_stmt = ( - select(DocumentSection) - .where(DocumentSection.document_id == target.document_id) - .where(DocumentSection.job_result_id == target.job_result_id) - .order_by(DocumentSection.sort_order, DocumentSection.section_id) + select(corpusStorage.DocumentSection) + .where(corpusStorage.DocumentSection.document_id == target.document_id) + .where(corpusStorage.DocumentSection.job_result_id == target.job_result_id) + .order_by(corpusStorage.DocumentSection.sort_order, corpusStorage.DocumentSection.section_id) ) doc_sections = list((await ctx.db.execute(section_stmt)).scalars().all()) if target.section_path is not None: diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py index 9cca417b5..4c1246315 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py @@ -19,11 +19,12 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage + from typing import Any from sqlalchemy import select -from shared.models.database.document import Document, DocumentSection from shared.services.retrieval.agent_tools.map_render import render_map from shared.services.retrieval.agent_tools.registry import ( ToolContext, @@ -76,6 +77,7 @@ }, ) async def outline(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(ctx.namespace) targets, scope_error = await resolve_scope( ctx.db, user_id=ctx.user_id, @@ -93,7 +95,7 @@ async def outline(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: documents = ( ( await ctx.db.execute( - select(Document).where(Document.document_id.in_(document_ids)) + select(corpusStorage.Document).where(corpusStorage.Document.document_id.in_(document_ids)) ) ) .scalars() @@ -106,10 +108,10 @@ async def outline(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: for target in targets: section_stmt = ( - select(DocumentSection) - .where(DocumentSection.document_id == target.document_id) - .where(DocumentSection.job_result_id == target.job_result_id) - .order_by(DocumentSection.sort_order, DocumentSection.section_id) + select(corpusStorage.DocumentSection) + .where(corpusStorage.DocumentSection.document_id == target.document_id) + .where(corpusStorage.DocumentSection.job_result_id == target.job_result_id) + .order_by(corpusStorage.DocumentSection.sort_order, corpusStorage.DocumentSection.section_id) ) doc_sections = list((await ctx.db.execute(section_stmt)).scalars().all()) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py index 69ce7f63c..f2b1a6805 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py @@ -2,12 +2,15 @@ from __future__ import annotations +from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext + +from shared.services.retrieval.corpus_storage import CorpusStorage + import sqlite3 from typing import Any from sqlalchemy import select -from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job_result import JobResult from shared.services.retrieval.agent_tools.registry import ( ToolContext, @@ -55,6 +58,7 @@ }, ) async def query_table(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(ctx.namespace) document_id = str(args.get("document_id") or "").strip() chunk_id = str(args.get("chunk_id") or "").strip() sql = normalize_select_sql(str(args.get("sql") or "")) @@ -69,36 +73,36 @@ async def query_table(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: document = ( await ctx.db.execute( select( - Document, JobResult.job_id, + corpusStorage.Document, JobResult.job_id, JobResult.document_metadata["result_raw_prefix"].as_string(), ) - .select_from(Document) - .outerjoin(JobResult, JobResult.id == Document.current_job_result_id) - .where(Document.document_id == document_id) - .where(Document.user_id == ctx.user_id) - .where(Document.namespace == ctx.namespace) - .where(Document.status == "active") - .where(ctx.document_scope.predicate(Document.document_id)) + .select_from(corpusStorage.Document) + .outerjoin(JobResult, JobResult.id == CorpusRevisionContext.build_revision_column(corpusStorage.Document)) + .where(corpusStorage.Document.document_id == document_id) + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(ctx.user_id)) + .where(corpusStorage.Document.namespace == ctx.namespace) + .where(corpusStorage.Document.status == "active") + .where(ctx.document_scope.predicate(corpusStorage.Document.document_id)) ) ).first() if document is None: return ToolResult(text="", error=f"unknown document_id: {document_id}") doc, job_id, result_raw_prefix = document - if not doc.current_job_result_id: + if not CorpusRevisionContext.resolve_revision(doc.document_id, doc.current_job_result_id): return ToolResult(text="", error=f"unknown document_id: {document_id}") row = ( await ctx.db.execute( - select(DocumentChunk, DocumentSection.section_path) - .select_from(DocumentChunk) + select(corpusStorage.DocumentChunk, corpusStorage.DocumentSection.section_path) + .select_from(corpusStorage.DocumentChunk) .outerjoin( - DocumentSection, - DocumentSection.section_id == DocumentChunk.section_id, + corpusStorage.DocumentSection, + corpusStorage.DocumentSection.section_id == corpusStorage.DocumentChunk.section_id, ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == doc.current_job_result_id) - .where(DocumentChunk.chunk_id == chunk_id) - .where(DocumentChunk.chunk_type == "table") + .where(corpusStorage.DocumentChunk.document_id == document_id) + .where(corpusStorage.DocumentChunk.job_result_id == CorpusRevisionContext.resolve_revision(doc.document_id, doc.current_job_result_id)) + .where(corpusStorage.DocumentChunk.chunk_id == chunk_id) + .where(corpusStorage.DocumentChunk.chunk_type == "table") ) ).first() if row is None: diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py index bd3b6632c..39c36480d 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py @@ -27,13 +27,16 @@ from __future__ import annotations +from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext + +from shared.services.retrieval.corpus_storage import CorpusStorage + import re from typing import Any from sqlalchemy import and_, or_, select from shared.models.database.document import ( - Document, DocumentChunk, DocumentSection, ) @@ -106,6 +109,7 @@ async def _resolve_same_as_markers( source_file_name_by_doc: dict[str, str], ) -> None: """Mutate ``page`` rows in place, replacing SAME-AS markers with owner text.""" + corpusStorage: CorpusStorage = CorpusStorage.resolve_document(next(iter(revision_by_doc), '')) matches_by_index: dict[int, list[tuple[str, str]]] = {} needed: set[tuple[str, str]] = set() for index, row in enumerate(rows): @@ -135,13 +139,13 @@ async def _resolve_same_as_markers( if not job_result_id: continue result = await db.execute( - select(DocumentChunk.content) - .select_from(DocumentChunk) - .join(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentSection.section_path == owner_path) - .where(DocumentChunk.chunk_type == "page") + select(corpusStorage.DocumentChunk.content) + .select_from(corpusStorage.DocumentChunk) + .join(corpusStorage.DocumentSection, corpusStorage.DocumentSection.section_id == corpusStorage.DocumentChunk.section_id) + .where(corpusStorage.DocumentChunk.document_id == document_id) + .where(corpusStorage.DocumentChunk.job_result_id == job_result_id) + .where(corpusStorage.DocumentSection.section_path == owner_path) + .where(corpusStorage.DocumentChunk.chunk_type == "page") ) content_row = result.first() owner_content[(document_id, owner_path)] = ( @@ -317,6 +321,7 @@ def _add(row: dict[str, Any]) -> None: }, ) async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(ctx.namespace) refs = args.get("refs") if not isinstance(refs, list) or not refs: return ToolResult(text="", error="read requires a non-empty refs list") @@ -333,19 +338,19 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: documents = ( ( await ctx.db.execute( - select(Document) - .where(Document.document_id.in_(document_ids)) - .where(Document.user_id == ctx.user_id) - .where(Document.namespace == ctx.namespace) - .where(Document.status == "active") - .where(ctx.document_scope.predicate(Document.document_id)) + select(corpusStorage.Document) + .where(corpusStorage.Document.document_id.in_(document_ids)) + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(ctx.user_id)) + .where(corpusStorage.Document.namespace == ctx.namespace) + .where(corpusStorage.Document.status == "active") + .where(ctx.document_scope.predicate(corpusStorage.Document.document_id)) ) ) .scalars() .all() ) revision_by_doc = { - d.document_id: d.current_job_result_id for d in documents if d.current_job_result_id + d.document_id: revision for d in documents if (revision := CorpusRevisionContext.resolve_revision(d.document_id, d.current_job_result_id)) } source_file_name_by_doc = {d.document_id: d.source_file_name or "" for d in documents} job_result_ids = sorted(set(revision_by_doc.values())) @@ -395,15 +400,15 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: if chunk_id: row = ( await ctx.db.execute( - select(DocumentChunk, DocumentSection.section_path) - .select_from(DocumentChunk) + select(corpusStorage.DocumentChunk, corpusStorage.DocumentSection.section_path) + .select_from(corpusStorage.DocumentChunk) .outerjoin( - DocumentSection, - DocumentSection.section_id == DocumentChunk.section_id, + corpusStorage.DocumentSection, + corpusStorage.DocumentSection.section_id == corpusStorage.DocumentChunk.section_id, ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.chunk_id == chunk_id) + .where(corpusStorage.DocumentChunk.document_id == document_id) + .where(corpusStorage.DocumentChunk.job_result_id == job_result_id) + .where(corpusStorage.DocumentChunk.chunk_id == chunk_id) ) ).first() if row is None: @@ -473,21 +478,21 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: path_clauses = [] for job in section_jobs: path_filter = ( - section_path_subtree_filter(job["resolved_path"]) + section_path_subtree_filter(job["resolved_path"], job["document_id"]) if mode == "descendants" - else section_path_anchor_filter(job["resolved_path"]) + else section_path_anchor_filter(job["resolved_path"], job["document_id"]) ) path_clauses.append( and_( - DocumentSection.document_id == job["document_id"], - DocumentSection.job_result_id == job["job_result_id"], + corpusStorage.DocumentSection.document_id == job["document_id"], + corpusStorage.DocumentSection.job_result_id == job["job_result_id"], path_filter, ) ) section_matches = list( ( await ctx.db.execute( - select(DocumentSection).where(or_(*path_clauses)) + select(corpusStorage.DocumentSection).where(or_(*path_clauses)) ) ) .scalars() @@ -498,14 +503,14 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: batched_chunks = list( ( await ctx.db.execute( - select(DocumentChunk).where( - DocumentChunk.document_id.in_( + select(corpusStorage.DocumentChunk).where( + corpusStorage.DocumentChunk.document_id.in_( {job["document_id"] for job in section_jobs} ), - DocumentChunk.job_result_id.in_( + corpusStorage.DocumentChunk.job_result_id.in_( {job["job_result_id"] for job in section_jobs} ), - DocumentChunk.section_id.in_(section_ids), + corpusStorage.DocumentChunk.section_id.in_(section_ids), # Body chunks only (text/page). image/table chunks share a # section_id with whichever section happens to store them # in the DB (always Root โ€” CORPUS_SCHEMA.md ยง3), which is @@ -514,7 +519,7 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: # below via hydrate_connected_target_rows. Without this # filter, reading Root would return every still-unmounted # asset in the document as spurious top-level entries. - DocumentChunk.chunk_type.in_(_BODY_CHUNK_TYPES), + corpusStorage.DocumentChunk.chunk_type.in_(_BODY_CHUNK_TYPES), ) ) ) @@ -610,6 +615,7 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: exclude_document_ids=[], document_scope=ctx.document_scope, exclude_sections=[], + revision_pins=CorpusRevisionContext.get_pins(), ) enriched_rows = await enrich_rows_with_retrieval_asset_url( [*base_rows, *connected_rows], diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/recall.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/recall.py index 56ff0c2e5..6ac6bf72a 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/recall.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/recall.py @@ -15,6 +15,8 @@ from __future__ import annotations +from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext + from typing import Any from shared.services.retrieval.agent_tools.asset_hosts import host_paths_for_hits @@ -123,6 +125,7 @@ async def recall(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: document_scope=document_scope, exclude_sections=[], chunk_types=chunk_types, + revision_pins=CorpusRevisionContext.get_pins(), section_targets=[ (target.document_id, target.section_path) for target in scope ] diff --git a/packages/shared-python/shared/services/retrieval/cache_service.py b/packages/shared-python/shared/services/retrieval/cache_service.py index ad766f5af..c26bdc181 100644 --- a/packages/shared-python/shared/services/retrieval/cache_service.py +++ b/packages/shared-python/shared/services/retrieval/cache_service.py @@ -80,6 +80,8 @@ def _cache_shape_digest( llm_vision_model: str | None = None, harness: str | None = None, include_document_ids: list[str] | None = None, + corpus_generation: int | None = None, + revision_digest: str | None = None, ) -> str: normalized_excludes = sorted(exclude_document_ids) normalized_sections = _normalize_exclude_sections(exclude_sections) @@ -97,6 +99,8 @@ def _cache_shape_digest( str(llm_text_model or ""), str(llm_vision_model or ""), str(harness or ""), + str(corpus_generation), + str(revision_digest or ""), ] ) payload = f"{query}|{top_k}|{'|'.join(normalized_excludes)}|{'|'.join(normalized_sections)}|{extra}" diff --git a/packages/shared-python/shared/services/retrieval/corpus_revision_context.py b/packages/shared-python/shared/services/retrieval/corpus_revision_context.py new file mode 100644 index 000000000..24d6af2bd --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/corpus_revision_context.py @@ -0,0 +1,46 @@ +"""Carry immutable revision pins across asynchronous harness/tool callbacks.""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar + +from sqlalchemy import case +from sqlalchemy.sql.elements import ColumnElement + +from shared.models.database.document import Document +from shared.models.database.demo_corpus import DemoDocument + +_CAPTURED_REVISIONS: ContextVar[Mapping[str, str] | None] = ContextVar( + "retrieval_corpus_revisions", default=None +) + + +class CorpusRevisionContext: + @staticmethod + @contextmanager + def bind(revisions: Mapping[str, str]) -> Iterator[None]: + token = _CAPTURED_REVISIONS.set(revisions) + try: + yield + finally: + _CAPTURED_REVISIONS.reset(token) + + @staticmethod + def get_pins() -> Mapping[str, str] | None: + return _CAPTURED_REVISIONS.get() + + @staticmethod + def resolve_revision(document_id: str, current_revision: str | None) -> str | None: + revisions: Mapping[str, str] | None = _CAPTURED_REVISIONS.get() + return current_revision if revisions is None else revisions.get(document_id) + + @staticmethod + def build_revision_column( + model: type[Document] | type[DemoDocument], + ) -> ColumnElement[str | None]: + revisions: Mapping[str, str] | None = _CAPTURED_REVISIONS.get() + if revisions is None: + return model.current_job_result_id.__clause_element__() + return case(dict(revisions), value=model.document_id, else_=None) diff --git a/packages/shared-python/shared/services/retrieval/corpus_storage.py b/packages/shared-python/shared/services/retrieval/corpus_storage.py new file mode 100644 index 000000000..e411b6e11 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/corpus_storage.py @@ -0,0 +1,159 @@ +"""Resolve the server-owned corpus target without replacing caller identity.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import Literal + +from shared.models.database import document as private, demo_corpus as demo +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace + +DEMO_NAMESPACE: str = "__knowhere_demo__" +DEMO_OWNER: str = "__knowhere_demo__" + + +@dataclass(frozen=True) +class CorpusStorage: + target: Literal["PRIVATE", "DEMO"] + + @property + def is_demo(self) -> bool: + return self.target == "DEMO" + + def resolve_owner(self, caller_id: str) -> str: + """Storage ownership only; auth, billing, telemetry keep the caller.""" + return DEMO_OWNER if self.is_demo else caller_id + + def compile_sql(self, statement: str) -> str: + """Map fixed internal table identifiers in existing lexical SQL.""" + if not self.is_demo: + return statement + return re.sub( + r"\b(" + "|".join(_TABLE_NAMES) + r")\b", + lambda match: "demo_" + match.group(0), + statement, + ) + + @property + def Document(self) -> type[private.Document] | type[demo.DemoDocument]: + return demo.DemoDocument if self.is_demo else private.Document + + @property + def DocumentSection( + self, + ) -> type[private.DocumentSection] | type[demo.DemoDocumentSection]: + return demo.DemoDocumentSection if self.is_demo else private.DocumentSection + + @property + def DocumentChunk( + self, + ) -> type[private.DocumentChunk] | type[demo.DemoDocumentChunk]: + return demo.DemoDocumentChunk if self.is_demo else private.DocumentChunk + + @property + def DocumentMapUnit( + self, + ) -> type[private.DocumentMapUnit] | type[demo.DemoDocumentMapUnit]: + return demo.DemoDocumentMapUnit if self.is_demo else private.DocumentMapUnit + + @property + def DocumentMapUnitToken( + self, + ) -> type[private.DocumentMapUnitToken] | type[demo.DemoDocumentMapUnitToken]: + return ( + demo.DemoDocumentMapUnitToken + if self.is_demo + else private.DocumentMapUnitToken + ) + + @property + def DocumentMapUnitIndex( + self, + ) -> type[private.DocumentMapUnitIndex] | type[demo.DemoDocumentMapUnitIndex]: + return ( + demo.DemoDocumentMapUnitIndex + if self.is_demo + else private.DocumentMapUnitIndex + ) + + @property + def RetrievalNamespaceGeneration( + self, + ) -> ( + type[private.RetrievalNamespaceGeneration] + | type[demo.DemoRetrievalNamespaceGeneration] + ): + return ( + demo.DemoRetrievalNamespaceGeneration + if self.is_demo + else private.RetrievalNamespaceGeneration + ) + + @property + def RetrievalServingRevisionManifest( + self, + ) -> ( + type[private.RetrievalServingRevisionManifest] + | type[demo.DemoRetrievalServingRevisionManifest] + ): + return ( + demo.DemoRetrievalServingRevisionManifest + if self.is_demo + else private.RetrievalServingRevisionManifest + ) + + @property + def RetrievalNamespaceMapSnapshot( + self, + ) -> ( + type[private.RetrievalNamespaceMapSnapshot] + | type[demo.DemoRetrievalNamespaceMapSnapshot] + ): + return ( + demo.DemoRetrievalNamespaceMapSnapshot + if self.is_demo + else private.RetrievalNamespaceMapSnapshot + ) + + @property + def GraphNode(self) -> type[private.GraphNode] | type[demo.DemoGraphNode]: + return demo.DemoGraphNode if self.is_demo else private.GraphNode + + @property + def GraphEdge(self) -> type[private.GraphEdge] | type[demo.DemoGraphEdge]: + return demo.DemoGraphEdge if self.is_demo else private.GraphEdge + + @property + def RetrievalHitStat( + self, + ) -> type[private.RetrievalHitStat] | type[demo.DemoRetrievalHitStat]: + return demo.DemoRetrievalHitStat if self.is_demo else private.RetrievalHitStat + + @classmethod + def resolve_namespace(cls, namespace: str | None) -> CorpusStorage: + return cls( + "DEMO" + if normalize_retrieval_namespace(namespace) == DEMO_NAMESPACE + else "PRIVATE" + ) + + @classmethod + def resolve_document(cls, document_id: str) -> CorpusStorage: + return cls("DEMO" if document_id.startswith("ddoc_") else "PRIVATE") + + +_TABLE_NAMES: tuple[str, ...] = ( + "retrieval_serving_revision_manifests", + "retrieval_namespace_map_snapshots", + "retrieval_namespace_generations", + "document_map_unit_indexes", + "document_map_unit_tokens", + "retrieval_hit_stats", + "document_map_units", + "document_sections", + "document_chunks", + "graph_nodes", + "graph_edges", + "documents", +) diff --git a/packages/shared-python/shared/services/retrieval/demo_authorization.py b/packages/shared-python/shared/services/retrieval/demo_authorization.py new file mode 100644 index 000000000..07081fa17 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/demo_authorization.py @@ -0,0 +1,44 @@ +"""Internal transaction authorization for shared-corpus writes and caller stats.""" + +from sqlalchemy import text +from sqlalchemy.orm import Session + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import PermissionDeniedException + + +def require_demo_maintainer(user_id: str) -> None: + if user_id not in settings.get_demo_maintainer_ids(): + raise PermissionDeniedException( + user_message="Only demo maintainers may modify the shared demo corpus.", + required_permission="demo_maintainer", + ) + + +def authorize_demo_transaction(db: Session, *, user_id: str) -> None: + require_demo_maintainer(user_id) + db.execute( + text("SELECT set_config('knowhere.demo_maintainer', :identity, true)"), + {"identity": user_id}, + ) + + +def authorize_demo_caller(db: Session, *, user_id: str) -> None: + db.execute( + text("SELECT set_config('knowhere.demo_caller', :identity, true)"), + {"identity": user_id}, + ) + + +def verify_runtime_database_role(db: Session) -> None: + hasBypass: bool = bool( + db.execute( + text( + "SELECT rolsuper OR rolbypassrls FROM pg_roles WHERE rolname = current_user" + ) + ).scalar_one() + ) + if hasBypass: + raise RuntimeError( + "API and Worker require a non-superuser database role without BYPASSRLS; use a separate migration role." + ) diff --git a/packages/shared-python/shared/services/retrieval/demo_examples.py b/packages/shared-python/shared/services/retrieval/demo_examples.py new file mode 100644 index 000000000..ceb0e60e2 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/demo_examples.py @@ -0,0 +1,91 @@ +"""Bind curated examples once per publication to unique valid content matches.""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import select +from sqlalchemy.orm import Session + +from shared.models.database.demo_corpus import ( + DemoDocument, + DemoDocumentChunk, + DemoDocumentSection, +) +from shared.models.database.job_result import JobResult + + +def bind_demo_examples(db: Session, *, document_id: str, job_result_id: str) -> None: + document: DemoDocument = db.execute( + select(DemoDocument).where(DemoDocument.document_id == document_id) + ).scalar_one() + examples: object = (document.catalog_metadata or {}).get("examples", []) + boundExamples: list[dict[str, Any]] = [] + if isinstance(examples, list) and examples: + rows = db.execute( + select(DemoDocumentChunk, DemoDocumentSection.section_path) + .outerjoin( + DemoDocumentSection, + DemoDocumentSection.section_id == DemoDocumentChunk.section_id, + ) + .where( + DemoDocumentChunk.document_id == document_id, + DemoDocumentChunk.job_result_id == job_result_id, + ) + ).all() + normalizedRows = [ + (chunk, path, " ".join(str(chunk.content or "").split()).casefold()) + for chunk, path in rows + ] + for example in examples: + if not isinstance(example, dict): + continue + citations: list[dict[str, Any]] = [] + for citation in example.get("citations", []): + needle: str = " ".join( + str(citation.get("content") or "").split() + ).casefold() + matches = [ + (chunk, path) + for chunk, path, content in normalizedRows + if needle and needle in content + ] + if len(matches) != 1: + citations = [] + break + chunk, path = matches[0] + pageNumbers: object = (chunk.chunk_metadata or {}).get("page_nums", []) + citations.append( + { + "description": str(citation.get("description") or ""), + "document_id": document_id, + "canonical_document_id": document_id, + "job_result_id": job_result_id, + "demo_source_id": document.demo_source_id, + "id": chunk.id, + "demo_chunk_id": chunk.id, + "chunk_id": chunk.chunk_id, + "section_path": path, + "content": str(citation.get("content") or ""), + "page_number": pageNumbers[0] + if isinstance(pageNumbers, list) and pageNumbers + else None, + } + ) + if citations: + boundExamples.append( + { + "id": example.get("id"), + "question": example.get("question"), + "answer": example.get("answer"), + "citations": citations, + } + ) + revision: JobResult = db.execute( + select(JobResult).where(JobResult.id == job_result_id) + ).scalar_one() + revision.document_metadata = { + **(revision.document_metadata or {}), + "bound_demo_examples": boundExamples, + } + db.flush() diff --git a/packages/shared-python/shared/services/retrieval/demo_job_scope.py b/packages/shared-python/shared/services/retrieval/demo_job_scope.py new file mode 100644 index 000000000..0ad3f4821 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/demo_job_scope.py @@ -0,0 +1,48 @@ +"""Verify durable ingestion ownership before parsing or publishing demos.""" + +from sqlalchemy import select +from sqlalchemy.orm import Session + +from shared.core.exceptions.domain_exceptions import PermissionDeniedException +from shared.models.database.demo_corpus import DemoDocument +from shared.models.database.job import Job +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace +from shared.services.retrieval.corpus_storage import CorpusStorage, DEMO_NAMESPACE +from shared.services.retrieval.demo_authorization import authorize_demo_transaction + + +def validate_job_corpus(db: Session, *, job: Job) -> CorpusStorage: + metadata: dict[str, object] = dict(job.job_metadata or {}) + storage: CorpusStorage = CorpusStorage.resolve_namespace( + str(metadata.get("namespace") or "") + ) + target: object = metadata.get("corpus_target") + documentId: str = str(metadata.get("document_id") or "") + if not storage.is_demo: + if target == "DEMO" or documentId.startswith("ddoc_"): + raise PermissionDeniedException( + user_message="Ingestion corpus target does not match the namespace." + ) + return storage + authorize_demo_transaction(db, user_id=str(job.user_id)) + if ( + target != "DEMO" + or normalize_retrieval_namespace(str(metadata.get("namespace") or "")) + != DEMO_NAMESPACE + ): + raise PermissionDeniedException( + user_message="Invalid internal demo ingestion target." + ) + document: DemoDocument | None = db.execute( + select(DemoDocument) + .where( + DemoDocument.document_id == documentId, + DemoDocument.demo_source_id == metadata.get("demo_source_id"), + ) + .with_for_update() + ).scalar_one_or_none() + if document is None or document.status == "archived": + raise PermissionDeniedException( + user_message="Demo ingestion source and document do not match an active source." + ) + return storage diff --git a/packages/shared-python/shared/services/retrieval/demo_query_scope.py b/packages/shared-python/shared/services/retrieval/demo_query_scope.py new file mode 100644 index 000000000..12a2e4ca7 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/demo_query_scope.py @@ -0,0 +1,35 @@ +"""Readiness errors are explicit; prepared subsets remain independently usable.""" + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.demo_not_ready import DemoNotReadyException +from shared.core.exceptions.domain_exceptions import NotFoundException +from shared.models.database.demo_corpus import DemoDocument +from shared.services.retrieval.execution.revision_pins import RetrievalRevisionPins + + +async def validate_demo_query_scope( + db: AsyncSession, + *, + include_document_ids: list[str] | None, + pins: RetrievalRevisionPins, +) -> None: + if include_document_ids: + documents = ( + await db.execute( + select(DemoDocument.document_id, DemoDocument.status).where( + DemoDocument.document_id.in_(include_document_ids) + ) + ) + ).all() + activeIds: set[str] = { + str(identity) for identity, status in documents if status == "active" + } + for identity in include_document_ids: + if identity not in activeIds: + raise NotFoundException(resource="Demo document", resource_id=identity) + if identity not in pins: + raise DemoNotReadyException(is_explicit_source=True) + if not pins and include_document_ids != []: + raise DemoNotReadyException(is_explicit_source=False) diff --git a/packages/shared-python/shared/services/retrieval/execution/plan.py b/packages/shared-python/shared/services/retrieval/execution/plan.py index 2b654a36f..e1c136b75 100644 --- a/packages/shared-python/shared/services/retrieval/execution/plan.py +++ b/packages/shared-python/shared/services/retrieval/execution/plan.py @@ -6,6 +6,10 @@ from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession +from hashlib import sha256 +from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext +from shared.services.retrieval.corpus_storage import CorpusStorage +from shared.services.retrieval.demo_query_scope import validate_demo_query_scope from shared.models.schemas.llm_config import LLMConfig from shared.services.ai.llm_overrides import ( @@ -136,7 +140,14 @@ async def _execute_with_overrides(self, request: RetrievalQuery) -> dict[str, An f"rerank={request.rerank} threshold={request.threshold}" ) + capturedDemoPins = None cache_extra = request.build_cache_extra() + if CorpusStorage.resolve_namespace(request.namespace).is_demo: + capturedDemoPins = await capture_revision_pins(request.db, user_id=request.user_id, namespace=request.namespace) + await validate_demo_query_scope(request.db, include_document_ids=request.include_document_ids, pins=capturedDemoPins) + cache_extra["corpus_generation"] = capturedDemoPins.generation + cache_extra["revision_digest"] = sha256(repr(sorted(capturedDemoPins.items())).encode()).hexdigest() + cache_version, cached_response = await _read_cached_response( user_id=request.user_id, namespace=request.namespace, @@ -152,7 +163,7 @@ async def _execute_with_overrides(self, request: RetrievalQuery) -> dict[str, An logger.debug(f" ๐Ÿ“ฆ Cache miss (version={cache_version}), running full pipeline") route_context = request.build_route_context() - revision_pins = await capture_revision_pins( + revision_pins = capturedDemoPins or await capture_revision_pins( request.db, user_id=request.user_id, namespace=request.namespace, @@ -168,9 +179,14 @@ async def _execute_with_overrides(self, request: RetrievalQuery) -> dict[str, An user_id=request.user_id, namespace=request.namespace, ) - outcome = await run_retrieval_route( - replace(route_context, revision_pins=revision_pins) - ) + if capturedDemoPins is not None: + await validate_demo_query_scope(request.db, include_document_ids=request.include_document_ids, pins=revision_pins) + cache_extra["corpus_generation"] = revision_pins.generation + cache_extra["revision_digest"] = sha256(repr(sorted(revision_pins.items())).encode()).hexdigest() + with CorpusRevisionContext.bind(revision_pins): + outcome = await run_retrieval_route( + replace(route_context, revision_pins=revision_pins, document_scope=route_context.document_scope.narrow(list(revision_pins))) + ) if cache_version is not None: await _write_cached_response( diff --git a/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py index 459ad2738..608a4391a 100644 --- a/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py +++ b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py @@ -120,6 +120,8 @@ def _merge_reference_projection( if row is None: row = _find_matching_row_for_ref(ref, rows) if row is not None: + if row.get("job_result_id"): + merged["job_result_id"] = row["job_result_id"] if row.get("asset_url"): merged["asset_url"] = row["asset_url"] page_nums = extract_page_nums(row) diff --git a/packages/shared-python/shared/services/retrieval/execution/revision_pins.py b/packages/shared-python/shared/services/retrieval/execution/revision_pins.py index e7e9efcc5..84f7f57c8 100644 --- a/packages/shared-python/shared/services/retrieval/execution/revision_pins.py +++ b/packages/shared-python/shared/services/retrieval/execution/revision_pins.py @@ -2,16 +2,16 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage + from collections.abc import Iterator, Mapping from dataclasses import dataclass from types import MappingProxyType -from sqlalchemy import select +from sqlalchemy import select, func from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import Document -from shared.models.database.document import RetrievalNamespaceGeneration @dataclass(frozen=True) @@ -41,25 +41,21 @@ async def capture_revision_pins( namespace: str, ) -> RetrievalRevisionPins: """Capture active document revisions in one database read transaction.""" + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) + generationQuery = select(corpusStorage.RetrievalNamespaceGeneration.generation).where(corpusStorage.RetrievalNamespaceGeneration.user_id == corpusStorage.resolve_owner(user_id), corpusStorage.RetrievalNamespaceGeneration.namespace == namespace).scalar_subquery() statement = ( - select(Document.document_id, Document.current_job_result_id) - .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.document_id) + select(corpusStorage.Document.document_id, corpusStorage.Document.current_job_result_id, func.coalesce(generationQuery, 0)) + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.Document.namespace == namespace) + .where(corpusStorage.Document.status == "active") + .where(corpusStorage.Document.current_job_result_id.is_not(None)) + .order_by(corpusStorage.Document.document_id) ) - try: - generation_result = await db.execute( - select(RetrievalNamespaceGeneration.generation) - .where(RetrievalNamespaceGeneration.user_id == user_id) - .where(RetrievalNamespaceGeneration.namespace == namespace) - ) - generation_row = generation_result.scalar_one_or_none() - except SQLAlchemyError: - await db.rollback() - generation_row = None - rows = (await db.execute(statement)).all() + if corpusStorage.is_demo: + statement = statement.join(corpusStorage.DocumentMapUnitIndex, (corpusStorage.DocumentMapUnitIndex.document_id == corpusStorage.Document.document_id) & (corpusStorage.DocumentMapUnitIndex.job_result_id == corpusStorage.Document.current_job_result_id)) + capturedRows = (await db.execute(statement)).all() + generation_row = capturedRows[0][2] if capturedRows else 0 + rows = [(identity, revision) for identity, revision, _generation in capturedRows] revisions = { str(document_id): str(job_result_id) for document_id, job_result_id in rows @@ -79,11 +75,12 @@ async def is_revision_generation_stable( pins: RetrievalRevisionPins, ) -> bool: """Return whether the namespace generation is unchanged since capture.""" + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) try: result = await db.execute( - select(RetrievalNamespaceGeneration.generation) - .where(RetrievalNamespaceGeneration.user_id == user_id) - .where(RetrievalNamespaceGeneration.namespace == namespace) + select(corpusStorage.RetrievalNamespaceGeneration.generation) + .where(corpusStorage.RetrievalNamespaceGeneration.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.RetrievalNamespaceGeneration.namespace == namespace) ) current_generation = result.scalar_one_or_none() except SQLAlchemyError: diff --git a/packages/shared-python/shared/services/retrieval/graph/query_service.py b/packages/shared-python/shared/services/retrieval/graph/query_service.py index f5d542b98..e7bdb6cc1 100644 --- a/packages/shared-python/shared/services/retrieval/graph/query_service.py +++ b/packages/shared-python/shared/services/retrieval/graph/query_service.py @@ -1,23 +1,26 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage +from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext + 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): +def _build_lexical_match_predicate(query: str, namespace: str = "default"): + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) like = f'%{query}%' return ( - DocumentChunk.content_lexical_text.ilike(like) - | DocumentChunk.path_lexical_text.ilike(like) + corpusStorage.DocumentChunk.content_lexical_text.ilike(like) + | corpusStorage.DocumentChunk.path_lexical_text.ilike(like) ) @@ -67,25 +70,26 @@ async def _find_documents_by_section( exclude_document_ids: set[str], exclude_sections: Iterable[dict[str, str]], ) -> list[str]: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) like = f'%{query}%' stmt = ( - select(DocumentSection.document_id) + select(corpusStorage.DocumentSection.document_id) .join( - Document, - (Document.document_id == DocumentSection.document_id) - & (Document.current_job_result_id == DocumentSection.job_result_id), + corpusStorage.Document, + (corpusStorage.Document.document_id == corpusStorage.DocumentSection.document_id) + & (CorpusRevisionContext.build_revision_column(corpusStorage.Document) == corpusStorage.DocumentSection.job_result_id), ) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.Document.namespace == namespace) + .where(corpusStorage.Document.status == 'active') .where( - DocumentSection.section_title.ilike(like) - | DocumentSection.section_path.ilike(like) + corpusStorage.DocumentSection.section_title.ilike(like) + | corpusStorage.DocumentSection.section_path.ilike(like) ) .distinct() ) if exclude_document_ids: - stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) + stmt = stmt.where(corpusStorage.Document.document_id.notin_(list(exclude_document_ids))) for item in exclude_sections or (): if not isinstance(item, dict): continue @@ -94,10 +98,10 @@ async def _find_documents_by_section( if excluded_document_id and excluded_path: stmt = stmt.where( ~( - (DocumentSection.document_id == excluded_document_id) + (corpusStorage.DocumentSection.document_id == excluded_document_id) & ( - (DocumentSection.section_path == excluded_path) - | DocumentSection.section_path.like(f'{excluded_path} / %') + (corpusStorage.DocumentSection.section_path == excluded_path) + | corpusStorage.DocumentSection.section_path.like(f'{excluded_path} / %') ) ) ) @@ -114,21 +118,22 @@ async def _find_documents_by_content( query: str, exclude_document_ids: set[str], ) -> list[str]: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) like = f'%{query}%' stmt = ( - select(Document.document_id) + select(corpusStorage.Document.document_id) .join( - DocumentChunk, - (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id), + corpusStorage.DocumentChunk, + (corpusStorage.DocumentChunk.document_id == corpusStorage.Document.document_id) + & (corpusStorage.DocumentChunk.job_result_id == CorpusRevisionContext.build_revision_column(corpusStorage.Document)), ) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where(DocumentChunk.content_lexical_text.ilike(like)) + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.Document.namespace == namespace) + .where(corpusStorage.Document.status == 'active') + .where(corpusStorage.DocumentChunk.content_lexical_text.ilike(like)) ) if exclude_document_ids: - stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) + stmt = stmt.where(corpusStorage.Document.document_id.notin_(list(exclude_document_ids))) result = await db.execute(stmt) seen: list[str] = [] for (document_id,) in result.all(): @@ -147,29 +152,30 @@ async def collect_candidate_chunks( top_k: int, exclude_sections: Iterable[dict[str, str]] = (), ) -> list[dict[str, Any]]: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) 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) + select(corpusStorage.Document, corpusStorage.DocumentChunk, corpusStorage.DocumentSection, JobResult) .join( - DocumentChunk, - (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id), + corpusStorage.DocumentChunk, + (corpusStorage.DocumentChunk.document_id == corpusStorage.Document.document_id) + & (corpusStorage.DocumentChunk.job_result_id == CorpusRevisionContext.build_revision_column(corpusStorage.Document)), ) .outerjoin( - DocumentSection, - DocumentSection.section_id == DocumentChunk.section_id, + corpusStorage.DocumentSection, + corpusStorage.DocumentSection.section_id == corpusStorage.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) + .join(JobResult, JobResult.id == corpusStorage.DocumentChunk.job_result_id) + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.Document.namespace == namespace) + .where(corpusStorage.Document.status == 'active') + .where(corpusStorage.Document.document_id.in_(list(entry_document_ids))) + .where(_build_lexical_match_predicate(query, namespace)) + .order_by(corpusStorage.DocumentChunk.sort_order) ) rows: list[dict[str, Any]] = [] offset = 0 diff --git a/packages/shared-python/shared/services/retrieval/graph/service.py b/packages/shared-python/shared/services/retrieval/graph/service.py index a16a82cd4..ca4d4437b 100644 --- a/packages/shared-python/shared/services/retrieval/graph/service.py +++ b/packages/shared-python/shared/services/retrieval/graph/service.py @@ -1,5 +1,7 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage + import logging from collections import defaultdict from dataclasses import dataclass @@ -9,12 +11,6 @@ from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.orm import Session -from shared.models.database.document import ( - Document, - DocumentChunk, - GraphEdge, - GraphNode, -) from shared.services.retrieval.publication_trace_stage import trace_publication_stage if TYPE_CHECKING: @@ -84,22 +80,23 @@ def publish_document_graph( top_summary: str | None = None, trace: PublicationTrace | None = None, ) -> None: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) with trace_publication_stage(trace, "graph_prepare"): document = db.execute( - select(Document).where(Document.document_id == document_id) + select(corpusStorage.Document).where(corpusStorage.Document.document_id == document_id) ).scalar_one_or_none() if document is None: return - if document.user_id != user_id or document.namespace != namespace: + if document.user_id != corpusStorage.resolve_owner(user_id) or document.namespace != namespace: raise ValueError( "Graph publication scope does not match the document owner" ) 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) + select(corpusStorage.DocumentChunk.chunk_type, corpusStorage.DocumentChunk.chunk_metadata) + .where(corpusStorage.DocumentChunk.document_id == document_id) + .where(corpusStorage.DocumentChunk.job_result_id == job_result_id) ).all() ) chunk_metadata_list = [row[1] or {} for row in chunk_meta_rows] @@ -122,13 +119,8 @@ def publish_document_graph( # โ”€โ”€ Clean up old graph data for this document โ”€โ”€ with trace_publication_stage(trace, "graph_persist"): - self.remove_document_graph( - db, - # Node IDs are global to a document. A namespace move must - # remove the node and its old edges before recreating it. - scope=None, - document_id=document_id, - ) + if not corpusStorage.is_demo: + self.remove_document_graph(db, scope=None, document_id=document_id) # Serialize typed entities as ["type:text", ...] for storage in node props # so peers can reconstruct the set without a separate schema. @@ -139,12 +131,12 @@ def publish_document_graph( ) # โ”€โ”€ Create document-level node (no section nodes โ€” aligned with KB KG) โ”€โ”€ - document_node_id = f"doc:{document_id}" + document_node_id = f"doc:{document_id}:{job_result_id}" if corpusStorage.is_demo else f"doc:{document_id}" with trace_publication_stage(trace, "graph_persist"): db.add( - GraphNode( + corpusStorage.GraphNode( node_id=document_node_id, - user_id=user_id, + user_id=corpusStorage.resolve_owner(user_id), namespace=namespace, node_kind="document", owner_document_id=document_id, @@ -174,18 +166,24 @@ def publish_document_graph( with trace_publication_stage(trace, "graph_prepare"): peer_statement = ( select( - GraphNode.node_id, - GraphNode.owner_document_id, - GraphNode.properties, + corpusStorage.GraphNode.node_id, + corpusStorage.GraphNode.owner_document_id, + corpusStorage.GraphNode.properties, ) - .where(GraphNode.user_id == user_id) - .where(GraphNode.namespace == namespace) - .where(GraphNode.node_kind == "document") - .where(GraphNode.owner_document_id != document_id) + .where(corpusStorage.GraphNode.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.GraphNode.namespace == namespace) + .where(corpusStorage.GraphNode.node_kind == "document") + .where(corpusStorage.GraphNode.owner_document_id != document_id) ) + if corpusStorage.is_demo: + peer_statement = peer_statement.join( + corpusStorage.Document, + (corpusStorage.Document.document_id == corpusStorage.GraphNode.owner_document_id) + & (corpusStorage.Document.current_job_result_id == corpusStorage.GraphNode.job_result_id), + ).where(corpusStorage.Document.status == "active") if candidate_terms: term_array = cast(candidate_terms, ARRAY(Text)) - properties_json = cast(GraphNode.properties, JSONB) + properties_json = cast(corpusStorage.GraphNode.properties, JSONB) peer_statement = peer_statement.where( or_( properties_json.op("->")("top_keywords") @@ -217,9 +215,9 @@ def publish_document_graph( edge_pair = tuple(sorted([document_id, peer_document_id])) db.add( - GraphEdge( - edge_id=f"related:{edge_pair[0]}<->{edge_pair[1]}", - user_id=user_id, + corpusStorage.GraphEdge( + edge_id=f"related:{document_node_id}<->{peer_node_id}" if corpusStorage.is_demo else f"related:{edge_pair[0]}<->{edge_pair[1]}", + user_id=corpusStorage.resolve_owner(user_id), namespace=namespace, edge_kind="related", source_node_id=document_node_id, @@ -314,25 +312,29 @@ def _build_edge_properties( def remove_document_graph( self, db: Session, *, scope: GraphScope | None, document_id: str ) -> None: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(scope.namespace) if scope is not None else CorpusStorage.resolve_document(document_id) + if corpusStorage.is_demo: + # Archived revisions remain stored; RLS hides the archived source. + return document_node_id = f"doc:{document_id}" - edge_delete = delete(GraphEdge).where( + edge_delete = delete(corpusStorage.GraphEdge).where( or_( - GraphEdge.owner_document_id == document_id, - GraphEdge.source_node_id == document_node_id, - GraphEdge.target_node_id == document_node_id, + corpusStorage.GraphEdge.owner_document_id == document_id, + corpusStorage.GraphEdge.source_node_id == document_node_id, + corpusStorage.GraphEdge.target_node_id == document_node_id, ) ) - node_delete = delete(GraphNode).where( - GraphNode.owner_document_id == document_id + node_delete = delete(corpusStorage.GraphNode).where( + corpusStorage.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, + corpusStorage.GraphEdge.user_id == corpusStorage.resolve_owner(scope.user_id), + corpusStorage.GraphEdge.namespace == scope.namespace, ) node_delete = node_delete.where( - GraphNode.user_id == scope.user_id, - GraphNode.namespace == scope.namespace, + corpusStorage.GraphNode.user_id == corpusStorage.resolve_owner(scope.user_id), + corpusStorage.GraphNode.namespace == scope.namespace, ) db.execute(edge_delete) db.execute(node_delete) diff --git a/packages/shared-python/shared/services/retrieval/hydration/assets.py b/packages/shared-python/shared/services/retrieval/hydration/assets.py index bec7ea168..d4d2f4b2e 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/assets.py +++ b/packages/shared-python/shared/services/retrieval/hydration/assets.py @@ -4,6 +4,13 @@ from typing import Any from loguru import logger +from sqlalchemy import select + +from shared.core.database import get_db_context +from shared.models.database.demo_corpus import DemoDocument, DemoDocumentMapUnitIndex +from shared.models.database.job_result import JobResult +from shared.services.retrieval.corpus_storage import CorpusStorage +from shared.services.storage.demo_asset_signer import DemoAssetSigner from shared.services.retrieval.hydration.row_utils import MEDIA_CHUNK_TYPES, normalize_chunk_type from shared.services.storage.page_pdf_crop import crop_source_pdf_pages @@ -14,6 +21,41 @@ PagePdfRequestKey = tuple[str, tuple[int, ...], str] +async def _load_demo_asset_signers(rows: list[dict[str, Any]]) -> None: + revisions: dict[str, str] = { + str(row["job_result_id"]): str(row["document_id"]) + for row in rows + if row.get("job_result_id") and CorpusStorage.resolve_document(str(row.get("document_id") or "")).is_demo + } + if not revisions: + return + # Load every needed revision in one query. An unpublished or archived + # revision never supplies a signing capability; demo RLS enforces this too. + async with get_db_context() as db: + records = (await db.execute( + select(JobResult.id, JobResult.job_id, JobResult.document_metadata) + .join(DemoDocument, DemoDocument.document_id == JobResult.demo_document_id) + .join(DemoDocumentMapUnitIndex, (DemoDocumentMapUnitIndex.document_id == DemoDocument.document_id) & (DemoDocumentMapUnitIndex.job_result_id == JobResult.id)) + .where(JobResult.id.in_(revisions), DemoDocument.status == "active") + )).all() + signers: dict[str, DemoAssetSigner] = { + str(revision): DemoAssetSigner(str(job), metadata or {}) + for revision, job, metadata in records + } + for row in rows: + if CorpusStorage.resolve_document(str(row.get("document_id") or "")).is_demo: + row["demo_asset_signer"] = signers.get(str(row.get("job_result_id") or "")) + + +def _is_demo_row(row: dict[str, Any]) -> bool: + return CorpusStorage.resolve_document(str(row.get("document_id") or "")).is_demo + + +def _sign_demo_asset(row: dict[str, Any], artifact_ref: str) -> str | None: + signer: object = row.get("demo_asset_signer") + return signer.generate_url(artifact_ref) if isinstance(signer, DemoAssetSigner) else None + + 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) @@ -94,6 +136,8 @@ def _coerce_page_nums(value: object) -> list[int]: def _resolve_page_pdf_request(row: dict[str, Any]) -> PagePdfRequestKey | None: + if _is_demo_row(row): + return None job_id = str(row.get("job_id") or "").strip() if not job_id or not _is_page_row(row): return None @@ -121,6 +165,8 @@ async def _generate_retrieval_asset_url( return None job_id, artifact_ref = request + if _is_demo_row(row): + return _sign_demo_asset(row, str(row.get("file_path") or "")) try: return get_result_storage().generate_artifact_url( **( @@ -141,6 +187,13 @@ async def _generate_page_citation_asset_url( row: dict[str, Any], log_context: str, ) -> str | None: + if _is_demo_row(row): + for asset in _iter_page_assets(_metadata_for_row(row).get("page_assets")): + if url := _sign_demo_asset(row, str(asset.get("artifact_ref") or "")): + return url + # Readers reuse the normalized PDF uploaded by the Worker. Creating + # cropped PDFs here would materialize new shared assets during reads. + return _sign_demo_asset(row, "source.pdf") if _is_page_row(row) else None direct_url = _resolve_direct_page_citation_asset_url(row) if direct_url: return direct_url @@ -182,6 +235,12 @@ async def _enrich_page_assets( job_id = str(row.get("job_id") or "").strip() for page_asset in page_assets: enriched = dict(page_asset) + if _is_demo_row(row): + enriched.pop("asset_url", None) + if url := _sign_demo_asset(row, str(enriched.get("artifact_ref") or "")): + enriched["asset_url"] = url + enriched_assets.append(enriched) + continue if not str(enriched.get("asset_url") or "").strip() and job_id: artifact_ref = _normalize_artifact_ref(enriched.get("artifact_ref")) if artifact_ref is not None: @@ -275,10 +334,13 @@ async def enrich_rows_with_retrieval_asset_url( *, log_context: str, ) -> list[dict[str, Any]]: + rows = [dict(row) for row in rows] + await _load_demo_asset_signers(rows) page_pdf_urls = await _build_page_pdf_url_lookup(rows, log_context=log_context) enriched_rows: list[dict[str, Any]] = [] for row in rows: enriched = dict(row) + enriched.pop("demo_asset_signer", None) page_assets = await _enrich_page_assets( row=row, log_context=log_context, @@ -310,6 +372,8 @@ async def build_retrieval_asset_url_map( *, log_context: str, ) -> dict[str, AssetUrlValue]: + rows = [dict(row) for row in rows] + await _load_demo_asset_signers(rows) page_pdf_urls = await _build_page_pdf_url_lookup(rows, log_context=log_context) url_map: dict[str, AssetUrlValue] = {} for row in rows: diff --git a/packages/shared-python/shared/services/retrieval/hydration/connected.py b/packages/shared-python/shared/services/retrieval/hydration/connected.py index 93dc41305..c29149740 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/connected.py +++ b/packages/shared-python/shared/services/retrieval/hydration/connected.py @@ -1,5 +1,7 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage + from shared.services.retrieval.document_scope import DocumentScope from collections.abc import Mapping @@ -8,7 +10,6 @@ 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, @@ -26,6 +27,7 @@ async def hydrate_connected_target_rows( document_scope: DocumentScope = DocumentScope(), revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: + corpusStorage: CorpusStorage = CorpusStorage.resolve_document(str(rows[0].get('document_id') or '') if rows else '') if db is None: return [] @@ -56,9 +58,9 @@ async def hydrate_connected_target_rows( revision_filters = [ and_( - DocumentChunk.document_id == document_id, - DocumentChunk.job_result_id == job_result_id, - DocumentChunk.chunk_id.in_(sorted(target_ids)), + corpusStorage.DocumentChunk.document_id == document_id, + corpusStorage.DocumentChunk.job_result_id == job_result_id, + corpusStorage.DocumentChunk.chunk_id.in_(sorted(target_ids)), ) for (document_id, job_result_id), target_ids in target_ids_by_revision.items() if target_ids @@ -72,21 +74,21 @@ async def hydrate_connected_target_rows( # relationship, loading the entire legacy job-chunk collection for # every connected revision during final hydration. select( - Document, DocumentChunk, DocumentSection, JobResult.job_id, + corpusStorage.Document, corpusStorage.DocumentChunk, corpusStorage.DocumentSection, JobResult.job_id, JobResult.document_metadata["result_raw_prefix"].as_string(), ) .join( - DocumentChunk, + corpusStorage.DocumentChunk, ( - (DocumentChunk.document_id == Document.document_id) + (corpusStorage.DocumentChunk.document_id == corpusStorage.Document.document_id) if revision_pins is None else and_( - DocumentChunk.document_id == Document.document_id, + corpusStorage.DocumentChunk.document_id == corpusStorage.Document.document_id, or_( *[ and_( - DocumentChunk.document_id == document_id, - DocumentChunk.job_result_id == job_result_id, + corpusStorage.DocumentChunk.document_id == document_id, + corpusStorage.DocumentChunk.job_result_id == job_result_id, ) for document_id, job_result_id in target_ids_by_revision ] @@ -94,11 +96,11 @@ async def hydrate_connected_target_rows( ) ), ) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(document_scope.predicate(Document.document_id)) + .outerjoin(corpusStorage.DocumentSection, corpusStorage.DocumentSection.section_id == corpusStorage.DocumentChunk.section_id) + .join(JobResult, JobResult.id == corpusStorage.DocumentChunk.job_result_id) + .where(document_scope.predicate(corpusStorage.Document.document_id)) .where(or_(*revision_filters)) - .order_by(DocumentChunk.sort_order) + .order_by(corpusStorage.DocumentChunk.sort_order) ) result = await db.execute(stmt) diff --git a/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py b/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py index 0c1e4430d..3f12020be 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py +++ b/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py @@ -57,7 +57,11 @@ def collect_evidence(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: for row in rows: composed = row.get("composed") if isinstance(composed, list): - parts.extend(composed) + revisionId: object = row.get("job_result_id") + parts.extend( + {**part, "job_result_id": revisionId} if revisionId else part + for part in composed + ) return parts diff --git a/packages/shared-python/shared/services/retrieval/hydration/path.py b/packages/shared-python/shared/services/retrieval/hydration/path.py index 5572f864b..bed5fbed5 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/path.py +++ b/packages/shared-python/shared/services/retrieval/hydration/path.py @@ -1,12 +1,14 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage +from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext + from typing import Any from loguru import logger 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.search.lexical_text import normalize_section_path from shared.services.retrieval.search.scoring import get_row_path @@ -87,24 +89,25 @@ async def _hydrate_outline_paths( namespace: str, document_id: str | None, ) -> list[dict[str, Any]]: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) outline_section_filters = [ - DocumentSection.section_path == path + corpusStorage.DocumentSection.section_path == path for path in outline_paths ] outline_stmt = ( - select(Document, DocumentSection) + select(corpusStorage.Document, corpusStorage.DocumentSection) .join( - DocumentSection, - (DocumentSection.document_id == Document.document_id) - & (DocumentSection.job_result_id == Document.current_job_result_id), + corpusStorage.DocumentSection, + (corpusStorage.DocumentSection.document_id == corpusStorage.Document.document_id) + & (corpusStorage.DocumentSection.job_result_id == CorpusRevisionContext.build_revision_column(corpusStorage.Document)), ) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.Document.namespace == namespace) + .where(corpusStorage.Document.status == 'active') .where(or_(*outline_section_filters)) ) if document_id: - outline_stmt = outline_stmt.where(Document.document_id == document_id) + outline_stmt = outline_stmt.where(corpusStorage.Document.document_id == document_id) rows: list[dict[str, Any]] = [] outline_result = await db.execute(outline_stmt) @@ -146,45 +149,46 @@ async def _hydrate_chunk_paths( namespace: str, document_id: str | None, ) -> list[dict[str, Any]]: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) section_path_filters = [] self_only_paths = {path for path in chunk_paths if mode_by_path.get(path) == 'self_only'} shallow_paths = {path for path in chunk_paths if mode_by_path.get(path) == 'shallow'} for path in chunk_paths: - section_path_filters.append(DocumentSection.section_path == path) + section_path_filters.append(corpusStorage.DocumentSection.section_path == path) if path in self_only_paths: pass # exact match only โ€” no descendants elif path in shallow_paths: # Direct children only: match "path / X" but not "path / X / Y" section_path_filters.append( and_( - DocumentSection.section_path.like(f'{path} / %'), - ~DocumentSection.section_path.like(f'{path} / % / %'), + corpusStorage.DocumentSection.section_path.like(f'{path} / %'), + ~corpusStorage.DocumentSection.section_path.like(f'{path} / % / %'), ) ) else: - section_path_filters.append(DocumentSection.section_path.like(f'{path} / %')) + section_path_filters.append(corpusStorage.DocumentSection.section_path.like(f'{path} / %')) stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) + select(corpusStorage.Document, corpusStorage.DocumentChunk, corpusStorage.DocumentSection, JobResult) .join( - DocumentChunk, - (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id), + corpusStorage.DocumentChunk, + (corpusStorage.DocumentChunk.document_id == corpusStorage.Document.document_id) + & (corpusStorage.DocumentChunk.job_result_id == CorpusRevisionContext.build_revision_column(corpusStorage.Document)), ) - .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') + .outerjoin(corpusStorage.DocumentSection, corpusStorage.DocumentSection.section_id == corpusStorage.DocumentChunk.section_id) + .join(JobResult, JobResult.id == corpusStorage.DocumentChunk.job_result_id) + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.Document.namespace == namespace) + .where(corpusStorage.Document.status == 'active') .where( or_( *section_path_filters, - DocumentChunk.source_chunk_path.in_(chunk_paths), + corpusStorage.DocumentChunk.source_chunk_path.in_(chunk_paths), ) ) ) if document_id: - stmt = stmt.where(Document.document_id == document_id) + stmt = stmt.where(corpusStorage.Document.document_id == document_id) result = await db.execute(stmt) rows: list[dict[str, Any]] = [] diff --git a/packages/shared-python/shared/services/retrieval/hydration/reference.py b/packages/shared-python/shared/services/retrieval/hydration/reference.py index 5a8999e40..823ce37bf 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/reference.py +++ b/packages/shared-python/shared/services/retrieval/hydration/reference.py @@ -1,12 +1,13 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage + from collections.abc import Mapping 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 ( ReferenceLookupKey, @@ -23,6 +24,7 @@ async def hydrate_referenced_chunk_rows( score_by_chunk_id: dict[str, float] | None = None, revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) if db is None or not refs: return [] @@ -50,17 +52,17 @@ async def hydrate_referenced_chunk_rows( return [] if revision_pins is None: - chunk_join = (DocumentChunk.document_id == Document.document_id) & ( - DocumentChunk.job_result_id == Document.current_job_result_id + chunk_join = (corpusStorage.DocumentChunk.document_id == corpusStorage.Document.document_id) & ( + corpusStorage.DocumentChunk.job_result_id == corpusStorage.Document.current_job_result_id ) else: chunk_join = and_( - DocumentChunk.document_id == Document.document_id, + corpusStorage.DocumentChunk.document_id == corpusStorage.Document.document_id, or_( *[ and_( - DocumentChunk.document_id == document_id, - DocumentChunk.job_result_id == str(revision_pins[document_id]), + corpusStorage.DocumentChunk.document_id == document_id, + corpusStorage.DocumentChunk.job_result_id == str(revision_pins[document_id]), ) for document_id in pinned_document_ids ] @@ -73,26 +75,26 @@ async def hydrate_referenced_chunk_rows( # relationship, loading the entire legacy job-chunk collection for # every referenced revision during final hydration. select( - Document, DocumentChunk, DocumentSection, JobResult.job_id, + corpusStorage.Document, corpusStorage.DocumentChunk, corpusStorage.DocumentSection, JobResult.job_id, JobResult.document_metadata["result_raw_prefix"].as_string(), ) - .join(DocumentChunk, chunk_join) + .join(corpusStorage.DocumentChunk, chunk_join) .outerjoin( - DocumentSection, DocumentSection.section_id == DocumentChunk.section_id + corpusStorage.DocumentSection, corpusStorage.DocumentSection.section_id == corpusStorage.DocumentChunk.section_id ) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) + .join(JobResult, JobResult.id == corpusStorage.DocumentChunk.job_result_id) + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.Document.namespace == namespace) .where( - Document.document_id.in_( + corpusStorage.Document.document_id.in_( document_ids if revision_pins is None else pinned_document_ids ) ) - .where(DocumentChunk.chunk_id.in_(chunk_ids)) - .order_by(DocumentChunk.sort_order) + .where(corpusStorage.DocumentChunk.chunk_id.in_(chunk_ids)) + .order_by(corpusStorage.DocumentChunk.sort_order) ) if revision_pins is None: - stmt = stmt.where(Document.status == "active") + stmt = stmt.where(corpusStorage.Document.status == "active") result = await db.execute(stmt) rows_by_key: dict[ReferenceLookupKey, dict[str, Any]] = {} diff --git a/packages/shared-python/shared/services/retrieval/hydration/row_utils.py b/packages/shared-python/shared/services/retrieval/hydration/row_utils.py index b6bd1b14d..eaff7cb70 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/row_utils.py +++ b/packages/shared-python/shared/services/retrieval/hydration/row_utils.py @@ -8,6 +8,7 @@ MEDIA_CHUNK_TYPES = {'image', 'table'} PUBLIC_RESULT_FIELDS = { + 'job_result_id', 'chunk_id', 'chunk_type', 'content', @@ -18,7 +19,7 @@ 'file_path', } PUBLIC_SOURCE_FIELDS = { - 'document_id', 'source_file_name', 'section_path', 'page_nums', + 'document_id', 'source_file_name', 'section_path', 'page_nums', 'job_result_id', } ReferenceLookupKey = tuple[str, str, str, str] diff --git a/packages/shared-python/shared/services/retrieval/map_unit_index.py b/packages/shared-python/shared/services/retrieval/map_unit_index.py index 6674c601b..dc1ed8b5b 100644 --- a/packages/shared-python/shared/services/retrieval/map_unit_index.py +++ b/packages/shared-python/shared/services/retrieval/map_unit_index.py @@ -2,6 +2,8 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage + from collections import Counter from hashlib import sha256 from typing import TYPE_CHECKING @@ -13,7 +15,6 @@ from shared.models.database.document import ( DocumentChunk, DocumentMapUnit, - DocumentMapUnitIndex, DocumentMapUnitToken, DocumentSection, ) @@ -32,7 +33,6 @@ from shared.services.retrieval.scoring.score_units import build_score_units from shared.services.retrieval.publication_models import DocumentPublicationScope from shared.services.retrieval.publication_preparation_cache import PublicationPreparationCache -from shared.services.retrieval.publication_strategy import resolve_publication_strategy from shared.services.retrieval.publication_token_copy import ( TokenCopyRow, insert_token_records_with_copy, @@ -44,7 +44,7 @@ __all__ = ["MAP_UNIT_INDEX_FORMAT_VERSION", "replace_document_map_units"] -_BULK_INSERT_BATCH_SIZE = 5_000 +_BULK_INSERT_BATCH_SIZE = 1_000 _TOKEN_ID_PREFIX_LENGTH = 12 _TOKEN_ID_SEQUENCE_WIDTH = 19 @@ -57,46 +57,47 @@ def replace_document_map_units( clear_existing: bool = True, ) -> None: """Build the derived index through the authoritative map-unit constructor.""" + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(scope.namespace) if clear_existing: with trace_publication_stage(trace, "tokens_persist"): db.execute( - delete(DocumentMapUnitToken).where( - DocumentMapUnitToken.map_unit_id.in_( - select(DocumentMapUnit.id) - .where(DocumentMapUnit.document_id == scope.document_id) - .where(DocumentMapUnit.job_result_id == scope.job_result_id) + delete(corpusStorage.DocumentMapUnitToken).where( + corpusStorage.DocumentMapUnitToken.map_unit_id.in_( + select(corpusStorage.DocumentMapUnit.id) + .where(corpusStorage.DocumentMapUnit.document_id == scope.document_id) + .where(corpusStorage.DocumentMapUnit.job_result_id == scope.job_result_id) ) ) ) with trace_publication_stage(trace, "map_units_persist"): db.execute( - delete(DocumentMapUnit) - .where(DocumentMapUnit.document_id == scope.document_id) - .where(DocumentMapUnit.job_result_id == scope.job_result_id) + delete(corpusStorage.DocumentMapUnit) + .where(corpusStorage.DocumentMapUnit.document_id == scope.document_id) + .where(corpusStorage.DocumentMapUnit.job_result_id == scope.job_result_id) ) db.execute( - delete(DocumentMapUnitIndex) - .where(DocumentMapUnitIndex.document_id == scope.document_id) - .where(DocumentMapUnitIndex.job_result_id == scope.job_result_id) + delete(corpusStorage.DocumentMapUnitIndex) + .where(corpusStorage.DocumentMapUnitIndex.document_id == scope.document_id) + .where(corpusStorage.DocumentMapUnitIndex.job_result_id == scope.job_result_id) ) with trace_publication_stage(trace, "serving_index_prepare"): section_models = list( db.scalars( - select(DocumentSection) - .where(DocumentSection.document_id == scope.document_id) - .where(DocumentSection.job_result_id == scope.job_result_id) - .order_by(DocumentSection.sort_order, DocumentSection.section_id) + select(corpusStorage.DocumentSection) + .where(corpusStorage.DocumentSection.document_id == scope.document_id) + .where(corpusStorage.DocumentSection.job_result_id == scope.job_result_id) + .order_by(corpusStorage.DocumentSection.sort_order, corpusStorage.DocumentSection.section_id) ) ) chunk_models = list( db.scalars( - select(DocumentChunk) - .where(DocumentChunk.document_id == scope.document_id) - .where(DocumentChunk.job_result_id == scope.job_result_id) + select(corpusStorage.DocumentChunk) + .where(corpusStorage.DocumentChunk.document_id == scope.document_id) + .where(corpusStorage.DocumentChunk.job_result_id == scope.job_result_id) .order_by( - DocumentChunk.sort_order, - DocumentChunk.chunk_id, - DocumentChunk.id, + corpusStorage.DocumentChunk.sort_order, + corpusStorage.DocumentChunk.chunk_id, + corpusStorage.DocumentChunk.id, ) ) ) @@ -113,8 +114,6 @@ def replace_document_map_units( persisted_count = 0 token_count = 0 map_unit_rows: list[dict[str, object]] = [] - token_rows: list[dict[str, object]] = [] - use_candidate_strategy = resolve_publication_strategy() == "candidate" token_hash_cache: dict[str, str] = {} candidate_token_rows: list[TokenCopyRow] = [] path_unit_df: Counter[str] = Counter() @@ -181,49 +180,34 @@ def replace_document_map_units( if token_hash is None: token_hash = sha256(token.encode("utf-8")).hexdigest() token_hash_cache[token] = token_hash - if use_candidate_strategy: - # Keep candidate rows as compact mutable records. Their - # IDs are assigned after locality sorting, so creating a - # dictionary here would only add memory and GC pressure. - candidate_token_rows.append( - [ - "dmut_candidate", - map_unit_id, - channel, - token, - token_hash, - frequency, - ] - ) - else: - token_rows.append( - { - "id": f"dmut_{uuid4().hex[:31]}", - "map_unit_id": map_unit_id, - "channel": channel, - "token": token, - "token_hash": token_hash, - "frequency": frequency, - } - ) + # Keep candidate rows as compact mutable records. Their + # IDs are assigned after locality sorting, so creating a + # dictionary here would only add memory and GC pressure. + candidate_token_rows.append( + [ + "dmut_candidate", + map_unit_id, + channel, + token, + token_hash, + frequency, + ] + ) token_count += len(frequencies) persisted_count += 1 if trace is not None: trace.record_count("map_units", persisted_count) trace.record_count("tokens", token_count) with trace_publication_stage(trace, "map_units_persist"): - _execute_bulk_insert(db, DocumentMapUnit, map_unit_rows) + _execute_bulk_insert(db, corpusStorage.DocumentMapUnit, map_unit_rows) with trace_publication_stage(trace, "tokens_persist"): - if use_candidate_strategy: - _prepare_candidate_token_rows( - candidate_token_rows, - job_result_id=scope.job_result_id, - ) - insert_token_records_with_copy(db, candidate_token_rows) - else: - _execute_bulk_insert(db, DocumentMapUnitToken, token_rows) + _prepare_candidate_token_rows( + candidate_token_rows, + job_result_id=scope.job_result_id, + ) + insert_token_records_with_copy(db, candidate_token_rows, corpus_storage=corpusStorage) with trace_publication_stage(trace, "statistics_prepare"): - index = DocumentMapUnitIndex( + index = corpusStorage.DocumentMapUnitIndex( id=f"dmui_{uuid4().hex}", document_id=scope.document_id, job_result_id=scope.job_result_id, diff --git a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py index e4689ff7f..d570646c1 100644 --- a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py @@ -8,6 +8,8 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage + from typing import TYPE_CHECKING, Any from sqlalchemy import select @@ -139,12 +141,13 @@ def _target_generation( Callers advance the generation after this write, in the same transaction. """ + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) generation = current_generation if generation is None: generation = db.execute( - select(RetrievalNamespaceGeneration) - .where(RetrievalNamespaceGeneration.user_id == user_id) - .where(RetrievalNamespaceGeneration.namespace == namespace) + select(corpusStorage.RetrievalNamespaceGeneration) + .where(corpusStorage.RetrievalNamespaceGeneration.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.RetrievalNamespaceGeneration.namespace == namespace) .with_for_update() ).scalar_one() return int(generation.generation) + 1 @@ -153,10 +156,11 @@ def _target_generation( def _load_snapshot_row( db: Session, *, user_id: str, namespace: str ) -> RetrievalNamespaceMapSnapshot | None: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) return db.execute( - select(RetrievalNamespaceMapSnapshot) - .where(RetrievalNamespaceMapSnapshot.user_id == user_id) - .where(RetrievalNamespaceMapSnapshot.namespace == namespace) + select(corpusStorage.RetrievalNamespaceMapSnapshot) + .where(corpusStorage.RetrievalNamespaceMapSnapshot.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.RetrievalNamespaceMapSnapshot.namespace == namespace) ).scalar_one_or_none() @@ -202,10 +206,11 @@ def _write_snapshot( format_version: int, target_generation: int, ) -> None: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) if row is None: db.add( - RetrievalNamespaceMapSnapshot( - user_id=user_id, + corpusStorage.RetrievalNamespaceMapSnapshot( + user_id=corpusStorage.resolve_owner(user_id), namespace=namespace, generation=target_generation, format_version=format_version, diff --git a/packages/shared-python/shared/services/retrieval/publication_chunk_copy.py b/packages/shared-python/shared/services/retrieval/publication_chunk_copy.py index d86218b02..0ec8019c3 100644 --- a/packages/shared-python/shared/services/retrieval/publication_chunk_copy.py +++ b/packages/shared-python/shared/services/retrieval/publication_chunk_copy.py @@ -11,7 +11,9 @@ from psycopg2.extensions import cursor as PsycopgCursor from psycopg2.extensions import get_wait_callback from psycopg2.extras import execute_values +from sqlalchemy import insert from sqlalchemy.orm import Session +from shared.models.database.demo_corpus import DemoDocumentChunk from shared.models.database.document import DocumentChunk from shared.services.jobs.lifecycle.publication_trace_sql import record_publication_sql @@ -123,8 +125,12 @@ def _copy_psycopg( cursor.close() -def _insert_chunk_batch(db: Session, chunks: list[DocumentChunk]) -> None: +def _insert_chunk_batch(db: Session, chunks: list[DocumentChunk | DemoDocumentChunk]) -> None: """Persist one bounded statement and record its actual SQL duration.""" + if chunks and isinstance(chunks[0], DemoDocumentChunk): + demoRecords = [{column: (getattr(chunk, column) if column != "created_at" else getattr(chunk, column) or datetime.now(timezone.utc).replace(tzinfo=None)) for column in _COLUMNS} for chunk in chunks] + db.execute(insert(DemoDocumentChunk.__table__), demoRecords) + return connection = db.connection() raw_connection = connection.connection driver = connection.dialect.driver @@ -176,7 +182,7 @@ def _insert_chunk_batch(db: Session, chunks: list[DocumentChunk]) -> None: ) -def insert_chunks_with_copy(db: Session, chunks: list[DocumentChunk]) -> None: +def insert_chunks_with_copy(db: Session, chunks: list[DocumentChunk | DemoDocumentChunk]) -> None: """Persist prepared chunks in bounded batches inside the caller's transaction.""" for start in range(0, len(chunks), _COPY_BATCH_SIZE): _insert_chunk_batch(db, chunks[start : start + _COPY_BATCH_SIZE]) diff --git a/packages/shared-python/shared/services/retrieval/publication_content.py b/packages/shared-python/shared/services/retrieval/publication_content.py index 106320f3a..13d60f17f 100644 --- a/packages/shared-python/shared/services/retrieval/publication_content.py +++ b/packages/shared-python/shared/services/retrieval/publication_content.py @@ -1,5 +1,7 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage + from typing import TYPE_CHECKING, Any, cast from uuid import uuid4 @@ -8,13 +10,10 @@ from shared.models.database.document import ( DocumentChunk, - DocumentMapUnit, - DocumentMapUnitIndex, DocumentSection, ) from shared.services.retrieval.map_unit_index import replace_document_map_units from shared.services.retrieval.publication_models import DocumentPublicationScope -from shared.services.retrieval.publication_strategy import resolve_publication_strategy from shared.services.retrieval.publication_trace_stage import trace_publication_stage from shared.services.retrieval.serving_manifest import persist_revision_serving_state from shared.services.retrieval.search.lexical_text import ( @@ -111,15 +110,11 @@ def replace_document_revision_content( ) ) with trace_publication_stage(trace, "chunks_persist"): - if resolve_publication_strategy() == "candidate": - from shared.services.retrieval.publication_chunk_copy import ( - insert_chunks_with_copy, - ) + from shared.services.retrieval.publication_chunk_copy import ( + insert_chunks_with_copy, + ) - insert_chunks_with_copy(db, chunk_rows) - else: - db.add_all(chunk_rows) - db.flush() + insert_chunks_with_copy(db, chunk_rows) # The revision cleanup above already removed any derived rows for this # immutable job result. Skip the constructor's defensive DELETE pass on # the hot publication path; standalone backfills retain the default. @@ -150,6 +145,7 @@ def __init__( self._sections_by_path: dict[str, DocumentSection] = {} def ensure_section(self, section_path: str) -> DocumentSection: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(self._scope.namespace) existing_section = self._sections_by_path.get(section_path) if existing_section is not None: return existing_section @@ -160,7 +156,7 @@ def ensure_section(self, section_path: str) -> DocumentSection: if ancestor_path in self._sections_by_path: continue - ancestor_section = DocumentSection( + ancestor_section = corpusStorage.DocumentSection( section_id=f"sec_{uuid4().hex[:12]}", user_id=self._scope.user_id, namespace=self._scope.namespace, @@ -197,29 +193,30 @@ def _delete_existing_revision_content( scope: DocumentPublicationScope, trace: PublicationTrace | None = None, ) -> None: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(scope.namespace) with trace_publication_stage(trace, "statistics_persist"): db.execute( - delete(DocumentMapUnitIndex) - .where(DocumentMapUnitIndex.document_id == scope.document_id) - .where(DocumentMapUnitIndex.job_result_id == scope.job_result_id) + delete(corpusStorage.DocumentMapUnitIndex) + .where(corpusStorage.DocumentMapUnitIndex.document_id == scope.document_id) + .where(corpusStorage.DocumentMapUnitIndex.job_result_id == scope.job_result_id) ) with trace_publication_stage(trace, "map_units_persist"): db.execute( - delete(DocumentMapUnit) - .where(DocumentMapUnit.document_id == scope.document_id) - .where(DocumentMapUnit.job_result_id == scope.job_result_id) + delete(corpusStorage.DocumentMapUnit) + .where(corpusStorage.DocumentMapUnit.document_id == scope.document_id) + .where(corpusStorage.DocumentMapUnit.job_result_id == scope.job_result_id) ) with trace_publication_stage(trace, "chunks_persist"): db.execute( - delete(DocumentChunk) - .where(DocumentChunk.document_id == scope.document_id) - .where(DocumentChunk.job_result_id == scope.job_result_id) + delete(corpusStorage.DocumentChunk) + .where(corpusStorage.DocumentChunk.document_id == scope.document_id) + .where(corpusStorage.DocumentChunk.job_result_id == scope.job_result_id) ) with trace_publication_stage(trace, "sections_persist"): db.execute( - delete(DocumentSection) - .where(DocumentSection.document_id == scope.document_id) - .where(DocumentSection.job_result_id == scope.job_result_id) + delete(corpusStorage.DocumentSection) + .where(corpusStorage.DocumentSection.document_id == scope.document_id) + .where(corpusStorage.DocumentSection.job_result_id == scope.job_result_id) ) @@ -232,11 +229,12 @@ def _build_document_chunk( scope: DocumentPublicationScope, fallback_sort_order: int, ) -> DocumentChunk: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(scope.namespace) 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( + return corpusStorage.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, diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index d5968deb1..13a435105 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -8,6 +8,9 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage +from shared.services.retrieval.demo_job_scope import validate_job_corpus + from datetime import datetime, timezone from typing import TYPE_CHECKING, Any from uuid import uuid4 @@ -40,7 +43,6 @@ lock_namespace_generation, ) from shared.services.retrieval.publication_trace_stage import trace_publication_stage -from shared.services.retrieval.publication_strategy import resolve_publication_strategy if TYPE_CHECKING: from shared.services.jobs.lifecycle.publication_trace import PublicationTrace @@ -66,21 +68,15 @@ def get_existing_document_scope( def _get_existing_document_scope( self, db: Session, *, job_id: str ) -> ExistingDocumentScope | None: - metadata, document_id, namespace = db.execute( - select( - Job.job_metadata, - Document.document_id, - Document.namespace, - ) - .select_from(Job) - .outerjoin( - Document, - Document.document_id == Job.job_metadata["document_id"].as_string(), - ) - .where(Job.job_id == job_id) - ).one_or_none() or (None, None, None) - if not isinstance(metadata, dict) or not document_id or namespace is None: + job: Job | None = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() + if job is None: return None + corpusStorage: CorpusStorage = validate_job_corpus(db, job=job) + metadata = job.job_metadata or {} + document = db.execute(select(corpusStorage.Document).where(corpusStorage.Document.document_id == metadata.get("document_id"), corpusStorage.Document.user_id == corpusStorage.resolve_owner(str(job.user_id)))).scalar_one_or_none() + if document is None: + return None + document_id, namespace = document.document_id, document.namespace return ExistingDocumentScope( document_id=str(document_id), @@ -125,6 +121,7 @@ def _publish_document_state_for_job( trace: PublicationTrace | None = None, ) -> PublishedDocumentState | None: + corpusStorage: CorpusStorage = validate_job_corpus(db, job=job) job_metadata = job.job_metadata or {} namespace = normalize_retrieval_namespace(job_metadata.get("namespace")) document_id = job_metadata.get("document_id") @@ -139,7 +136,7 @@ def _publish_document_state_for_job( # create another document or rewrite serving state. Locking the row # also serializes concurrent replays of the same revision. existing_document_id = db.execute( - select(JobResult.document_id) + select(JobResult.demo_document_id if corpusStorage.is_demo else JobResult.document_id) .where(JobResult.id == job_result_id) .with_for_update() ).scalar_one_or_none() @@ -156,10 +153,9 @@ def _publish_document_state_for_job( skipped_all_duplicate=True, ) - if resolve_publication_strategy() == "candidate": - # Keep bulk GIN updates inside the publication transaction while - # allowing background cleanup between publications. - db.execute(text("SET LOCAL gin_pending_list_limit = '64MB'")) + # Keep bulk GIN updates inside the publication transaction while + # allowing background cleanup between publications. + db.execute(text("SET LOCAL gin_pending_list_limit = '64MB'")) with trace_publication_stage(trace, "chunks_prepare"): deduped_chunks = deduplicate_chunks_by_source_path(chunks) @@ -196,10 +192,11 @@ def _publish_document_state_for_job( db, job_result_id=job_result_id, document_id=document.document_id, + is_demo=corpusStorage.is_demo, ) namespace = normalize_retrieval_namespace(namespace or document.namespace) scope = DocumentPublicationScope( - user_id=str(job.user_id), + user_id=corpusStorage.resolve_owner(str(job.user_id)), namespace=namespace, document_id=document.document_id, job_result_id=job_result_id, @@ -215,6 +212,10 @@ def _publish_document_state_for_job( trace=trace, ) + if corpusStorage.is_demo: + from shared.services.retrieval.demo_examples import bind_demo_examples + bind_demo_examples(db, document_id=scope.document_id, job_result_id=job_result_id) + if update_namespace_snapshot: self.update_namespace_snapshot( db, @@ -297,15 +298,16 @@ def _upsert_document_revision( document_metadata: dict[str, Any], trace: PublicationTrace | None = None, ) -> tuple[Document | None, str | None]: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) document = None previous_namespace: str | None = None if document_id: with trace_publication_stage(trace, "existing_document_lock_wait"): document = db.execute( - select(Document) + select(corpusStorage.Document) .where( - Document.document_id == document_id, - Document.user_id == str(job.user_id), + corpusStorage.Document.document_id == document_id, + corpusStorage.Document.user_id == corpusStorage.resolve_owner(str(job.user_id)), ) .with_for_update() ).scalar_one_or_none() @@ -314,9 +316,9 @@ def _upsert_document_revision( with trace_publication_stage(trace, "revision_update"): if document is None: - document = Document( + document = corpusStorage.Document( document_id=document_id or f"doc_{uuid4().hex[:12]}", - user_id=str(job.user_id), + user_id=corpusStorage.resolve_owner(str(job.user_id)), namespace=namespace, status="active", current_job_result_id=job_result_id, @@ -326,6 +328,8 @@ def _upsert_document_revision( ) db.add(document) else: + if corpusStorage.is_demo and document.status == "archived": + return None, previous_namespace if self._is_stale_document_completion( db, document=document, @@ -357,11 +361,12 @@ def _bind_job_result_document( *, job_result_id: str, document_id: str, + is_demo: bool = False, ) -> None: db.execute( update(JobResult) .where(JobResult.id == job_result_id) - .values(document_id=document_id) + .values(**({"demo_document_id": document_id} if is_demo else {"document_id": document_id})) ) def publish_document_graph( @@ -395,12 +400,13 @@ def _publish_document_graph_for_job( trace: PublicationTrace | None = None, ) -> None: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace((job.job_metadata or {}).get('namespace')) metadata = job.job_metadata or {} namespace = normalize_retrieval_namespace(metadata.get("namespace")) document_id = metadata.get("document_id") if not document_id: document = db.execute( - select(Document).where(Document.current_job_result_id == job_result_id) + select(corpusStorage.Document).where(corpusStorage.Document.current_job_result_id == job_result_id) ).scalar_one_or_none() document_id = document.document_id if document else None if not document_id: diff --git a/packages/shared-python/shared/services/retrieval/publication_strategy.py b/packages/shared-python/shared/services/retrieval/publication_strategy.py deleted file mode 100644 index c6aa6fb15..000000000 --- a/packages/shared-python/shared/services/retrieval/publication_strategy.py +++ /dev/null @@ -1,64 +0,0 @@ -"""Resolve and announce the process-wide Publication strategy. - -Worker and API processes announce the selected strategy once at startup. The -resolver deliberately accepts no override arguments: the temporary strategy is -global, and request-, namespace-, or percentage-based routing is out of scope. -""" - -from __future__ import annotations - -from loguru import logger - -from shared.core.config import settings -from shared.core.config.publication import ( - DEFAULT_PUBLICATION_STRATEGY, - PUBLICATION_STRATEGIES, - PUBLICATION_STRATEGY_ENV_VAR, - PublicationStrategy, - PublicationStrategyError, - normalize_publication_strategy, -) - -__all__ = [ - "DEFAULT_PUBLICATION_STRATEGY", - "PUBLICATION_STRATEGIES", - "PUBLICATION_STRATEGY_ENV_VAR", - "PublicationStrategy", - "PublicationStrategyAnnouncer", - "PublicationStrategyError", - "announce_publication_strategy", - "normalize_publication_strategy", - "publication_strategy_announcer", - "resolve_publication_strategy", -] - - -def resolve_publication_strategy() -> PublicationStrategy: - """Return the process-wide strategy; nothing may override it per request.""" - return normalize_publication_strategy(settings.KNOWHERE_PUBLICATION_STRATEGY) - - -class PublicationStrategyAnnouncer: - """Announce the selected strategy exactly once per process.""" - - def __init__(self) -> None: - self._has_announced: bool = False - - def announce(self, *, service_name: str) -> PublicationStrategy: - """Log the selected strategy once and return it.""" - strategy = resolve_publication_strategy() - if not self._has_announced: - self._has_announced = True - logger.info( - f"{service_name} publication strategy: {strategy} " - f"({PUBLICATION_STRATEGY_ENV_VAR})" - ) - return strategy - - -publication_strategy_announcer = PublicationStrategyAnnouncer() - - -def announce_publication_strategy(service_name: str) -> PublicationStrategy: - """Announce the strategy for a service through the process singleton.""" - return publication_strategy_announcer.announce(service_name=service_name) diff --git a/packages/shared-python/shared/services/retrieval/publication_token_copy.py b/packages/shared-python/shared/services/retrieval/publication_token_copy.py index 83aadf5fd..8243fc2a7 100644 --- a/packages/shared-python/shared/services/retrieval/publication_token_copy.py +++ b/packages/shared-python/shared/services/retrieval/publication_token_copy.py @@ -11,7 +11,9 @@ from psycopg2.extensions import cursor as PsycopgCursor from psycopg2.extensions import get_wait_callback from psycopg2.extras import execute_values +from sqlalchemy import insert from sqlalchemy.orm import Session +from shared.services.retrieval.corpus_storage import CorpusStorage from shared.services.jobs.lifecycle.publication_trace_sql import record_publication_sql @@ -125,10 +127,15 @@ def insert_token_rows_with_copy( def insert_token_records_with_copy( db: Session, rows: Sequence[TokenCopyRow], + *, corpus_storage: CorpusStorage = CorpusStorage("PRIVATE"), ) -> None: """COPY compact token rows within the publication transaction.""" if not rows: return + if corpus_storage.is_demo: + for records in _iter_batches(rows): + db.execute(insert(corpus_storage.DocumentMapUnitToken.__table__), [dict(zip(_TOKEN_COLUMNS, record, strict=True)) for record in records]) + return connection = db.connection() driver = connection.dialect.driver raw_connection = connection.connection diff --git a/packages/shared-python/shared/services/retrieval/scoring/map_lighting.py b/packages/shared-python/shared/services/retrieval/scoring/map_lighting.py index ced48f737..0ad4b85cb 100644 --- a/packages/shared-python/shared/services/retrieval/scoring/map_lighting.py +++ b/packages/shared-python/shared/services/retrieval/scoring/map_lighting.py @@ -17,17 +17,14 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage + from dataclasses import dataclass from hashlib import sha256 from sqlalchemy import select, tuple_ from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import ( - DocumentMapUnit, - DocumentMapUnitIndex, - DocumentMapUnitToken, -) from shared.services.retrieval.scoring.knowhere_hybrid import ( MAP_UNIT_INDEX_FORMAT_VERSION, PersistedScoreCorpus, @@ -67,6 +64,7 @@ async def load_leaf_unit_scores( instead of folding unscored. A section that simply has no unit of its own is absent from the mapping and scores ``0.0``. """ + corpusStorage: CorpusStorage = CorpusStorage.resolve_document(next(iter(revision_by_document), '')) query_tokens = tokenize_query_for_ranker(query) if not query_tokens or not section_ids or not revision_by_document: return None @@ -75,8 +73,8 @@ async def load_leaf_unit_scores( unit_rows = ( ( await db.execute( - select(DocumentMapUnit).where( - DocumentMapUnit.section_id.in_(section_ids) + select(corpusStorage.DocumentMapUnit).where( + corpusStorage.DocumentMapUnit.section_id.in_(section_ids) ) ) ) @@ -94,14 +92,14 @@ async def load_leaf_unit_scores( map_unit_ids = [row.id for row in unit_rows] token_rows = await db.execute( select( - DocumentMapUnitToken.map_unit_id, - DocumentMapUnitToken.channel, - DocumentMapUnitToken.token, - DocumentMapUnitToken.frequency, + corpusStorage.DocumentMapUnitToken.map_unit_id, + corpusStorage.DocumentMapUnitToken.channel, + corpusStorage.DocumentMapUnitToken.token, + corpusStorage.DocumentMapUnitToken.frequency, ).where( - DocumentMapUnitToken.map_unit_id.in_(map_unit_ids), - DocumentMapUnitToken.channel.in_(_SCORE_CHANNELS), - DocumentMapUnitToken.token_hash.in_(token_hashes), + corpusStorage.DocumentMapUnitToken.map_unit_id.in_(map_unit_ids), + corpusStorage.DocumentMapUnitToken.channel.in_(_SCORE_CHANNELS), + corpusStorage.DocumentMapUnitToken.token_hash.in_(token_hashes), ) ) frequencies: dict[tuple[str, str], dict[str, int]] = {} @@ -114,10 +112,10 @@ async def load_leaf_unit_scores( index_rows = ( ( await db.execute( - select(DocumentMapUnitIndex).where( + select(corpusStorage.DocumentMapUnitIndex).where( tuple_( - DocumentMapUnitIndex.document_id, - DocumentMapUnitIndex.job_result_id, + corpusStorage.DocumentMapUnitIndex.document_id, + corpusStorage.DocumentMapUnitIndex.job_result_id, ).in_(revision_pairs) ) ) diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py index d0ad50a2d..c3f319eaa 100644 --- a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -26,6 +26,7 @@ from loguru import logger from sqlalchemy import text +from shared.services.retrieval.corpus_storage import CorpusStorage from sqlalchemy.ext.asyncio import AsyncSession from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows @@ -239,6 +240,7 @@ async def map_unit_discovery( **_kwargs: Any, ) -> DiscoveryResult: """Score the whole in-scope corpus via the persisted map-unit BM25 scorer.""" + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) t0 = time.monotonic() if db is None: raise ValueError("database session required for map_unit_discovery") @@ -261,7 +263,7 @@ async def map_unit_discovery( signal_paths or [], filter_mode ) section_clause, section_params = _build_section_subtree_clause(section_targets) - params: dict[str, Any] = {"user_id": user_id, "namespace": namespace} + params: dict[str, Any] = {"user_id": corpusStorage.resolve_owner(user_id), "namespace": namespace} params.update(revision_params) params.update(document_scope_params) params.update(type_params) @@ -275,7 +277,9 @@ async def map_unit_discovery( section_clause, exclude_sections, exclude_document_ids, - document_scope.include is not None, + document_scope.include is not None and ( + revision_pins is None or set(document_scope.include) != set(revision_pins) + ), document_scope.exclude, ) ) @@ -308,7 +312,7 @@ async def map_unit_discovery( "token_hashes": query_token_hashes, } stage_started = time.monotonic() - unit_result = await db.execute(text(unit_statement), params) + unit_result = await db.execute(text(corpusStorage.compile_sql(unit_statement)), params) unit_rows = [dict(row._mapping) for row in unit_result.all()] unit_rows = [ row @@ -332,7 +336,7 @@ async def map_unit_discovery( type_clause=type_clause, ) frequency_query = text( - "WITH matching_tokens AS MATERIALIZED (" + corpusStorage.compile_sql("WITH matching_tokens AS MATERIALIZED (" "SELECT map_unit_id, channel, token, frequency " "FROM document_map_unit_tokens " "WHERE channel = ANY(:channels) " @@ -345,7 +349,7 @@ async def map_unit_discovery( FROM matching_tokens JOIN scoped_units ON scoped_units.map_unit_id = matching_tokens.map_unit_id - """ + """) ) stage_started = time.monotonic() frequency_result = await db.execute( @@ -417,7 +421,7 @@ async def map_unit_discovery( """ ) stage_started = time.monotonic() - index_result = await db.execute(text(index_statement), params) + index_result = await db.execute(text(corpusStorage.compile_sql(index_statement)), params) index_parts = [ ( float(path_idf or 0.0), @@ -461,7 +465,7 @@ async def map_unit_discovery( else: stage_started = time.monotonic() revision_result = await db.execute( - text(cte + "SELECT DISTINCT document_id, job_result_id FROM scoped_units"), + text(corpusStorage.compile_sql(cte + "SELECT DISTINCT document_id, job_result_id FROM scoped_units")), params, ) expected_revisions = { @@ -488,7 +492,7 @@ async def map_unit_discovery( ( await db.execute( text( - _SCOPED_UNIT_IDS_CTE.format( + corpusStorage.compile_sql(_SCOPED_UNIT_IDS_CTE.format( revision_join=revision_join, revision_clause=revision_clause, document_scope_clause=document_scope_clause, @@ -503,7 +507,7 @@ async def map_unit_discovery( JOIN scoped_units ON scoped_units.map_unit_id = tokens.map_unit_id ) AS token_count - """ + """) ), params, ) @@ -609,7 +613,7 @@ async def map_unit_discovery( if key not in {"channels", "token_hashes"} } full_unit_result = await db.execute( - text(cte + "SELECT * FROM scoped_units"), full_unit_params + text(corpusStorage.compile_sql(cte + "SELECT * FROM scoped_units")), full_unit_params ) unit_rows = [dict(row._mapping) for row in full_unit_result.all()] unit_rows = [ @@ -784,6 +788,7 @@ async def _hydrate_winning_units( revision_pins: Mapping[str, str] | None, ) -> list[dict[str, Any]]: """Map each winning leaf to its one chunk; asset requests follow connect_to.""" + corpusStorage: CorpusStorage = CorpusStorage.resolve_document(str(next(iter(rows_by_unit_id.values()), {}).get("document_id") or "")) if not ranked_unit_ids: return [] @@ -799,7 +804,7 @@ async def _hydrate_winning_units( result = await session.execute( text( - "SELECT dc.chunk_id, dc.document_id, dc.section_id, dc.chunk_type, " + corpusStorage.compile_sql("SELECT dc.chunk_id, dc.document_id, dc.section_id, dc.chunk_type, " "dc.content, dc.source_chunk_path, dc.file_path, dc.chunk_metadata, " "dc.job_result_id, dc.sort_order, ds.section_path, d.source_file_name, " "jr.job_id, " @@ -811,7 +816,7 @@ async def _hydrate_winning_units( "WHERE dc.document_id = ANY(:document_ids) " "AND dc.job_result_id = ANY(:job_result_ids) " "AND dc.section_id = ANY(:section_ids) " - "ORDER BY dc.sort_order, dc.chunk_id" + "ORDER BY dc.sort_order, dc.chunk_id") ), { "document_ids": document_ids, diff --git a/packages/shared-python/shared/services/retrieval/search/ranking.py b/packages/shared-python/shared/services/retrieval/search/ranking.py index c4906e461..eb8d25105 100644 --- a/packages/shared-python/shared/services/retrieval/search/ranking.py +++ b/packages/shared-python/shared/services/retrieval/search/ranking.py @@ -1,5 +1,7 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage + import math from collections.abc import Mapping from typing import Any @@ -8,7 +10,6 @@ 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 @@ -29,6 +30,10 @@ async def load_chunk_importance_scores( rows: list[dict[str, Any]], revision_pins: Mapping[str, str] | None = None, ) -> dict[str, float]: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) + if corpusStorage.is_demo: + from shared.services.retrieval.demo_authorization import authorize_demo_caller + await db.run_sync(lambda session: authorize_demo_caller(session, user_id=user_id)) chunk_ids = sorted( {str(row.get("chunk_id") or "").strip() for row in rows if row.get("chunk_id")} ) @@ -36,19 +41,19 @@ async def load_chunk_importance_scores( return {} stmt = ( select( - RetrievalHitStat.chunk_id, - RetrievalHitStat.hit_count, - RetrievalHitStat.last_hit_at, - RetrievalHitStat.created_at, + corpusStorage.RetrievalHitStat.chunk_id, + corpusStorage.RetrievalHitStat.hit_count, + corpusStorage.RetrievalHitStat.last_hit_at, + corpusStorage.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)) + .where(corpusStorage.RetrievalHitStat.user_id == user_id) + .where(corpusStorage.RetrievalHitStat.namespace == namespace) + .where(corpusStorage.RetrievalHitStat.hit_kind == "chunk") + .where(corpusStorage.RetrievalHitStat.chunk_id.in_(chunk_ids)) ) if revision_pins is not None: pinned_document_ids = {str(document_id) for document_id in revision_pins} - stmt = stmt.where(RetrievalHitStat.document_id.in_(pinned_document_ids)) + stmt = stmt.where(corpusStorage.RetrievalHitStat.document_id.in_(pinned_document_ids)) result = await db.execute(stmt) importance_scores: dict[str, float] = {} for chunk_id, hit_count, last_hit_at, created_at in result.all(): diff --git a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py index 69c581fb0..76271428b 100644 --- a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py +++ b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py @@ -1,5 +1,7 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage + from shared.services.retrieval.document_scope import DocumentScope from collections.abc import Mapping @@ -9,12 +11,6 @@ from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.exc import SQLAlchemyError -from shared.models.database.document import ( - Document, - DocumentChunk, - DocumentSection, - RetrievalServingRevisionManifest, -) from shared.models.database.job_result import JobResult from shared.services.retrieval.search.section_filters import is_excluded_section from shared.services.retrieval.serving_manifest import decode_serving_manifest @@ -27,18 +23,19 @@ async def count_manifest_chunks( revision_pins: Mapping[str, str], ) -> int | None: """Count chunks exactly from complete pinned manifests when available.""" + corpusStorage: CorpusStorage = CorpusStorage.resolve_document(next(iter(revision_pins), '')) if not revision_pins: return None statement = select( - RetrievalServingRevisionManifest.document_id, - RetrievalServingRevisionManifest.job_result_id, - RetrievalServingRevisionManifest.payload_zlib, - RetrievalServingRevisionManifest.checksum, - RetrievalServingRevisionManifest.format_version, + corpusStorage.RetrievalServingRevisionManifest.document_id, + corpusStorage.RetrievalServingRevisionManifest.job_result_id, + corpusStorage.RetrievalServingRevisionManifest.payload_zlib, + corpusStorage.RetrievalServingRevisionManifest.checksum, + corpusStorage.RetrievalServingRevisionManifest.format_version, ).where( tuple_( - RetrievalServingRevisionManifest.document_id, - RetrievalServingRevisionManifest.job_result_id, + corpusStorage.RetrievalServingRevisionManifest.document_id, + corpusStorage.RetrievalServingRevisionManifest.job_result_id, ).in_(list(revision_pins.items())) ) try: @@ -82,33 +79,34 @@ async def count_scoped_chunks( revision_pins: Mapping[str, str] | None = None, max_count: int | None = None, ) -> int: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) if revision_pins is None: stmt = ( - select(DocumentChunk.id) + select(corpusStorage.DocumentChunk.id) .join( - Document, - (Document.document_id == DocumentChunk.document_id) - & (Document.current_job_result_id == DocumentChunk.job_result_id), + corpusStorage.Document, + (corpusStorage.Document.document_id == corpusStorage.DocumentChunk.document_id) + & (corpusStorage.Document.current_job_result_id == corpusStorage.DocumentChunk.job_result_id), ) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.Document.namespace == namespace) + .where(corpusStorage.Document.status == 'active') ) else: stmt = ( - select(DocumentChunk.id) - .join(Document, Document.document_id == DocumentChunk.document_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) + select(corpusStorage.DocumentChunk.id) + .join(corpusStorage.Document, corpusStorage.Document.document_id == corpusStorage.DocumentChunk.document_id) + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.Document.namespace == namespace) .where( - tuple_(DocumentChunk.document_id, DocumentChunk.job_result_id).in_( + tuple_(corpusStorage.DocumentChunk.document_id, corpusStorage.DocumentChunk.job_result_id).in_( list(revision_pins.items()) ) ) ) - stmt = stmt.where(document_scope.excluding(exclude_document_ids).predicate(Document.document_id)) + stmt = stmt.where(document_scope.excluding(exclude_document_ids).predicate(corpusStorage.Document.document_id)) if allowed_chunk_types is not None: - stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types))) + stmt = stmt.where(func.lower(corpusStorage.DocumentChunk.chunk_type).in_(list(allowed_chunk_types))) if max_count is not None: stmt = stmt.limit(max_count) @@ -132,32 +130,33 @@ async def load_all_scoped_chunks( filter_mode: str, revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) if revision_pins is None: chunk_join = ( - (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id) + (corpusStorage.DocumentChunk.document_id == corpusStorage.Document.document_id) + & (corpusStorage.DocumentChunk.job_result_id == corpusStorage.Document.current_job_result_id) ) else: chunk_join = and_( - DocumentChunk.document_id == Document.document_id, - tuple_(DocumentChunk.document_id, DocumentChunk.job_result_id).in_( + corpusStorage.DocumentChunk.document_id == corpusStorage.Document.document_id, + tuple_(corpusStorage.DocumentChunk.document_id, corpusStorage.DocumentChunk.job_result_id).in_( list(revision_pins.items()) ), ) stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, chunk_join) - .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) - .order_by(DocumentChunk.sort_order) + select(corpusStorage.Document, corpusStorage.DocumentChunk, corpusStorage.DocumentSection, JobResult) + .join(corpusStorage.DocumentChunk, chunk_join) + .outerjoin(corpusStorage.DocumentSection, corpusStorage.DocumentSection.section_id == corpusStorage.DocumentChunk.section_id) + .join(JobResult, JobResult.id == corpusStorage.DocumentChunk.job_result_id) + .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.Document.namespace == namespace) + .order_by(corpusStorage.DocumentChunk.sort_order) ) if revision_pins is None: - stmt = stmt.where(Document.status == 'active') - stmt = stmt.where(document_scope.excluding(exclude_document_ids).predicate(Document.document_id)) + stmt = stmt.where(corpusStorage.Document.status == 'active') + stmt = stmt.where(document_scope.excluding(exclude_document_ids).predicate(corpusStorage.Document.document_id)) if allowed_chunk_types is not None: - stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types))) + stmt = stmt.where(func.lower(corpusStorage.DocumentChunk.chunk_type).in_(list(allowed_chunk_types))) result = await db.execute(stmt) rows: list[dict[str, Any]] = [] diff --git a/packages/shared-python/shared/services/retrieval/serving_generation.py b/packages/shared-python/shared/services/retrieval/serving_generation.py index d222f0165..283ec0a2b 100644 --- a/packages/shared-python/shared/services/retrieval/serving_generation.py +++ b/packages/shared-python/shared/services/retrieval/serving_generation.py @@ -2,6 +2,8 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage + from hashlib import sha256 from sqlalchemy import select @@ -18,9 +20,11 @@ def lock_namespace_generation( namespace: str, ) -> RetrievalNamespaceGeneration: """Create if needed, then lock and return one namespace generation row.""" + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) + user_id = corpusStorage.resolve_owner(user_id) generation_id = f"rng_{sha256(f'{user_id}:{namespace}'.encode()).hexdigest()}" db.execute( - insert(RetrievalNamespaceGeneration) + insert(corpusStorage.RetrievalNamespaceGeneration) .values( id=generation_id, user_id=user_id, @@ -29,15 +33,15 @@ def lock_namespace_generation( ) .on_conflict_do_nothing( index_elements=[ - RetrievalNamespaceGeneration.user_id, - RetrievalNamespaceGeneration.namespace, + corpusStorage.RetrievalNamespaceGeneration.user_id, + corpusStorage.RetrievalNamespaceGeneration.namespace, ] ) ) generation = db.execute( - select(RetrievalNamespaceGeneration) - .where(RetrievalNamespaceGeneration.user_id == user_id) - .where(RetrievalNamespaceGeneration.namespace == namespace) + select(corpusStorage.RetrievalNamespaceGeneration) + .where(corpusStorage.RetrievalNamespaceGeneration.user_id == corpusStorage.resolve_owner(user_id)) + .where(corpusStorage.RetrievalNamespaceGeneration.namespace == namespace) .with_for_update() ).scalar_one() return generation diff --git a/packages/shared-python/shared/services/retrieval/serving_manifest.py b/packages/shared-python/shared/services/retrieval/serving_manifest.py index f04d601c7..bf8d493c3 100644 --- a/packages/shared-python/shared/services/retrieval/serving_manifest.py +++ b/packages/shared-python/shared/services/retrieval/serving_manifest.py @@ -2,6 +2,8 @@ from __future__ import annotations +from shared.services.retrieval.corpus_storage import CorpusStorage + import hashlib import json import time @@ -11,12 +13,6 @@ from sqlalchemy import delete, select from sqlalchemy.orm import Session -from shared.models.database.document import ( - Document, - DocumentChunk, - DocumentSection, - RetrievalServingRevisionManifest, -) from shared.models.database.job_result import JobResult from shared.services.retrieval.publication_models import DocumentPublicationScope from shared.services.retrieval.publication_trace_stage import trace_publication_stage @@ -38,27 +34,28 @@ def build_revision_serving_payload( scope: DocumentPublicationScope, ) -> dict[str, Any]: """Build ordered metadata for one published document revision.""" + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(scope.namespace) document_source_file_name, job_id = db.execute( - select(Document.source_file_name, JobResult.job_id) - .select_from(Document) + select(corpusStorage.Document.source_file_name, JobResult.job_id) + .select_from(corpusStorage.Document) .join(JobResult, JobResult.id == scope.job_result_id) - .where(Document.document_id == scope.document_id) + .where(corpusStorage.Document.document_id == scope.document_id) ).one() sections = list( db.scalars( - select(DocumentSection) - .where(DocumentSection.document_id == scope.document_id) - .where(DocumentSection.job_result_id == scope.job_result_id) - .order_by(DocumentSection.sort_order, DocumentSection.section_id) + select(corpusStorage.DocumentSection) + .where(corpusStorage.DocumentSection.document_id == scope.document_id) + .where(corpusStorage.DocumentSection.job_result_id == scope.job_result_id) + .order_by(corpusStorage.DocumentSection.sort_order, corpusStorage.DocumentSection.section_id) ) ) chunks = list( db.scalars( - select(DocumentChunk) - .where(DocumentChunk.document_id == scope.document_id) - .where(DocumentChunk.job_result_id == scope.job_result_id) + select(corpusStorage.DocumentChunk) + .where(corpusStorage.DocumentChunk.document_id == scope.document_id) + .where(corpusStorage.DocumentChunk.job_result_id == scope.job_result_id) .order_by( - DocumentChunk.sort_order, DocumentChunk.chunk_id, DocumentChunk.id + corpusStorage.DocumentChunk.sort_order, corpusStorage.DocumentChunk.chunk_id, corpusStorage.DocumentChunk.id ) ) ) @@ -135,6 +132,7 @@ def persist_revision_serving_state( Returns the manifest payload so callers can patch the namespace-level MAP snapshot without rebuilding it. """ + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(scope.namespace) byte_counts: dict[str, int] | None = {} if trace is not None else None with trace_publication_stage(trace, "serving_manifest_prepare"): manifest_payload = build_revision_serving_payload(db, scope=scope) @@ -146,14 +144,14 @@ def persist_revision_serving_state( trace.record_count("manifest_uncompressed_bytes", byte_counts["uncompressed"]) with trace_publication_stage(trace, "serving_manifest_persist"): db.execute( - delete(RetrievalServingRevisionManifest) - .where(RetrievalServingRevisionManifest.document_id == scope.document_id) + delete(corpusStorage.RetrievalServingRevisionManifest) + .where(corpusStorage.RetrievalServingRevisionManifest.document_id == scope.document_id) .where( - RetrievalServingRevisionManifest.job_result_id == scope.job_result_id + corpusStorage.RetrievalServingRevisionManifest.job_result_id == scope.job_result_id ) ) db.add( - RetrievalServingRevisionManifest( + corpusStorage.RetrievalServingRevisionManifest( user_id=scope.user_id, namespace=scope.namespace, document_id=scope.document_id, diff --git a/packages/shared-python/shared/services/retrieval/stats/service.py b/packages/shared-python/shared/services/retrieval/stats/service.py index 211956022..80451231d 100644 --- a/packages/shared-python/shared/services/retrieval/stats/service.py +++ b/packages/shared-python/shared/services/retrieval/stats/service.py @@ -6,6 +6,8 @@ from uuid import uuid4 from sqlalchemy import text +from shared.services.retrieval.corpus_storage import CorpusStorage +from shared.services.retrieval.demo_authorization import authorize_demo_caller from sqlalchemy.ext.asyncio import AsyncSession _HALF_LIFE_DAYS = 30.0 @@ -67,6 +69,9 @@ async def _upsert_hit_stat( chunk_id: str | None, now: datetime, ) -> None: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) + if corpusStorage.is_demo: + await db.run_sync(lambda session: authorize_demo_caller(session, user_id=user_id)) params = { "id": f"rhs_{uuid4().hex[:12]}", "user_id": user_id, @@ -76,10 +81,10 @@ async def _upsert_hit_stat( "now": now, } if chunk_id is None: - await db.execute(_UPSERT_DOCUMENT_HIT_SQL, params) + await db.execute(text(corpusStorage.compile_sql(str(_UPSERT_DOCUMENT_HIT_SQL))), params) else: params["chunk_id"] = chunk_id - await db.execute(_UPSERT_CHUNK_HIT_SQL, params) + await db.execute(text(corpusStorage.compile_sql(str(_UPSERT_CHUNK_HIT_SQL))), params) async def record_retrieval_hits( diff --git a/packages/shared-python/shared/services/storage/adapters/s3_adapter.py b/packages/shared-python/shared/services/storage/adapters/s3_adapter.py index 45b352d62..09c2b51f3 100644 --- a/packages/shared-python/shared/services/storage/adapters/s3_adapter.py +++ b/packages/shared-python/shared/services/storage/adapters/s3_adapter.py @@ -195,6 +195,11 @@ def generate_presigned_url( if content_type: params["ContentType"] = content_type + if method.upper() == "GET" and headers: + if contentType := headers.get("Content-Type"): + params["ResponseContentType"] = contentType + if disposition := headers.get("Content-Disposition"): + params["ResponseContentDisposition"] = disposition if method.upper() == "PUT": url = self.s3_client.generate_presigned_url( "put_object", Params=params, ExpiresIn=expiration diff --git a/packages/shared-python/shared/services/storage/demo_asset_signer.py b/packages/shared-python/shared/services/storage/demo_asset_signer.py new file mode 100644 index 000000000..8f06341a1 --- /dev/null +++ b/packages/shared-python/shared/services/storage/demo_asset_signer.py @@ -0,0 +1,54 @@ +"""Sign only retained, published demo assets without creating derivative files.""" + +from __future__ import annotations + +import mimetypes +from dataclasses import dataclass +from typing import Mapping +from urllib.parse import quote + +from shared.services.storage.job_file_storage import JobFileStorage +from shared.services.storage.result_storage import get_result_storage + + +@dataclass(frozen=True) +class DemoAssetSigner: + job_id: str + revision_metadata: Mapping[str, object] + + def generate_url(self, artifact_ref: str) -> str | None: + storage = get_result_storage() + normalized: str | None = ( + "source.pdf" + if artifact_ref == "source.pdf" + else storage.normalize_artifact_ref(artifact_ref) + ) + allowed_assets: object = self.revision_metadata.get("demo_asset_manifest") + if ( + normalized != artifact_ref + or not isinstance(allowed_assets, list) + or artifact_ref not in allowed_assets + ): + return None + raw_prefix: object = self.revision_metadata.get("result_raw_prefix") + key: str = storage.build_raw_key( + job_id=self.job_id, + relative_path=artifact_ref, + raw_prefix=str(raw_prefix) if raw_prefix else None, + ) + file_name: str = artifact_ref.rsplit("/", 1)[-1] + mime_type: str = ( + mimetypes.guess_type(file_name)[0] or "application/octet-stream" + ) + files: JobFileStorage = JobFileStorage() + return files.storage_adapter.generate_presigned_url( + key, + expiration=3600, + bucket=files.results_bucket, + method="GET", + headers={ + "Content-Type": mime_type, + "Content-Disposition": "inline; filename*=UTF-8''" + + quote(file_name, safe=""), + }, + ) diff --git a/packages/shared-python/shared/services/storage/result_storage.py b/packages/shared-python/shared/services/storage/result_storage.py index 8583e19ee..4a57bfe7c 100644 --- a/packages/shared-python/shared/services/storage/result_storage.py +++ b/packages/shared-python/shared/services/storage/result_storage.py @@ -41,6 +41,8 @@ def upload( ) -> UploadedResultBundle: raise NotImplementedError + def build_raw_key(self, *, job_id: str, relative_path: str, raw_prefix: str | None = None) -> str: ... + def generate_artifact_url( self, *, diff --git a/packages/shared-python/shared/testing/contract_runtime.py b/packages/shared-python/shared/testing/contract_runtime.py index f5fbbd5f4..fca0df1b5 100644 --- a/packages/shared-python/shared/testing/contract_runtime.py +++ b/packages/shared-python/shared/testing/contract_runtime.py @@ -377,6 +377,10 @@ def clear_application_modules() -> None: sys.modules.pop(module_name, None) continue + if module_name == "shared.services.retrieval" or module_name.startswith("shared.services.retrieval."): + sys.modules.pop(module_name, None) + continue + if module_name == "shared.services.redis" or module_name.startswith( "shared.services.redis." ): @@ -820,3 +824,22 @@ def reset_contract_object_storage() -> None: parents=True, exist_ok=True, ) + + +async def configure_runtime_database_role(monkeypatch: MonkeyPatch, adminUrl: str) -> None: + engine = create_async_engine(adminUrl) + try: + async with engine.begin() as connection: + await connection.execute(text("DO $$ BEGIN IF NOT EXISTS(SELECT FROM pg_roles WHERE rolname = 'contract_runtime') THEN CREATE ROLE contract_runtime LOGIN NOSUPERUSER NOBYPASSRLS; END IF; END $$")) + await connection.execute(text("GRANT USAGE, CREATE ON SCHEMA public TO contract_runtime")) + await connection.execute(text("GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public TO contract_runtime")) + await connection.execute(text("GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO contract_runtime")) + await connection.execute(text('ALTER TABLE "user" OWNER TO contract_runtime')) + # Reset destroys test content; reseed only the immutable directory. + import json + entries = json.loads((_API_ROOT / "alembic/demo-corpus-catalog.json").read_text()) + for entry in entries: + await connection.execute(text("INSERT INTO demo_documents (document_id, demo_source_id, title, category_id, catalog_metadata, user_id, namespace, status, source_file_name, parse_track, created_at, updated_at) VALUES (:document_id, :demo_source_id, :title, :category_id, CAST(:catalog_metadata AS json), '__knowhere_demo__', '__knowhere_demo__', 'active', :file_name, 'chunk', now(), now())"), {**entry, "catalog_metadata": json.dumps(entry["catalog_metadata"])}) + finally: + await engine.dispose() + monkeypatch.setenv("DATABASE_URL", make_url(adminUrl).set(username="contract_runtime", password=None).render_as_string(hide_password=False)) diff --git a/scripts/backup-demo-originals.py b/scripts/backup-demo-originals.py new file mode 100644 index 000000000..77ab49b22 --- /dev/null +++ b/scripts/backup-demo-originals.py @@ -0,0 +1,145 @@ +"""Back up authoritative S3 originals and fail closed on missing/invalid files.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from datetime import datetime, timezone +from pathlib import Path +from typing import Any +from zipfile import ZipFile + +import boto3 +import fitz +from docx import Document + + +def validate_original(path: Path) -> dict[str, object]: + if path.suffix.lower() == ".pdf": + with fitz.open(path) as document: + if document.is_encrypted or document.page_count < 1: + raise ValueError( + "PDF cannot be opened without a password or has no pages" + ) + for page in document: + page.get_text() + return {"format": "pdf", "pages": document.page_count, "is_valid": True} + if path.suffix.lower() == ".docx": + with ZipFile(path) as archive: + if ( + archive.testzip() is not None + or "word/document.xml" not in archive.namelist() + ): + raise ValueError("Invalid DOCX archive") + document = Document(str(path)) + return { + "format": "docx", + "paragraphs": len(document.paragraphs), + "tables": len(document.tables), + "is_valid": True, + } + raise ValueError("Original format is not supported by the backup validator") + + +def hash_original(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as original: + for block in iter(lambda: original.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def back_up_originals( + *, directory: Path, sources: list[dict[str, Any]], profile: str, bucket: str +) -> None: + directory.mkdir(parents=True, exist_ok=True) + client = boto3.Session(profile_name=profile).client("s3", region_name="us-west-1") + inventoryPath: Path = directory / "inventory.json" + priorInventory: list[dict[str, Any]] = ( + json.loads(inventoryPath.read_text()) if inventoryPath.exists() else [] + ) + inventory: list[dict[str, Any]] = [] + for source in sources: + sourceId: str = str(source["source_id"]) + fileName: str = str(source["file_name"]) + if ( + Path(fileName).name != fileName + or "/" in sourceId + or "\\" in sourceId + or sourceId in (".", "..") + ): + raise ValueError("Unsafe backup source path") + path: Path = directory / sourceId / fileName + path.parent.mkdir(parents=True, exist_ok=True) + entry: dict[str, Any] = { + **source, + "relative_path": path.relative_to(directory).as_posix(), + } + try: + head = client.head_object(Bucket=bucket, Key=source["s3_key"]) + size: int = int(head["ContentLength"]) + etag: str = str(head["ETag"]) + prior = next( + (item for item in priorInventory if item["source_id"] == sourceId), {} + ) + if ( + not path.exists() + or path.stat().st_size != size + or prior.get("etag") != etag + ): + temporary: Path = path.with_name(path.name + ".partial") + client.download_file(bucket, source["s3_key"], str(temporary)) + if temporary.stat().st_size != size: + raise ValueError("Downloaded original size does not match S3 HEAD") + temporary.replace(path) + entry.update( + { + "size_bytes": size, + "sha256": hash_original(path), + "etag": etag, + "validation": validate_original(path), + "status": "verified", + } + ) + print(f"{sourceId}: verified {size} bytes", flush=True) + except Exception as error: + entry.update({"status": "failed", "error_type": type(error).__name__}) + print(f"{sourceId}: failed ({type(error).__name__})", flush=True) + inventory.append(entry) + temporaryInventory = inventoryPath.with_suffix(".partial") + temporaryInventory.write_text(json.dumps(inventory, indent=2) + "\n") + temporaryInventory.replace(inventoryPath) + if len(inventory) != len(sources) or any( + item["status"] != "verified" for item in inventory + ): + raise RuntimeError("Backup is incomplete; do not switch demo publication") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--directory", + type=Path, + default=Path.home() + / "knowhere-demo-originals" + / datetime.now(timezone.utc).strftime("%Y%m%d"), + ) + parser.add_argument( + "--sources", + type=Path, + default=Path(__file__).with_name("demo-original-sources.json"), + ) + parser.add_argument("--profile", default="knowhere") + parser.add_argument("--bucket", default="knowhere-storage-staging") + arguments = parser.parse_args() + back_up_originals( + directory=arguments.directory, + sources=json.loads(arguments.sources.read_text()), + profile=arguments.profile, + bucket=arguments.bucket, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/demo-original-sources.json b/scripts/demo-original-sources.json new file mode 100644 index 000000000..cb9c4b950 --- /dev/null +++ b/scripts/demo-original-sources.json @@ -0,0 +1,191 @@ +[ + { + "source_id": "demo-tsla-q4-2025", + "original_job_id": "job_424aa26c4f99", + "s3_key": "uploads/job_424aa26c4f99.pdf", + "file_name": "TSLA-Q4-2025-Update.pdf", + "asset_directory": "tsla-q4-2025" + }, + { + "source_id": "demo-spacex-s1", + "original_job_id": "job_d69e5bcaa3ad", + "s3_key": "uploads/job_d69e5bcaa3ad.pdf", + "file_name": "spacex-s1.pdf", + "asset_directory": "spacex-s1" + }, + { + "source_id": "demo-financial-nvda-q1-fy27-earnings-call", + "original_job_id": "job_af4ebba4cd03", + "s3_key": "uploads/job_af4ebba4cd03.pdf", + "file_name": "NVDA Q1 FY27 Earnings Call Transcript.pdf", + "asset_directory": "financial-nvda-q1-fy27-earnings-call" + }, + { + "source_id": "demo-financial-nvda-q4-fy26-earnings-call", + "original_job_id": "job_1c90dc114eec", + "s3_key": "uploads/job_1c90dc114eec.pdf", + "file_name": "NVDA Q4 FY26 Earnings Call Transcript.pdf", + "asset_directory": "financial-nvda-q4-fy26-earnings-call" + }, + { + "source_id": "demo-financial-nvda-q1-fy27-cfo-commentary", + "original_job_id": "job_fd2e1fbf5bd3", + "s3_key": "uploads/job_fd2e1fbf5bd3.pdf", + "file_name": "NVIDIA Q1 FY27 CFO Commentary.pdf", + "asset_directory": "financial-nvda-q1-fy27-cfo-commentary" + }, + { + "source_id": "demo-financial-nvda-q1-fy27-results", + "original_job_id": "job_2fa89f92637f", + "s3_key": "uploads/job_2fa89f92637f.pdf", + "file_name": "NVIDIA Q1 FY27 Results.pdf", + "asset_directory": "financial-nvda-q1-fy27-results" + }, + { + "source_id": "demo-financial-nvda-q1-fy27-presentation", + "original_job_id": "job_983e1b5fd4c4", + "s3_key": "uploads/job_983e1b5fd4c4.pdf", + "file_name": "NVIDIA Q1 FY27 Quarterly Presentation.pdf", + "asset_directory": "financial-nvda-q1-fy27-presentation" + }, + { + "source_id": "demo-financial-nvda-fy26-annual-report", + "original_job_id": "job_7cef64d9b61e", + "s3_key": "uploads/job_7cef64d9b61e.pdf", + "file_name": "NVIDIA FY26 Annual Report.pdf", + "asset_directory": "financial-nvda-fy26-annual-report" + }, + { + "source_id": "demo-financial-micron-report-530bd7ed", + "original_job_id": "job_d46a260e0c32", + "s3_key": "uploads/job_d46a260e0c32.pdf", + "file_name": "Micron investor report 530bd7ed.pdf", + "asset_directory": "financial-micron-report-530bd7ed" + }, + { + "source_id": "demo-financial-micron-report-9c0becf5", + "original_job_id": "job_b3f1ce4985c8", + "s3_key": "uploads/job_b3f1ce4985c8.pdf", + "file_name": "Micron investor report 9c0becf5.pdf", + "asset_directory": "financial-micron-report-9c0becf5" + }, + { + "source_id": "demo-financial-goog-10-k-2025", + "original_job_id": "job_57cae6fdb820", + "s3_key": "uploads/job_57cae6fdb820.pdf", + "file_name": "GOOG 10-K 2025.pdf", + "asset_directory": "financial-goog-10-k-2025" + }, + { + "source_id": "demo-financial-meta-q1-2026-results", + "original_job_id": "job_e28e0306a02a", + "s3_key": "uploads/job_e28e0306a02a.pdf", + "file_name": "Meta Q1 2026 Exhibit 99.1.pdf", + "asset_directory": "financial-meta-q1-2026-results" + }, + { + "source_id": "demo-financial-meta-q1-2026-presentation", + "original_job_id": "job_08cdaef3d6af", + "s3_key": "uploads/job_08cdaef3d6af.pdf", + "file_name": "Meta Q1 2026 Earnings Presentation.pdf", + "asset_directory": "financial-meta-q1-2026-presentation" + }, + { + "source_id": "demo-financial-meta-q1-2026-earnings-call", + "original_job_id": "job_fc7d8542b353", + "s3_key": "uploads/job_fc7d8542b353.pdf", + "file_name": "Meta Q1 2026 Earnings Call Transcript.pdf", + "asset_directory": "financial-meta-q1-2026-earnings-call" + }, + { + "source_id": "demo-financial-microsoft-2025-annual-report", + "original_job_id": "job_b9eb8b0f3f50", + "s3_key": "uploads/job_b9eb8b0f3f50.docx", + "file_name": "Microsoft 2025 Annual Report.docx", + "asset_directory": "financial-microsoft-2025-annual-report" + }, + { + "source_id": "demo-financial-alibaba-fy2026-interim-report", + "original_job_id": "job_8aacc19b54f8", + "s3_key": "uploads/job_8aacc19b54f8.pdf", + "file_name": "Alibaba Fiscal Year 2026 Interim Report.pdf", + "asset_directory": "financial-alibaba-fy2026-interim-report" + }, + { + "source_id": "demo-research-attention-is-all-you-need", + "original_job_id": "job_b583b4c5c4c0", + "s3_key": "uploads/job_b583b4c5c4c0.pdf", + "file_name": "Attention Is All You Need.pdf", + "asset_directory": "research-attention-is-all-you-need" + }, + { + "source_id": "demo-research-rag-survey", + "original_job_id": "job_9f032493aca9", + "s3_key": "uploads/job_9f032493aca9.pdf", + "file_name": "Retrieval-Augmented Generation for Large Language Models: A Survey.pdf", + "asset_directory": "research-rag-survey" + }, + { + "source_id": "demo-research-rag-realized", + "original_job_id": "job_e9c30a85c60b", + "s3_key": "uploads/job_e9c30a85c60b.pdf", + "file_name": "Retrieval-Augmented Generation Realized.pdf", + "asset_directory": "research-rag-realized" + }, + { + "source_id": "demo-research-toolformer", + "original_job_id": "job_8e54b45fba36", + "s3_key": "uploads/job_8e54b45fba36.pdf", + "file_name": "Toolformer: Language Models Can Teach Themselves to Use Tools.pdf", + "asset_directory": "research-toolformer" + }, + { + "source_id": "demo-research-ai-agents-overview", + "original_job_id": "job_36b03afb3724", + "s3_key": "uploads/job_36b03afb3724.pdf", + "file_name": "AI Agents Overview.pdf", + "asset_directory": "research-ai-agents-overview" + }, + { + "source_id": "demo-stem-jurafsky-transformers", + "original_job_id": "job_3a655ead76a2", + "s3_key": "uploads/job_3a655ead76a2.pdf", + "file_name": "Speech and Language Processing, Chapter 8: Transformers.pdf", + "asset_directory": "stem-jurafsky-transformers" + }, + { + "source_id": "demo-stem-transformers-tutorial", + "original_job_id": "job_0ad9b3c443ad", + "s3_key": "uploads/job_0ad9b3c443ad.pdf", + "file_name": "Transformers Tutorial.pdf", + "asset_directory": "stem-transformers-tutorial" + }, + { + "source_id": "demo-stem-deep-learning-transformer-network", + "original_job_id": "job_98728272ac53", + "s3_key": "uploads/job_98728272ac53.pdf", + "file_name": "Deep Learning - Transformer Network.pdf", + "asset_directory": "stem-deep-learning-transformer-network" + }, + { + "source_id": "demo-stem-probability-theory-lecture-notes", + "original_job_id": "job_459e062c475a", + "s3_key": "uploads/job_459e062c475a.pdf", + "file_name": "Probability Theory Lecture Notes.pdf", + "asset_directory": "stem-probability-theory-lecture-notes" + }, + { + "source_id": "demo-stem-introduction-statistical-learning-theory", + "original_job_id": "job_e2f0b8c6961c", + "s3_key": "uploads/job_e2f0b8c6961c.pdf", + "file_name": "Introduction to Statistical Learning Theory.pdf", + "asset_directory": "stem-introduction-statistical-learning-theory" + }, + { + "source_id": "demo-stem-statistical-learning", + "original_job_id": "job_d5b4e9750b05", + "s3_key": "uploads/job_d5b4e9750b05.pdf", + "file_name": "Statistical Learning Notes.pdf", + "asset_directory": "stem-statistical-learning" + } +] diff --git a/scripts/publication_benchmark/README.md b/scripts/publication_benchmark/README.md index 6977c446e..1706220d6 100644 --- a/scripts/publication_benchmark/README.md +++ b/scripts/publication_benchmark/README.md @@ -75,10 +75,9 @@ The example commands show one matrix point. Re-run the runner for the other owner (`async`), strategy (`candidate`), mode (`warm`), and case identifiers. Shell alternation syntax must not be used in command lines. -Each strategy sample must run in its own Python process. Do not import the -runner once and switch `KNOWHERE_PUBLICATION_STRATEGY` between samples: the -shared application settings are process-scoped, and a reused process can -otherwise report one strategy while executing another. +The optimized publication implementation is unconditional. Active runners accept +only `candidate`; baseline labels remain in historical reports for comparison. +Run each sample in its own process to isolate settings and connection pools. `formal_campaign_report.py` validates the split 59-round compact-template campaign across 236 independent clones. It checks sample identity, frozen diff --git a/scripts/publication_benchmark/capacity_execution.py b/scripts/publication_benchmark/capacity_execution.py index 0cc1dfd6d..6fce6091b 100644 --- a/scripts/publication_benchmark/capacity_execution.py +++ b/scripts/publication_benchmark/capacity_execution.py @@ -6,7 +6,6 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from hashlib import sha256 -import os import subprocess from threading import Barrier, Event from time import monotonic @@ -248,10 +247,9 @@ def run_capacity_batch(batch: CapacityBatch) -> dict[str, Any]: run_id=batch.run_id, strategy=batch.strategy, ) - from shared.core.config import settings - settings.KNOWHERE_PUBLICATION_STRATEGY = batch.strategy - os.environ["KNOWHERE_PUBLICATION_STRATEGY"] = batch.strategy + if batch.strategy != "candidate": + raise ValueError("Only the optimized publication implementation is available.") quiet_samples = asyncio.run( run_open_loop_probes( database_url=batch.database_url, diff --git a/scripts/publication_benchmark/guards.py b/scripts/publication_benchmark/guards.py index a2f709a82..d296d389e 100644 --- a/scripts/publication_benchmark/guards.py +++ b/scripts/publication_benchmark/guards.py @@ -148,7 +148,7 @@ def assert_strategy_matches_environment( if str(requested_strategy).strip().lower() != str(environment_value).strip().lower(): raise BenchmarkGuardError( f"{purpose}: requested strategy {requested_strategy!r} does not " - f"match KNOWHERE_PUBLICATION_STRATEGY {environment_value!r}" + f"match the available publication implementation {environment_value!r}" ) diff --git a/scripts/publication_benchmark/publication_execution.py b/scripts/publication_benchmark/publication_execution.py index db7e33fc6..b003e5d4f 100644 --- a/scripts/publication_benchmark/publication_execution.py +++ b/scripts/publication_benchmark/publication_execution.py @@ -83,7 +83,6 @@ def build_publication_environment( "DS_URL": "https://benchmark.invalid/v1", "TELEMETRY_ENABLED": "false", "LOGFIRE_TOKEN": "", - "KNOWHERE_PUBLICATION_STRATEGY": strategy, } diff --git a/scripts/publication_benchmark/run_publication.py b/scripts/publication_benchmark/run_publication.py index 2b28cd64b..b31bf2834 100644 --- a/scripts/publication_benchmark/run_publication.py +++ b/scripts/publication_benchmark/run_publication.py @@ -4,7 +4,7 @@ --owner sync \ --db-url-file .benchmarks/publication/clones//database-url \ --input .benchmarks/publication/inputs/spacex-s1-production \ - --strategy baseline \ + --strategy candidate \ --mode cold """ @@ -12,7 +12,6 @@ import argparse import hashlib -import os import sys from pathlib import Path from typing import Any, Sequence @@ -52,7 +51,6 @@ resolve_source_content_digest, ) -STRATEGY_ENVIRONMENT_VARIABLE: str = "KNOWHERE_PUBLICATION_STRATEGY" def default_run_id() -> str: @@ -147,19 +145,7 @@ def run_publication_point( f"clone {resolved_run_id!r} is in state {clone_record.state!r}" ) - os.environ[STRATEGY_ENVIRONMENT_VARIABLE] = strategy - assert_strategy_matches_environment( - strategy, - os.environ[STRATEGY_ENVIRONMENT_VARIABLE], - purpose="run_publication", - ) - from shared.core.config import settings as live_settings - - assert_strategy_matches_environment( - strategy, - str(live_settings.KNOWHERE_PUBLICATION_STRATEGY), - purpose="run_publication(live settings)", - ) + assert_strategy_matches_environment(strategy, "candidate", purpose="run_publication") scope_suffix = hashlib.sha256(resolved_run_id.encode("utf-8")).hexdigest()[:16] scope = PublicationScope( @@ -339,7 +325,7 @@ def build_argument_parser() -> argparse.ArgumentParser: parser.add_argument( "--strategy", required=True, - choices=("baseline", "candidate"), + choices=("candidate",), ) parser.add_argument("--mode", required=True, choices=PUBLICATION_MODES) parser.add_argument("--run-id", default=None) diff --git a/scripts/publication_benchmark/run_record.py b/scripts/publication_benchmark/run_record.py index d0cb61e40..dda934e72 100644 --- a/scripts/publication_benchmark/run_record.py +++ b/scripts/publication_benchmark/run_record.py @@ -230,17 +230,25 @@ def resolve_source_content_digest(repository_root: Path) -> str: capture_output=True, check=True, ) + deleted = subprocess.run( + ["git", "ls-files", "--deleted", "-z"], + cwd=repository_root, capture_output=True, check=True, + ) except (OSError, subprocess.CalledProcessError) as error: raise RunRecordError("could not list benchmark source files") from error paths = sorted(set(completed.stdout.split(b"\0")) - {b""}) + deleted_paths = set(deleted.stdout.split(b"\0")) - {b""} digest = hashlib.sha256() for raw_path in paths: if raw_path == b".benchmarks" or raw_path.startswith(b".benchmarks/"): continue path = repository_root / os.fsdecode(raw_path) try: - file_mode = path.lstat().st_mode - if stat.S_ISLNK(file_mode): + file_mode = path.lstat().st_mode if raw_path not in deleted_paths else None + if file_mode is None: + content = b"" + file_kind = b"deleted" + elif stat.S_ISLNK(file_mode): content = os.readlink(path).encode("utf-8", errors="surrogateescape") file_kind = b"link" elif stat.S_ISREG(file_mode): diff --git a/scripts/restore-demo-originals.py b/scripts/restore-demo-originals.py new file mode 100644 index 000000000..be5cf5bb0 --- /dev/null +++ b/scripts/restore-demo-originals.py @@ -0,0 +1,317 @@ +"""Manually restore verified demo originals through the v2 HTTP jobs API.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import httpx + + +@dataclass(frozen=True) +class RestoreRequest: + api_url: str + api_key: str + backup_directory: Path + state_directory: Path + timeout_seconds: float + can_reparse_ready: bool + can_retry_failed: bool + + +def write_state(path: Path, state: dict[str, Any]) -> None: + temporary: Path = path.with_suffix(".partial") + temporary.write_text(json.dumps(state, indent=2) + "\n") + temporary.replace(path) + + +def validate_backup(entry: dict[str, Any], directory: Path) -> Path: + path: Path = (directory / str(entry["relative_path"])).resolve() + if ( + not path.is_relative_to(directory.resolve()) + or entry.get("status") != "verified" + ): + raise ValueError("Original is outside the verified backup inventory") + if not isinstance(entry.get("validation"), dict) or not entry["validation"].get( + "is_valid" + ): + raise ValueError("Original format validation is missing") + digest = hashlib.sha256() + with path.open("rb") as original: + for block in iter(lambda: original.read(1024 * 1024), b""): + digest.update(block) + if ( + path.stat().st_size != entry["size_bytes"] + or digest.hexdigest() != entry["sha256"] + ): + raise ValueError("Original size or SHA-256 differs from the verified inventory") + return path + + +def request_json( + client: httpx.Client, method: str, path: str, **arguments: Any +) -> dict[str, Any]: + response: httpx.Response | None = None + for attempt in range(4): + try: + response = client.request(method, path, **arguments) + break + except httpx.TransportError: + if method != "GET" or attempt == 3: + raise + time.sleep(min(2**attempt, 8)) + if response is None: + raise RuntimeError("API request attempts completed without a response") + if response.is_error: + # Never persist response bodies, credentials, storage keys or signed URLs. + try: + errorCode: str = str( + response.json().get("error", {}).get("code", "HTTP_ERROR") + ) + except (ValueError, AttributeError): + errorCode = "HTTP_ERROR" + raise RuntimeError(f"API returned {response.status_code} {errorCode}") + payload: object = response.json() + if not isinstance(payload, dict): + raise ValueError("API response must be an object") + return payload + + +def restore_source(entry: dict[str, Any], request: RestoreRequest) -> dict[str, Any]: + sourceId: str = str(entry["source_id"]) + statePath: Path = request.state_directory / ( + hashlib.sha256(sourceId.encode()).hexdigest() + ".json" + ) + state: dict[str, Any] = ( + json.loads(statePath.read_text()) + if statePath.exists() + else {"source_id": sourceId} + ) + startedAt: float = time.monotonic() + try: + originalPath: Path = validate_backup(entry, request.backup_directory) + with httpx.Client( + base_url=request.api_url.rstrip("/") + "/", + headers={"Authorization": "Bearer " + request.api_key}, + timeout=60, + ) as client: + job: dict[str, Any] | None = None + if state.get("job_id"): + job = request_json(client, "GET", "api/v2/jobs/" + state["job_id"]) + if ( + job.get("namespace") != "__knowhere_demo__" + or job.get("data_id") != sourceId + or state.get("sha256") != entry["sha256"] + ): + raise ValueError( + "Recorded job does not match this source and verified original" + ) + if job.get("status") == "done" and request.can_reparse_ready: + state["previous_job_ids"] = [ + *state.get("previous_job_ids", []), + state.pop("job_id"), + ] + job = None + if job is not None and job.get("status") == "failed": + if not request.can_retry_failed: + raise RuntimeError( + "Prior job failed; use --retry-failed to create another revision" + ) + state["previous_job_ids"] = [ + *state.get("previous_job_ids", []), + state.pop("job_id"), + ] + job = None + if job is None: + createStarted: float = time.monotonic() + job = request_json( + client, + "POST", + "api/v2/jobs", + json={ + "namespace": "__knowhere_demo__", + "source_type": "file", + "file_name": entry["file_name"], + "data_id": sourceId, + }, + ) + state.update( + { + "job_id": job["job_id"], + "document_id": job["document_id"], + "status": "waiting-file", + "sha256": entry["sha256"], + "create_seconds": round(time.monotonic() - createStarted, 3), + } + ) + write_state(statePath, state) + if job.get("status") == "waiting-file" or job.get("upload_url"): + if not job.get("upload_url"): + job.update( + request_json( + client, + "POST", + f"api/v2/jobs/{state['job_id']}/upload-url", + json={}, + ) + ) + uploadUrl: str = str(job.get("upload_url") or "") + if not uploadUrl: + raise RuntimeError( + "Waiting job has no upload URL; inspect it before resuming" + ) + uploadHeaders: dict[str, str] = job.get("upload_headers") or { + "Content-Type": "application/pdf" + if originalPath.suffix.lower() == ".pdf" + else "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + } + # Use a separate client so the API bearer never reaches storage. + uploadStarted: float = time.monotonic() + with ( + originalPath.open("rb") as original, + httpx.Client(timeout=600) as storageClient, + ): + uploaded: httpx.Response = storageClient.put( + uploadUrl, + content=original, + headers={ + **uploadHeaders, + "Content-Length": str(originalPath.stat().st_size), + }, + ) + uploaded.raise_for_status() + state["upload_seconds"] = round(time.monotonic() - uploadStarted, 3) + confirmStarted: float = time.monotonic() + request_json( + client, + "POST", + f"api/v2/jobs/{state['job_id']}/confirm-upload", + json={}, + ) + state["confirm_seconds"] = round(time.monotonic() - confirmStarted, 3) + write_state(statePath, state) + interval: float = 2 + while True: + job = request_json(client, "GET", f"api/v2/jobs/{state['job_id']}") + state["status"] = str(job.get("status")) + write_state(statePath, state) + if state["status"] in ("done", "failed"): + break + if time.monotonic() - startedAt > request.timeout_seconds: + raise TimeoutError( + "Polling deadline reached; job is retained for --resume" + ) + time.sleep(interval) + interval = min(interval * 1.5, 30) + if state["status"] != "done": + raise RuntimeError( + "Parser job failed; inspect the authorized job response" + ) + catalog: dict[str, Any] = request_json(client, "GET", "api/v1/demo/catalog") + source: dict[str, Any] | None = next( + ( + item + for item in catalog["sources"] + if item["demo_source_id"] == sourceId + ), + None, + ) + if ( + source is None + or source["canonical_document_id"] != state["document_id"] + ): + raise RuntimeError("Completed job is absent from the READY catalog") + revisionId: str = str(job.get("job_result_id") or "") + if not revisionId: + raise RuntimeError( + "Completed job did not return its immutable revision" + ) + chunks: dict[str, Any] = request_json( + client, + "GET", + f"api/v2/documents/{state['document_id']}/chunks", + params={"job_result_id": revisionId, "page_size": 1}, + ) + state.update( + { + "status": "ready", + "job_result_id": revisionId, + "chunk_count": chunks["pagination"]["total"], + "elapsed_seconds": round(time.monotonic() - startedAt, 3), + } + ) + state.pop("error_type", None) + except Exception as error: + state["error_type"] = type(error).__name__ + state["status"] = "restore_failed" + print(f"{sourceId}: failed ({type(error).__name__})", flush=True) + else: + print(f"{sourceId}: ready ({state['chunk_count']} chunks)", flush=True) + write_state(statePath, state) + return state + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--api-url", required=True, help="Origin URL, without /api") + parser.add_argument("--api-key-env", default="KNOWHERE_DEMO_API_KEY") + parser.add_argument("--inventory", type=Path, required=True) + parser.add_argument( + "--source", + action="append", + help="Source ID; repeat to select several, omit for the full inventory", + ) + parser.add_argument("--state-directory", type=Path, required=True) + parser.add_argument("--concurrency", type=int, default=2, choices=range(1, 9)) + parser.add_argument("--timeout-seconds", type=float, default=14400) + parser.add_argument( + "--resume", + action="store_true", + help="Reuse recorded jobs (always on; provided for explicit operator intent)", + ) + parser.add_argument( + "--reparse-ready", + action="store_true", + help="Create a new revision for already completed recorded jobs", + ) + parser.add_argument("--retry-failed", action="store_true") + arguments: argparse.Namespace = parser.parse_args() + apiKey: str = os.environ.get(arguments.api_key_env, "") + if not apiKey: + parser.error("Maintainer API key environment variable is empty") + inventory: list[dict[str, Any]] = json.loads(arguments.inventory.read_text()) + if any(entry.get("status") != "verified" for entry in inventory): + parser.error("Backup inventory is incomplete; restoration is blocked") + sources: set[str] = set(arguments.source or []) + if sources - {entry["source_id"] for entry in inventory}: + parser.error("Requested source is absent from the inventory") + selected: list[dict[str, Any]] = [ + entry for entry in inventory if not sources or entry["source_id"] in sources + ] + arguments.state_directory.mkdir(parents=True, exist_ok=True) + request: RestoreRequest = RestoreRequest( + arguments.api_url, + apiKey, + arguments.inventory.parent, + arguments.state_directory, + arguments.timeout_seconds, + arguments.reparse_ready, + arguments.retry_failed, + ) + with ThreadPoolExecutor(max_workers=arguments.concurrency) as executor: + results: list[dict[str, Any]] = list( + executor.map(lambda entry: restore_source(entry, request), selected) + ) + write_state(arguments.state_directory / "report.json", {"sources": results}) + return 0 if all(result["status"] == "ready" for result in results) else 1 + + +if __name__ == "__main__": + raise SystemExit(main())