diff --git a/docs/AUTHENTICATION.md b/docs/AUTHENTICATION.md index d507857..37d0ff3 100644 --- a/docs/AUTHENTICATION.md +++ b/docs/AUTHENTICATION.md @@ -62,18 +62,16 @@ as a reverse proxy to handle OIDC authentication. This approach: ### Architecture -```text -┌──────────┐ ┌───────────────┐ ┌─────────────────────┐ -│ User │────▶│ OAuth2 Proxy │────▶│ Soliplex Ingester │ -│ Browser │ │ (OIDC) │ │ (API + UI) │ -└──────────┘ └───────────────┘ └─────────────────────┘ - │ - ▼ - ┌──────────────┐ - │ OIDC Provider │ - │ (Keycloak, │ - │ Auth0, etc) │ - └──────────────┘ +```mermaid +graph LR + User([User Browser]) -->|1. Access App| OAuth2[OAuth2 Proxy
OIDC] + OAuth2 -->|2. Authenticated
+ Headers| Ingester[Soliplex Ingester
API + UI] + OAuth2 -->|3. Verify Token| OIDC[OIDC Provider
Keycloak, Auth0,
Okta, etc.] + + style User fill:#4a90e2,stroke:#2e5c8a,color:#fff + style OAuth2 fill:#3ba13b,stroke:#2d7a2d,color:#fff + style Ingester fill:#f39c12,stroke:#c27d0e,color:#fff + style OIDC fill:#9b59b6,stroke:#7a4691,color:#fff ``` --- diff --git a/pyproject.toml b/pyproject.toml index 6f8bfaa..19872de 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,6 @@ dependencies = [ "aiosqlite<0.23.0", "alembic>=1.17.2", "async-lru>=2.0.5", - "docling-core>=2.51.1", "fastapi[standard]>=0.125.0", "fsspec>=2025.12.0", "haiku-rag-slim>=0.30.2", diff --git a/src/soliplex/ingester/cli.py b/src/soliplex/ingester/cli.py index 1d648dc..b7a5cd1 100644 --- a/src/soliplex/ingester/cli.py +++ b/src/soliplex/ingester/cli.py @@ -316,6 +316,63 @@ def list_batches(): asyncio.run(_list_batches()) +async def _check_db(db_name: str, lancedb_dir: str | None): + from .lib.models import Database + from .lib.operations import check_rag_db_consistency + + await Database.initialize() + result = await check_rag_db_consistency(db_name, lancedb_dir) + + print("\n[bold]Database Consistency Check[/bold]") + print(f" db_name: {result['db_name']}") + print(f" lancedb_dir: {result['lancedb_dir']}") + print(f" db_path: {result['db_path']}") + + if "error" in result: + print(f"\n[red]Error:[/red] {result['error']}") + + print("\n[bold]Summary:[/bold]") + print(f" DocumentDB records: {result['documentdb_count']}") + print(f" LanceDB documents: {result['lancedb_count']}") + print(f" Matched: {result['matched']}") + + if result["in_documentdb_only"]: + print(f"\n[yellow]In DocumentDB but NOT in LanceDB ({len(result['in_documentdb_only'])}):[/yellow]") + for doc in result["in_documentdb_only"]: + print(f" - rag_id: {doc['rag_id']}") + print(f" uri: {doc['uri']}") + print(f" hash: {doc['doc_hash']}") + else: + print("\n[green]No documents in DocumentDB missing from LanceDB[/green]") + + if result["in_lancedb_only"]: + print(f"\n[yellow]In LanceDB but NOT in DocumentDB ({len(result['in_lancedb_only'])}):[/yellow]") + for doc in result["in_lancedb_only"]: + print(f" - rag_id: {doc['rag_id']}") + print(f" uri: {doc.get('uri', 'N/A')}") + print(f" title: {doc.get('title', 'N/A')}") + else: + print("\n[green]No documents in LanceDB missing from DocumentDB[/green]") + + if not result["in_documentdb_only"] and not result["in_lancedb_only"] and "error" not in result: + print("\n[green]✓ Database is consistent[/green]") + + +@app.command("check-db") +def check_db( + db_name: str = typer.Argument(..., help="Database name (data_dir) to check"), + lancedb_dir: str = typer.Option(None, "--lancedb-dir", "-l", help="LanceDB directory (uses default if not specified)"), +): + """ + Check consistency between LanceDB database and DocumentDB records. + + Compares documents in the actual LanceDB database with records tracked + in DocumentDB to find discrepancies. + """ + validate_settings(dump=False) + asyncio.run(_check_db(db_name, lancedb_dir)) + + @app.command( "serve", ) diff --git a/src/soliplex/ingester/lib/models.py b/src/soliplex/ingester/lib/models.py index 242156c..7333e5b 100644 --- a/src/soliplex/ingester/lib/models.py +++ b/src/soliplex/ingester/lib/models.py @@ -141,6 +141,21 @@ def doc_hash(data: bytes) -> str: return f"sha256-{hex_digest}" +class DocumentDB(SQLModel, table=True): + id: int | None = Field( + default=None, + primary_key=True, + sa_column_kwargs={"autoincrement": True}, + ) + doc_hash: str = Field(default=None) # leave off foreign key due to possible missing documents + source: str = Field(default=None) + db_name: str = Field(default=None) + lancedb_dir: str = Field(default=None) + rag_id: str | None = Field(default=None) + chunk_count: int | None = Field(default=None) + created_date: datetime.datetime = Field(default=None) + + class DocumentBatch(SQLModel, table=True): """ A batch of documents to be ingested diff --git a/src/soliplex/ingester/lib/operations.py b/src/soliplex/ingester/lib/operations.py index d6458c8..d6bd0c5 100644 --- a/src/soliplex/ingester/lib/operations.py +++ b/src/soliplex/ingester/lib/operations.py @@ -9,6 +9,7 @@ from . import dal from . import models +from . import rag from .config import get_settings logger = logging.getLogger(__name__) @@ -543,18 +544,42 @@ async def delete_orphaned_documents(): """ Delete orphaned documents that have no URI pointing to them. - Also deletes orphaned history records for consistency. - Uses SQLModel ORM for cross-database compatibility. + Also deletes orphaned history records, HaikuRAG entries, and DocumentDB + records for consistency. Uses SQLModel ORM for cross-database compatibility. Returns ------- dict - Statistics about deleted records + Statistics about deleted records including: + - deleted_documents: Number of Document records deleted + - deleted_history: Number of DocumentURIHistory records deleted + - deleted_rag_entries: Number of HaikuRAG documents deleted + - deleted_documentdb_records: Number of DocumentDB records deleted """ + deleted_rag_entries = 0 + deleted_documentdb_records = 0 + async with models.get_session() as session: # Subquery: Get all hashes that are referenced referenced_hashes_subq = select(models.DocumentURI.doc_hash).distinct().subquery() + # Find orphaned document hashes first (before deletion) + orphaned_q = select(models.Document.hash).where( + models.Document.hash.not_in(select(referenced_hashes_subq.c.doc_hash)) + ) + orphaned_result = await session.exec(orphaned_q) + orphaned_hashes = list(orphaned_result.all()) + + # Delete from HaikuRAG for each orphaned hash (outside the session to avoid conflicts) + for doc_hash in orphaned_hashes: + rag_stats = await rag.delete_from_rag_by_hash(doc_hash) + deleted_rag_entries += rag_stats["deleted_rag_entries"] + deleted_documentdb_records += rag_stats["deleted_documentdb_records"] + + async with models.get_session() as session: + # Subquery: Get all hashes that are referenced (re-create for the new session) + referenced_hashes_subq = select(models.DocumentURI.doc_hash).distinct().subquery() + # Delete orphaned documents q1 = delete(models.Document).where(models.Document.hash.not_in(select(referenced_hashes_subq.c.doc_hash))) result1 = await session.exec(q1) @@ -569,11 +594,16 @@ async def delete_orphaned_documents(): await session.commit() - logger.info(f"Deleted {deleted_docs} orphaned documents and {deleted_history} history records") + logger.info( + f"Deleted {deleted_docs} orphaned documents, {deleted_history} history records, " + f"{deleted_rag_entries} RAG entries, {deleted_documentdb_records} DocumentDB records" + ) return { "deleted_documents": deleted_docs, "deleted_history": deleted_history, + "deleted_rag_entries": deleted_rag_entries, + "deleted_documentdb_records": deleted_documentdb_records, } @@ -779,6 +809,8 @@ async def delete_document_uri_by_uri(uri: str, source: str) -> dict[str, int]: deleted_workflow_runs = 0 deleted_run_steps = 0 deleted_lifecycle_history = 0 + deleted_rag_entries = 0 + deleted_documentdb_records = 0 # Step 3: If this is the only URI, cascade delete everything if uri_count == 1: @@ -818,6 +850,11 @@ async def delete_document_uri_by_uri(uri: str, source: str) -> dict[str, int]: extra=log_context(doc_hash=doc_hash, action="delete_document_uri_by_uri"), ) + # Delete from HaikuRAG and DocumentDB + rag_delete_stats = await rag.delete_from_rag_by_hash(doc_hash) + deleted_rag_entries = rag_delete_stats["deleted_rag_entries"] + deleted_documentdb_records = rag_delete_stats["deleted_documentdb_records"] + # Delete Document doc_delete_q = delete(models.Document).where(models.Document.hash == doc_hash) doc_result = await session.exec(doc_delete_q) @@ -842,6 +879,8 @@ async def delete_document_uri_by_uri(uri: str, source: str) -> dict[str, int]: "deleted_workflow_runs": deleted_workflow_runs, "deleted_run_steps": deleted_run_steps, "deleted_lifecycle_history": deleted_lifecycle_history, + "deleted_rag_entries": deleted_rag_entries, + "deleted_documentdb_records": deleted_documentdb_records, "total_deleted": ( deleted_document_uris + deleted_uri_history @@ -849,10 +888,268 @@ async def delete_document_uri_by_uri(uri: str, source: str) -> dict[str, int]: + deleted_workflow_runs + deleted_run_steps + deleted_lifecycle_history + + deleted_rag_entries + + deleted_documentdb_records ), } +async def list_documentdb_databases() -> list[dict]: + """ + List distinct lancedb_dir/db_name combinations from DocumentDB with document counts. + + Returns + ------- + list[dict] + List of dictionaries containing: + - lancedb_dir: The LanceDB directory path + - db_name: The database name (data_dir) + - document_count: Number of documents in this database + - total_chunks: Total chunks across all documents + """ + from sqlalchemy import func + + async with models.get_session() as session: + # Query distinct lancedb_dir/db_name with counts + q = ( + select( + models.DocumentDB.lancedb_dir, + models.DocumentDB.db_name, + func.count(models.DocumentDB.id).label("document_count"), + func.sum(models.DocumentDB.chunk_count).label("total_chunks"), + ) + .group_by(models.DocumentDB.lancedb_dir, models.DocumentDB.db_name) + .order_by(models.DocumentDB.lancedb_dir, models.DocumentDB.db_name) + ) + + result = await session.exec(q) + rows = result.all() + + databases = [] + for row in rows: + databases.append( + { + "lancedb_dir": row.lancedb_dir, + "db_name": row.db_name, + "document_count": row.document_count, + "total_chunks": row.total_chunks or 0, + } + ) + + return databases + + +async def list_documents_in_rag_db( + db_name: str, + lancedb_dir: str | None = None, +) -> list[dict]: + """ + List documents in a specific RAG database from DocumentDB records. + + Joins DocumentDB with DocumentURI and Document to provide full document information. + + Parameters + ---------- + db_name : str + The database name (data_dir) to filter by + lancedb_dir : str | None + The LanceDB directory path. If None, uses default from settings. + + Returns + ------- + list[dict] + List of dictionaries containing document information: + - doc_hash: The document hash + - rag_id: The HaikuRAG document ID + - chunk_count: Number of chunks for this document + - created_date: When the document was added to RAG + - uri: The document URI (from DocumentURI) + - source: The source system + - mime_type: Document MIME type + - file_size: Document file size in bytes + """ + if lancedb_dir is None: + lancedb_dir = get_settings().lancedb_dir + + async with models.get_session() as session: + # Query DocumentDB joined with DocumentURI and Document + q = ( + select( + models.DocumentDB.doc_hash, + models.DocumentDB.rag_id, + models.DocumentDB.chunk_count, + models.DocumentDB.created_date, + models.DocumentDB.source, + models.DocumentURI.uri, + models.Document.mime_type, + models.Document.file_size, + ) + .join(models.DocumentURI, models.DocumentDB.doc_hash == models.DocumentURI.doc_hash, isouter=True) + .join(models.Document, models.DocumentDB.doc_hash == models.Document.hash, isouter=True) + .where(models.DocumentDB.db_name == db_name) + .where(models.DocumentDB.lancedb_dir == lancedb_dir) + .order_by(models.DocumentDB.created_date.desc()) + ) + + result = await session.exec(q) + rows = result.all() + + documents = [] + for row in rows: + documents.append( + { + "doc_hash": row.doc_hash, + "rag_id": row.rag_id, + "chunk_count": row.chunk_count or 0, + "created_date": row.created_date.isoformat() if row.created_date else None, + "source": row.source, + "uri": row.uri, + "mime_type": row.mime_type, + "file_size": row.file_size, + } + ) + + return documents + + +async def check_rag_db_consistency( + db_name: str, + lancedb_dir: str | None = None, +) -> dict: + """ + Check consistency between LanceDB database and DocumentDB records. + + Compares documents in the actual LanceDB database with records tracked + in DocumentDB to find discrepancies. + + Parameters + ---------- + db_name : str + The database name (data_dir) to check + lancedb_dir : str | None + The LanceDB directory path. If None, uses default from settings. + + Returns + ------- + dict + Dictionary containing: + - db_name: The database name checked + - lancedb_dir: The LanceDB directory path + - db_path: Full path to the LanceDB database + - in_documentdb_only: List of documents in DocumentDB but not in LanceDB + - in_lancedb_only: List of documents in LanceDB but not in DocumentDB + - matched: Number of documents that exist in both + - documentdb_count: Total documents in DocumentDB for this database + - lancedb_count: Total documents in LanceDB database + """ + from pathlib import Path + + from haiku.rag.client import HaikuRAG + from haiku.rag.config import get_config + + settings = get_settings() + if lancedb_dir is None: + lancedb_dir = settings.lancedb_dir + + # Build db_path + if lancedb_dir.startswith("s3://"): + if lancedb_dir.endswith("/"): + db_path = f"{lancedb_dir}{db_name}" + else: + db_path = f"{lancedb_dir}/{db_name}" + else: + db_path = Path(lancedb_dir) / db_name + + # Get documents from DocumentDB + documentdb_docs = await list_documents_in_rag_db(db_name, lancedb_dir) + documentdb_rag_ids = {doc["rag_id"] for doc in documentdb_docs if doc["rag_id"]} + + # Get documents from LanceDB + lancedb_docs = [] + lancedb_rag_ids = set() + + try: + config = get_config() + async with HaikuRAG( + db_path=db_path, + config=config, + read_only=True, + ) as client: + documents = await client.list_documents() + for doc in documents: + doc_info = { + "rag_id": doc.id, + "uri": doc.uri, + "title": getattr(doc, "title", None), + "created_at": doc.created_at.isoformat() if getattr(doc, "created_at", None) else None, + "chunk_count": getattr(doc, "chunk_count", None), + "metadata": getattr(doc, "metadata", {}), + } + lancedb_docs.append(doc_info) + if doc.id: + lancedb_rag_ids.add(doc.id) + except Exception as e: + # Database might not exist or be accessible + return { + "db_name": db_name, + "lancedb_dir": lancedb_dir, + "db_path": str(db_path), + "error": f"Could not access LanceDB database: {e}", + "in_documentdb_only": [ + { + "rag_id": doc["rag_id"], + "doc_hash": doc["doc_hash"], + "uri": doc["uri"], + "source": doc["source"], + } + for doc in documentdb_docs + ], + "in_lancedb_only": [], + "matched": 0, + "documentdb_count": len(documentdb_docs), + "lancedb_count": 0, + } + + # Find discrepancies + in_documentdb_only_ids = documentdb_rag_ids - lancedb_rag_ids + in_lancedb_only_ids = lancedb_rag_ids - documentdb_rag_ids + matched_ids = documentdb_rag_ids & lancedb_rag_ids + + # Build detailed lists + in_documentdb_only = [ + { + "rag_id": doc["rag_id"], + "doc_hash": doc["doc_hash"], + "uri": doc["uri"], + "source": doc["source"], + } + for doc in documentdb_docs + if doc["rag_id"] in in_documentdb_only_ids + ] + + lancedb_by_rag_id = {doc["rag_id"]: doc for doc in lancedb_docs} + in_lancedb_only = [ + { + "rag_id": rag_id, + "uri": lancedb_by_rag_id[rag_id].get("uri"), + "title": lancedb_by_rag_id[rag_id].get("title"), + "metadata": lancedb_by_rag_id[rag_id].get("metadata"), + } + for rag_id in in_lancedb_only_ids + ] + + return { + "db_name": db_name, + "lancedb_dir": lancedb_dir, + "db_path": str(db_path), + "in_documentdb_only": in_documentdb_only, + "in_lancedb_only": in_lancedb_only, + "matched": len(matched_ids), + "documentdb_count": len(documentdb_docs), + "lancedb_count": len(lancedb_docs), + } + + async def delete_documents_by_hashes(doc_hashes: list[str]) -> dict[str, int]: """ Delete documents by their hashes with cascading deletion. diff --git a/src/soliplex/ingester/lib/rag.py b/src/soliplex/ingester/lib/rag.py index 10d0468..3cfcf1f 100644 --- a/src/soliplex/ingester/lib/rag.py +++ b/src/soliplex/ingester/lib/rag.py @@ -1,5 +1,6 @@ import asyncio import copy +import datetime import itertools import logging import pathlib @@ -12,10 +13,13 @@ from haiku.rag.embeddings import embed_chunks from haiku.rag.store.engine import DocumentRecord from haiku.rag.store.models.chunk import Chunk +from sqlalchemy import delete as sa_delete from . import models from .config import get_settings +from .models import DocumentDB from .models import StepConfig +from .models import get_session logger = logging.getLogger(__name__) @@ -220,3 +224,126 @@ async def save_to_rag( docling_document=docling_document, ) return new_doc.id + + +async def create_document_db_record( + doc_hash: str, + source: str, + step_config: StepConfig, + rag_id: str, + chunk_count: int, +) -> DocumentDB: + """ + Create a DocumentDB record to track a document stored in HaikuRAG. + + Parameters + ---------- + doc_hash : str + The document hash (sha256) + source : str + The source system identifier + step_config : StepConfig + The step configuration containing data_dir + rag_id : str + The ID returned by HaikuRAG after import + chunk_count : int + Number of chunks stored for this document + + Returns + ------- + DocumentDB + The created DocumentDB record + """ + env = get_settings() + config_dict = step_config.config_json + db_name = config_dict.get("data_dir", "") + lancedb_dir = env.lancedb_dir + + async with get_session() as session: + record = DocumentDB( + doc_hash=doc_hash, + source=source, + db_name=db_name, + lancedb_dir=lancedb_dir, + rag_id=rag_id, + chunk_count=chunk_count, + created_date=datetime.datetime.now(datetime.UTC), + ) + session.add(record) + await session.flush() + await session.refresh(record) + session.expunge(record) + # session commits automatically on context exit + return record + + +async def delete_from_rag_by_hash(doc_hash: str) -> dict[str, int]: + """ + Delete all HaikuRAG entries and DocumentDB records for a document hash. + + Queries DocumentDB for all records matching the hash, deletes each + document from its respective HaikuRAG database, then removes the + DocumentDB records. + + Parameters + ---------- + doc_hash : str + The document hash to delete + + Returns + ------- + dict[str, int] + Statistics containing: + - deleted_rag_entries: Number of HaikuRAG documents deleted + - deleted_documentdb_records: Number of DocumentDB records deleted + """ + deleted_rag_entries = 0 + deleted_documentdb_records = 0 + + async with get_session() as session: + # Find all DocumentDB records for this hash + from sqlmodel import select + + q = select(DocumentDB).where(DocumentDB.doc_hash == doc_hash) + result = await session.exec(q) + records = result.all() + + for record in records: + # Reconstruct the db_path + if record.lancedb_dir.startswith("s3://"): + if record.lancedb_dir.endswith("/"): + db_path = f"{record.lancedb_dir}{record.db_name}" + else: + db_path = f"{record.lancedb_dir}/{record.db_name}" + else: + import pathlib + + db_path = pathlib.Path(record.lancedb_dir) / record.db_name + + # Try to delete from HaikuRAG + if record.rag_id: + try: + # Build minimal config for deletion + config = build_embed_config(HRConfig, {"model": "dummy", "vector_dim": 1, "provider": "ollama"}) + config = build_storage_config(config, {"data_dir": record.db_name}) + + async with _rag_lock: + async with HaikuRAG(config=config, create=False, db_path=db_path) as client: + await client.delete_document(record.rag_id) + deleted_rag_entries += 1 + logger.info(f"Deleted document {record.rag_id} from HaikuRAG at {db_path}") + except Exception as e: + logger.warning(f"Failed to delete document {record.rag_id} from HaikuRAG at {db_path}: {e}") + + # Delete all DocumentDB records for this hash + if records: + delete_q = sa_delete(DocumentDB).where(DocumentDB.doc_hash == doc_hash) + delete_result = await session.exec(delete_q) + deleted_documentdb_records = delete_result.rowcount # type: ignore + + # session commits automatically on context exit + + return { + "deleted_rag_entries": deleted_rag_entries, + "deleted_documentdb_records": deleted_documentdb_records, + } diff --git a/src/soliplex/ingester/lib/wf/operations.py b/src/soliplex/ingester/lib/wf/operations.py index fec6118..2eeff6f 100644 --- a/src/soliplex/ingester/lib/wf/operations.py +++ b/src/soliplex/ingester/lib/wf/operations.py @@ -2,6 +2,7 @@ import json import logging +import docling_core.types.doc.document as docling_docs import opendal import yaml from sqlalchemy import Integer @@ -159,6 +160,8 @@ async def get_step_config_ids(param_id: str) -> dict[WorkflowStepType, int]: step_config = param_set.config[st] else: step_config = {} + if st == WorkflowStepType.PARSE: + step_config["docling_version"] = docling_docs.CURRENT_VERSION cuml_cfg = cuml_cfg.copy() cuml_cfg.update({st.value: step_config}) cuml_str = json.dumps(cuml_cfg, indent=4) diff --git a/src/soliplex/ingester/lib/workflow.py b/src/soliplex/ingester/lib/workflow.py index 7a467aa..4725860 100644 --- a/src/soliplex/ingester/lib/workflow.py +++ b/src/soliplex/ingester/lib/workflow.py @@ -514,6 +514,16 @@ async def save_to_rag( embed_config, _log_con, ) + + # Create DocumentDB record to track this document in HaikuRAG + await rag.create_document_db_record( + doc_hash=doc_hash, + source=source_uri.source, + step_config=step_config, + rag_id=rag_id, + chunk_count=len(chunk_objs), + ) + await doc_ops.add_history_for_hash(doc_hash, "ingested", hist_meta={"haiku_id": rag_id}, batch_id=batch_id) logger.info(f"save_to_rag completed {source} {batch_id} {doc_hash}") diff --git a/src/soliplex/ingester/server/routes/lancedb.py b/src/soliplex/ingester/server/routes/lancedb.py index d0503f9..42986a1 100644 --- a/src/soliplex/ingester/server/routes/lancedb.py +++ b/src/soliplex/ingester/server/routes/lancedb.py @@ -12,6 +12,7 @@ from haiku.rag.app import HaikuRAGApp from haiku.rag.config import get_config +from soliplex.ingester.lib import operations from soliplex.ingester.lib.auth import get_current_user from soliplex.ingester.lib.config import get_settings @@ -114,6 +115,98 @@ async def list_databases(): } +@lancedb_router.get("/list-dbs", status_code=status.HTTP_200_OK, summary="List tracked LanceDB databases") +async def list_tracked_databases(response: Response): + """ + List all LanceDB databases tracked in the DocumentDB table. + + Returns distinct lancedb_dir/db_name combinations with document counts. + This shows databases that have had documents ingested via the workflow, + as tracked by DocumentDB records. + """ + try: + databases = await operations.list_documentdb_databases() + return { + "status": "ok", + "database_count": len(databases), + "databases": databases, + } + except Exception as e: + logger.exception("Error listing tracked databases", exc_info=e) + response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR + return { + "status": "error", + "error": str(e), + } + + +@lancedb_router.get("/list-docs", status_code=status.HTTP_200_OK, summary="List documents in a tracked database") +async def list_tracked_documents( + response: Response, + db_name: str = Query(..., description="Database name (data_dir) to list documents from"), + lancedb_dir: str | None = Query(None, description="LanceDB directory. If not specified, uses default from settings."), +): + """ + List all documents tracked in a specific RAG database. + + Returns documents from the DocumentDB table joined with DocumentURI and Document + information for a given db_name (data_dir) and optional lancedb_dir. + """ + try: + documents = await operations.list_documents_in_rag_db(db_name, lancedb_dir) + settings = get_settings() + effective_lancedb_dir = lancedb_dir if lancedb_dir else settings.lancedb_dir + return { + "status": "ok", + "db_name": db_name, + "lancedb_dir": effective_lancedb_dir, + "document_count": len(documents), + "documents": documents, + } + except Exception as e: + logger.exception("Error listing tracked documents", exc_info=e) + response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR + return { + "status": "error", + "error": str(e), + } + + +@lancedb_router.get("/check-db", status_code=status.HTTP_200_OK, summary="Check database consistency") +async def check_database_consistency( + response: Response, + db_name: str = Query(..., description="Database name (data_dir) to check"), + lancedb_dir: str | None = Query(None, description="LanceDB directory. If not specified, uses default from settings."), +): + """ + Check consistency between a LanceDB database and DocumentDB records. + + Compares documents in the actual LanceDB database files with records tracked + in DocumentDB to find discrepancies: + - Documents tracked in DocumentDB but missing from LanceDB + - Documents in LanceDB but not tracked in DocumentDB + """ + try: + result = await operations.check_rag_db_consistency(db_name, lancedb_dir) + except Exception as e: + logger.exception("Error checking database consistency", exc_info=e) + response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR + return { + "status": "error", + "error": str(e), + } + else: + # Determine overall status + has_issues = ( + len(result.get("in_documentdb_only", [])) > 0 or len(result.get("in_lancedb_only", [])) > 0 or "error" in result + ) + + return { + "status": "warning" if has_issues else "ok", + **result, + } + + @lancedb_router.get("/info", status_code=status.HTTP_200_OK, summary="Get LanceDB database info") async def get_info( response: Response, diff --git a/tests/unit/test_documentdb.py b/tests/unit/test_documentdb.py new file mode 100644 index 0000000..5714cd0 --- /dev/null +++ b/tests/unit/test_documentdb.py @@ -0,0 +1,1032 @@ +""" +Unit tests for DocumentDB record creation and deletion. + +Tests the DocumentDB tracking of documents stored in HaikuRAG, +including record creation, deletion, and cascading through +delete_document_uri_by_uri. +""" + +import tempfile +from pathlib import Path +from unittest.mock import AsyncMock +from unittest.mock import MagicMock +from unittest.mock import patch + +import pytest +from sqlmodel import select + +import soliplex.ingester.lib.operations as doc_ops +from soliplex.ingester.lib import rag +from soliplex.ingester.lib.models import DocumentDB +from soliplex.ingester.lib.models import StepConfig +from soliplex.ingester.lib.models import get_session + + +@pytest.fixture +def mock_settings(): + """Create a mock settings object""" + settings = MagicMock() + settings.lancedb_dir = "/tmp/lancedb" + return settings + + +@pytest.fixture +def mock_step_config(): + """Create a mock step config with data_dir""" + step_config = MagicMock(spec=StepConfig) + step_config.config_json = {"data_dir": "test-project"} + return step_config + + +@pytest.mark.asyncio +async def test_create_document_db_record(db, mock_settings, mock_step_config): + """Test that create_document_db_record creates a record with correct fields.""" + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + record = await rag.create_document_db_record( + doc_hash="sha256-test123", + source="test_source", + step_config=mock_step_config, + rag_id="haiku-rag-id-456", + chunk_count=42, + ) + + # Verify returned record has correct values + assert record.doc_hash == "sha256-test123" + assert record.source == "test_source" + assert record.db_name == "test-project" + assert record.lancedb_dir == "/tmp/lancedb" + assert record.rag_id == "haiku-rag-id-456" + assert record.chunk_count == 42 + assert record.created_date is not None + assert record.id is not None + + # Verify record persisted to database + async with get_session() as session: + q = select(DocumentDB).where(DocumentDB.id == record.id) + result = await session.exec(q) + db_record = result.first() + assert db_record is not None + assert db_record.doc_hash == "sha256-test123" + assert db_record.rag_id == "haiku-rag-id-456" + + +@pytest.mark.asyncio +async def test_create_document_db_record_s3_lancedb_dir(db, mock_step_config): + """Test that create_document_db_record handles S3 lancedb_dir correctly.""" + mock_settings_s3 = MagicMock() + mock_settings_s3.lancedb_dir = "s3://my-bucket/lancedb" + + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings_s3): + record = await rag.create_document_db_record( + doc_hash="sha256-s3test", + source="s3_source", + step_config=mock_step_config, + rag_id="s3-rag-id", + chunk_count=10, + ) + + assert record.lancedb_dir == "s3://my-bucket/lancedb" + assert record.db_name == "test-project" + + +@pytest.mark.asyncio +async def test_multiple_documentdb_records_for_same_hash(db, mock_settings): + """Test that multiple DocumentDB records can exist for the same doc_hash (1:many mapping).""" + step_config1 = MagicMock(spec=StepConfig) + step_config1.config_json = {"data_dir": "project-a"} + + step_config2 = MagicMock(spec=StepConfig) + step_config2.config_json = {"data_dir": "project-b"} + + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + # Create two records for the same doc_hash but different data_dirs + await rag.create_document_db_record( + doc_hash="sha256-shared", + source="source_a", + step_config=step_config1, + rag_id="rag-id-1", + chunk_count=10, + ) + + await rag.create_document_db_record( + doc_hash="sha256-shared", + source="source_b", + step_config=step_config2, + rag_id="rag-id-2", + chunk_count=20, + ) + + # Verify both records exist + async with get_session() as session: + q = select(DocumentDB).where(DocumentDB.doc_hash == "sha256-shared") + result = await session.exec(q) + records = result.all() + assert len(records) == 2 + db_names = {r.db_name for r in records} + assert db_names == {"project-a", "project-b"} + + +@pytest.mark.asyncio +async def test_delete_from_rag_by_hash_no_records(db): + """Test that delete_from_rag_by_hash returns zeros when no records exist.""" + result = await rag.delete_from_rag_by_hash("sha256-nonexistent") + + assert result["deleted_rag_entries"] == 0 + assert result["deleted_documentdb_records"] == 0 + + +@pytest.mark.asyncio +async def test_delete_from_rag_by_hash_single_record(db, mock_settings, mock_step_config): + """Test deletion of a single DocumentDB record and its HaikuRAG entry.""" + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + # Create a DocumentDB record + await rag.create_document_db_record( + doc_hash="sha256-todelete", + source="test_source", + step_config=mock_step_config, + rag_id="rag-to-delete", + chunk_count=5, + ) + + # Mock HaikuRAG client for deletion + mock_client = MagicMock() + mock_client.delete_document = AsyncMock() + + with ( + patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings), + patch("soliplex.ingester.lib.rag.HaikuRAG") as mock_haiku_rag, + patch("soliplex.ingester.lib.rag.build_embed_config"), + patch("soliplex.ingester.lib.rag.build_storage_config"), + ): + mock_haiku_rag.return_value.__aenter__ = AsyncMock(return_value=mock_client) + mock_haiku_rag.return_value.__aexit__ = AsyncMock(return_value=None) + + result = await rag.delete_from_rag_by_hash("sha256-todelete") + + # Verify HaikuRAG deletion was called + mock_client.delete_document.assert_called_once_with("rag-to-delete") + + # Verify statistics + assert result["deleted_rag_entries"] == 1 + assert result["deleted_documentdb_records"] == 1 + + # Verify record removed from database + async with get_session() as session: + q = select(DocumentDB).where(DocumentDB.doc_hash == "sha256-todelete") + db_result = await session.exec(q) + assert db_result.first() is None + + +@pytest.mark.asyncio +async def test_delete_from_rag_by_hash_multiple_records(db, mock_settings): + """Test deletion of multiple DocumentDB records for the same hash.""" + step_config1 = MagicMock(spec=StepConfig) + step_config1.config_json = {"data_dir": "project-a"} + + step_config2 = MagicMock(spec=StepConfig) + step_config2.config_json = {"data_dir": "project-b"} + + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + # Create multiple records + await rag.create_document_db_record( + doc_hash="sha256-multi", + source="source_a", + step_config=step_config1, + rag_id="rag-id-a", + chunk_count=10, + ) + await rag.create_document_db_record( + doc_hash="sha256-multi", + source="source_b", + step_config=step_config2, + rag_id="rag-id-b", + chunk_count=20, + ) + + # Mock HaikuRAG client for deletion + mock_client = MagicMock() + mock_client.delete_document = AsyncMock() + + with ( + patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings), + patch("soliplex.ingester.lib.rag.HaikuRAG") as mock_haiku_rag, + patch("soliplex.ingester.lib.rag.build_embed_config"), + patch("soliplex.ingester.lib.rag.build_storage_config"), + ): + mock_haiku_rag.return_value.__aenter__ = AsyncMock(return_value=mock_client) + mock_haiku_rag.return_value.__aexit__ = AsyncMock(return_value=None) + + result = await rag.delete_from_rag_by_hash("sha256-multi") + + # Verify both HaikuRAG deletions were called + assert mock_client.delete_document.call_count == 2 + + # Verify statistics + assert result["deleted_rag_entries"] == 2 + assert result["deleted_documentdb_records"] == 2 + + # Verify all records removed from database + async with get_session() as session: + q = select(DocumentDB).where(DocumentDB.doc_hash == "sha256-multi") + db_result = await session.exec(q) + assert len(db_result.all()) == 0 + + +@pytest.mark.asyncio +async def test_delete_from_rag_by_hash_rag_failure_still_deletes_db_records(db, mock_settings, mock_step_config): + """Test that DocumentDB records are deleted even if HaikuRAG deletion fails.""" + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + await rag.create_document_db_record( + doc_hash="sha256-ragfail", + source="test_source", + step_config=mock_step_config, + rag_id="rag-will-fail", + chunk_count=5, + ) + + # Mock HaikuRAG client to raise an exception + mock_client = MagicMock() + mock_client.delete_document = AsyncMock(side_effect=Exception("Connection failed")) + + with ( + patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings), + patch("soliplex.ingester.lib.rag.HaikuRAG") as mock_haiku_rag, + patch("soliplex.ingester.lib.rag.build_embed_config"), + patch("soliplex.ingester.lib.rag.build_storage_config"), + ): + mock_haiku_rag.return_value.__aenter__ = AsyncMock(return_value=mock_client) + mock_haiku_rag.return_value.__aexit__ = AsyncMock(return_value=None) + + result = await rag.delete_from_rag_by_hash("sha256-ragfail") + + # RAG deletion failed, so count is 0 + assert result["deleted_rag_entries"] == 0 + # But DocumentDB record should still be deleted + assert result["deleted_documentdb_records"] == 1 + + # Verify record removed from database + async with get_session() as session: + q = select(DocumentDB).where(DocumentDB.doc_hash == "sha256-ragfail") + db_result = await session.exec(q) + assert db_result.first() is None + + +@pytest.mark.asyncio +async def test_delete_from_rag_by_hash_no_rag_id(db, mock_settings, mock_step_config): + """Test deletion of a DocumentDB record with no rag_id (skip RAG deletion).""" + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + # Create a record without rag_id + async with get_session() as session: + import datetime + + record = DocumentDB( + doc_hash="sha256-noragid", + source="test_source", + db_name="test-project", + lancedb_dir="/tmp/lancedb", + rag_id=None, # No rag_id + chunk_count=5, + created_date=datetime.datetime.now(datetime.UTC), + ) + session.add(record) + await session.commit() + + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + result = await rag.delete_from_rag_by_hash("sha256-noragid") + + # No RAG deletion since no rag_id + assert result["deleted_rag_entries"] == 0 + # But DocumentDB record should be deleted + assert result["deleted_documentdb_records"] == 1 + + +@pytest.mark.asyncio +async def test_delete_document_uri_includes_rag_deletion(db, mock_settings, mock_step_config): + """Test that delete_document_uri_by_uri includes RAG deletion in statistics.""" + # Create test document + batch_id = await doc_ops.new_batch("test_source", "Test Batch") + test_uri = "/tmp/rag_delete_test.pdf" + test_bytes = b"test bytes for rag deletion" + doc_uri, doc = await doc_ops.create_document_from_uri( + test_uri, "test_source", "application/pdf", test_bytes, batch_id=batch_id + ) + + # Create a DocumentDB record for this document + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + await rag.create_document_db_record( + doc_hash=doc.hash, + source="test_source", + step_config=mock_step_config, + rag_id="test-rag-id", + chunk_count=10, + ) + + # Mock HaikuRAG client for deletion + mock_client = MagicMock() + mock_client.delete_document = AsyncMock() + + with ( + patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings), + patch("soliplex.ingester.lib.rag.HaikuRAG") as mock_haiku_rag, + patch("soliplex.ingester.lib.rag.build_embed_config"), + patch("soliplex.ingester.lib.rag.build_storage_config"), + ): + mock_haiku_rag.return_value.__aenter__ = AsyncMock(return_value=mock_client) + mock_haiku_rag.return_value.__aexit__ = AsyncMock(return_value=None) + + result = await doc_ops.delete_document_uri_by_uri(test_uri, "test_source") + + # Verify RAG deletion statistics are included + assert "deleted_rag_entries" in result + assert "deleted_documentdb_records" in result + assert result["deleted_rag_entries"] == 1 + assert result["deleted_documentdb_records"] == 1 + + # Verify RAG deletion was called + mock_client.delete_document.assert_called_once_with("test-rag-id") + + # Verify total includes RAG counts + expected_total = ( + result["deleted_document_uris"] + + result["deleted_uri_history"] + + result["deleted_documents"] + + result["deleted_workflow_runs"] + + result["deleted_run_steps"] + + result["deleted_lifecycle_history"] + + result["deleted_rag_entries"] + + result["deleted_documentdb_records"] + ) + assert result["total_deleted"] == expected_total + + +@pytest.mark.asyncio +async def test_delete_document_uri_multiple_references_no_rag_deletion(db, mock_settings, mock_step_config): + """Test that RAG entries are NOT deleted when multiple URIs reference the document.""" + # Create test data with two URIs pointing to the same document + batch_id = await doc_ops.new_batch("test_source", "Test Batch") + test_bytes = b"shared document bytes" + + # Create first URI + test_uri1 = "/tmp/multi_rag_test1.pdf" + doc_uri1, doc = await doc_ops.create_document_from_uri( + test_uri1, "test_source", "application/pdf", test_bytes, batch_id=batch_id + ) + + # Create second URI pointing to the same document (same bytes = same hash) + test_uri2 = "/tmp/multi_rag_test2.pdf" + doc_uri2, doc2 = await doc_ops.create_document_from_uri( + test_uri2, "test_source", "application/pdf", test_bytes, batch_id=batch_id + ) + + # Create a DocumentDB record + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + await rag.create_document_db_record( + doc_hash=doc.hash, + source="test_source", + step_config=mock_step_config, + rag_id="shared-rag-id", + chunk_count=10, + ) + + # Delete only the first URI + result = await doc_ops.delete_document_uri_by_uri(test_uri1, "test_source") + + # RAG entries should NOT be deleted since document still has another reference + assert result["deleted_rag_entries"] == 0 + assert result["deleted_documentdb_records"] == 0 + assert result["deleted_documents"] == 0 + + # Verify DocumentDB record still exists + async with get_session() as session: + q = select(DocumentDB).where(DocumentDB.doc_hash == doc.hash) + db_result = await session.exec(q) + assert db_result.first() is not None + + +@pytest.mark.asyncio +async def test_delete_orphaned_documents_includes_rag_deletion(db, mock_settings, mock_step_config): + """Test that delete_orphaned_documents includes RAG deletion in statistics.""" + # Create a document without a URI (orphaned) + # We need to directly insert a Document record + from soliplex.ingester.lib.models import Document + + doc_hash = "sha256-orphaned-test" + async with get_session() as session: + doc = Document(hash=doc_hash, mime_type="application/pdf", file_size=100) + session.add(doc) + await session.commit() + + # Create a DocumentDB record for this orphaned document + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + await rag.create_document_db_record( + doc_hash=doc_hash, + source="test_source", + step_config=mock_step_config, + rag_id="orphaned-rag-id", + chunk_count=5, + ) + + # Mock HaikuRAG client for deletion + mock_client = MagicMock() + mock_client.delete_document = AsyncMock() + + with ( + patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings), + patch("soliplex.ingester.lib.rag.HaikuRAG") as mock_haiku_rag, + patch("soliplex.ingester.lib.rag.build_embed_config"), + patch("soliplex.ingester.lib.rag.build_storage_config"), + ): + mock_haiku_rag.return_value.__aenter__ = AsyncMock(return_value=mock_client) + mock_haiku_rag.return_value.__aexit__ = AsyncMock(return_value=None) + + result = await doc_ops.delete_orphaned_documents() + + # Verify RAG deletion statistics are included + assert "deleted_rag_entries" in result + assert "deleted_documentdb_records" in result + assert result["deleted_rag_entries"] == 1 + assert result["deleted_documentdb_records"] == 1 + assert result["deleted_documents"] == 1 + + # Verify RAG deletion was called + mock_client.delete_document.assert_called_once_with("orphaned-rag-id") + + +@pytest.mark.asyncio +async def test_list_documentdb_databases_empty(db): + """Test list_documentdb_databases returns empty list when no records exist.""" + result = await doc_ops.list_documentdb_databases() + assert result == [] + + +@pytest.mark.asyncio +async def test_list_documentdb_databases_with_records(db, mock_settings): + """Test list_documentdb_databases returns correct counts per database.""" + step_config1 = MagicMock(spec=StepConfig) + step_config1.config_json = {"data_dir": "project-a"} + + step_config2 = MagicMock(spec=StepConfig) + step_config2.config_json = {"data_dir": "project-b"} + + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + # Create 2 documents in project-a + await rag.create_document_db_record( + doc_hash="sha256-doc1", + source="source1", + step_config=step_config1, + rag_id="rag-1", + chunk_count=10, + ) + await rag.create_document_db_record( + doc_hash="sha256-doc2", + source="source2", + step_config=step_config1, + rag_id="rag-2", + chunk_count=20, + ) + + # Create 1 document in project-b + await rag.create_document_db_record( + doc_hash="sha256-doc3", + source="source3", + step_config=step_config2, + rag_id="rag-3", + chunk_count=15, + ) + + result = await doc_ops.list_documentdb_databases() + + assert len(result) == 2 + + # Find project-a + project_a = next((r for r in result if r["db_name"] == "project-a"), None) + assert project_a is not None + assert project_a["lancedb_dir"] == "/tmp/lancedb" + assert project_a["document_count"] == 2 + assert project_a["total_chunks"] == 30 + + # Find project-b + project_b = next((r for r in result if r["db_name"] == "project-b"), None) + assert project_b is not None + assert project_b["lancedb_dir"] == "/tmp/lancedb" + assert project_b["document_count"] == 1 + assert project_b["total_chunks"] == 15 + + +@pytest.mark.asyncio +async def test_list_documents_in_rag_db_empty(db, mock_settings): + """Test list_documents_in_rag_db returns empty list when no matching records.""" + with patch("soliplex.ingester.lib.operations.get_settings", return_value=mock_settings): + result = await doc_ops.list_documents_in_rag_db("nonexistent-db") + assert result == [] + + +@pytest.mark.asyncio +async def test_list_documents_in_rag_db_with_records(db, mock_settings): + """Test list_documents_in_rag_db returns correct document information.""" + from soliplex.ingester.lib.models import Document + from soliplex.ingester.lib.models import DocumentURI + + step_config = MagicMock(spec=StepConfig) + step_config.config_json = {"data_dir": "test-project"} + + # Create Document and DocumentURI records + async with get_session() as session: + doc = Document(hash="sha256-listdocs1", mime_type="application/pdf", file_size=12345) + session.add(doc) + await session.flush() + + doc_uri = DocumentURI( + doc_hash="sha256-listdocs1", + uri="/path/to/doc1.pdf", + source="test_source", + version=1, + ) + session.add(doc_uri) + + # Create DocumentDB record + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + await rag.create_document_db_record( + doc_hash="sha256-listdocs1", + source="test_source", + step_config=step_config, + rag_id="rag-listdocs-1", + chunk_count=25, + ) + + with patch("soliplex.ingester.lib.operations.get_settings", return_value=mock_settings): + result = await doc_ops.list_documents_in_rag_db("test-project") + + assert len(result) == 1 + doc_info = result[0] + assert doc_info["doc_hash"] == "sha256-listdocs1" + assert doc_info["rag_id"] == "rag-listdocs-1" + assert doc_info["chunk_count"] == 25 + assert doc_info["source"] == "test_source" + assert doc_info["uri"] == "/path/to/doc1.pdf" + assert doc_info["mime_type"] == "application/pdf" + assert doc_info["file_size"] == 12345 + assert doc_info["created_date"] is not None + + +@pytest.mark.asyncio +async def test_list_documents_in_rag_db_filters_by_lancedb_dir(db, mock_settings): + """Test list_documents_in_rag_db correctly filters by lancedb_dir.""" + step_config = MagicMock(spec=StepConfig) + step_config.config_json = {"data_dir": "shared-db"} + + # Create record with default lancedb_dir + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + await rag.create_document_db_record( + doc_hash="sha256-default-dir", + source="source1", + step_config=step_config, + rag_id="rag-default", + chunk_count=10, + ) + + # Create record with different lancedb_dir + mock_settings_s3 = MagicMock() + mock_settings_s3.lancedb_dir = "s3://bucket/lancedb" + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings_s3): + await rag.create_document_db_record( + doc_hash="sha256-s3-dir", + source="source2", + step_config=step_config, + rag_id="rag-s3", + chunk_count=20, + ) + + # Query with default lancedb_dir + with patch("soliplex.ingester.lib.operations.get_settings", return_value=mock_settings): + result_default = await doc_ops.list_documents_in_rag_db("shared-db") + + assert len(result_default) == 1 + assert result_default[0]["doc_hash"] == "sha256-default-dir" + + # Query with explicit S3 lancedb_dir + with patch("soliplex.ingester.lib.operations.get_settings", return_value=mock_settings): + result_s3 = await doc_ops.list_documents_in_rag_db("shared-db", lancedb_dir="s3://bucket/lancedb") + + assert len(result_s3) == 1 + assert result_s3[0]["doc_hash"] == "sha256-s3-dir" + + +@pytest.mark.asyncio +async def test_check_rag_db_consistency_lancedb_not_accessible(db, mock_settings): + """Test check_rag_db_consistency handles inaccessible LanceDB gracefully.""" + step_config = MagicMock(spec=StepConfig) + step_config.config_json = {"data_dir": "test-check-db"} + + # Create a DocumentDB record + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + await rag.create_document_db_record( + doc_hash="sha256-check1", + source="test_source", + step_config=step_config, + rag_id="rag-check-1", + chunk_count=10, + ) + + # Mock HaikuRAG to raise an error (database doesn't exist) + with ( + patch("soliplex.ingester.lib.operations.get_settings", return_value=mock_settings), + patch("haiku.rag.client.HaikuRAG") as mock_haiku_rag, + patch("haiku.rag.config.get_config"), + ): + mock_haiku_rag.return_value.__aenter__ = AsyncMock(side_effect=Exception("Database not found")) + + result = await doc_ops.check_rag_db_consistency("test-check-db") + + # Should return error but still list DocumentDB records + assert "error" in result + assert result["documentdb_count"] == 1 + assert result["lancedb_count"] == 0 + assert len(result["in_documentdb_only"]) == 1 + assert result["in_documentdb_only"][0]["rag_id"] == "rag-check-1" + + +@pytest.mark.asyncio +async def test_check_rag_db_consistency_all_matched(db, mock_settings): + """Test check_rag_db_consistency when all documents match.""" + step_config = MagicMock(spec=StepConfig) + step_config.config_json = {"data_dir": "matched-db"} + + # Create DocumentDB records + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + await rag.create_document_db_record( + doc_hash="sha256-matched1", + source="test_source", + step_config=step_config, + rag_id="rag-matched-1", + chunk_count=10, + ) + await rag.create_document_db_record( + doc_hash="sha256-matched2", + source="test_source", + step_config=step_config, + rag_id="rag-matched-2", + chunk_count=20, + ) + + # Mock HaikuRAG to return matching documents + mock_doc1 = MagicMock() + mock_doc1.id = "rag-matched-1" + mock_doc1.uri = "/path/to/doc1.pdf" + mock_doc1.title = "Document 1" + mock_doc1.created_at = None + mock_doc1.chunk_count = 10 + mock_doc1.metadata = {} + + mock_doc2 = MagicMock() + mock_doc2.id = "rag-matched-2" + mock_doc2.uri = "/path/to/doc2.pdf" + mock_doc2.title = "Document 2" + mock_doc2.created_at = None + mock_doc2.chunk_count = 20 + mock_doc2.metadata = {} + + mock_client = MagicMock() + mock_client.list_documents = AsyncMock(return_value=[mock_doc1, mock_doc2]) + + with ( + patch("soliplex.ingester.lib.operations.get_settings", return_value=mock_settings), + patch("haiku.rag.client.HaikuRAG") as mock_haiku_rag, + patch("haiku.rag.config.get_config"), + ): + mock_haiku_rag.return_value.__aenter__ = AsyncMock(return_value=mock_client) + mock_haiku_rag.return_value.__aexit__ = AsyncMock(return_value=None) + + result = await doc_ops.check_rag_db_consistency("matched-db") + + assert "error" not in result + assert result["documentdb_count"] == 2 + assert result["lancedb_count"] == 2 + assert result["matched"] == 2 + assert len(result["in_documentdb_only"]) == 0 + assert len(result["in_lancedb_only"]) == 0 + + +@pytest.mark.asyncio +async def test_check_rag_db_consistency_with_discrepancies(db, mock_settings): + """Test check_rag_db_consistency detects discrepancies in both directions.""" + step_config = MagicMock(spec=StepConfig) + step_config.config_json = {"data_dir": "discrepancy-db"} + + # Create DocumentDB records + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + # This one will match + await rag.create_document_db_record( + doc_hash="sha256-common", + source="test_source", + step_config=step_config, + rag_id="rag-common", + chunk_count=10, + ) + # This one is only in DocumentDB + await rag.create_document_db_record( + doc_hash="sha256-docdb-only", + source="test_source", + step_config=step_config, + rag_id="rag-docdb-only", + chunk_count=15, + ) + + # Mock HaikuRAG to return different documents + mock_doc_common = MagicMock() + mock_doc_common.id = "rag-common" + mock_doc_common.uri = "/path/to/common.pdf" + mock_doc_common.title = "Common Document" + mock_doc_common.created_at = None + mock_doc_common.chunk_count = 10 + mock_doc_common.metadata = {} + + mock_doc_lance_only = MagicMock() + mock_doc_lance_only.id = "rag-lance-only" + mock_doc_lance_only.uri = "/path/to/lance-only.pdf" + mock_doc_lance_only.title = "LanceDB Only Document" + mock_doc_lance_only.created_at = None + mock_doc_lance_only.chunk_count = 25 + mock_doc_lance_only.metadata = {"extra": "data"} + + mock_client = MagicMock() + mock_client.list_documents = AsyncMock(return_value=[mock_doc_common, mock_doc_lance_only]) + + with ( + patch("soliplex.ingester.lib.operations.get_settings", return_value=mock_settings), + patch("haiku.rag.client.HaikuRAG") as mock_haiku_rag, + patch("haiku.rag.config.get_config"), + ): + mock_haiku_rag.return_value.__aenter__ = AsyncMock(return_value=mock_client) + mock_haiku_rag.return_value.__aexit__ = AsyncMock(return_value=None) + + result = await doc_ops.check_rag_db_consistency("discrepancy-db") + + assert "error" not in result + assert result["documentdb_count"] == 2 + assert result["lancedb_count"] == 2 + assert result["matched"] == 1 + + # Check in_documentdb_only + assert len(result["in_documentdb_only"]) == 1 + assert result["in_documentdb_only"][0]["rag_id"] == "rag-docdb-only" + assert result["in_documentdb_only"][0]["doc_hash"] == "sha256-docdb-only" + + # Check in_lancedb_only + assert len(result["in_lancedb_only"]) == 1 + assert result["in_lancedb_only"][0]["rag_id"] == "rag-lance-only" + assert result["in_lancedb_only"][0]["uri"] == "/path/to/lance-only.pdf" + + +@pytest.fixture +def temp_lancedb_dir(): + """Create a temporary directory for LanceDB testing.""" + with tempfile.TemporaryDirectory() as tmpdir: + yield tmpdir + + +def create_embedded_chunk(content: str, order: int, vector_dim: int = 2560): + """Create a chunk with pre-computed fake embeddings to avoid calling embedding service. + + Note: HaikuRAG requires at least 2560 dimensions for the vector embedding. + """ + import hashlib + + from haiku.rag.store.models.chunk import Chunk + + # Create a deterministic fake embedding based on content hash + hash_bytes = hashlib.sha256(content.encode()).digest() + # Generate fake embedding from hash (normalized to small values) + # Start with values derived from the hash + fake_embedding = [(b / 255.0 - 0.5) * 0.1 for b in hash_bytes] + # Pad to full dimension by repeating the pattern + while len(fake_embedding) < vector_dim: + fake_embedding.extend(fake_embedding[: vector_dim - len(fake_embedding)]) + fake_embedding = fake_embedding[:vector_dim] + + return Chunk( + content=content, + metadata={}, + order=order, + embedding=fake_embedding, + ) + + +@pytest.mark.asyncio +async def test_check_db_with_real_lancedb(db, temp_lancedb_dir): + """ + Integration test that creates actual LanceDB entries via HaikuRAG directly + and verifies check_rag_db_consistency works with real data. + """ + from docling_core.types.doc.document import DoclingDocument + from haiku.rag.client import HaikuRAG + from haiku.rag.config import get_config + + db_name = "test-integration-db" + + # Create mock settings pointing to temp directory + mock_settings = MagicMock() + mock_settings.lancedb_dir = temp_lancedb_dir + + # Create step_config + step_config = MagicMock(spec=StepConfig) + step_config.config_json = {"data_dir": db_name} + + # Create minimal docling document + docling_doc = DoclingDocument(name="test-doc") + + db_path = Path(temp_lancedb_dir) / db_name + config = get_config() + + # Create two test documents directly in LanceDB with pre-embedded chunks + rag_ids = [] + test_doc_hashes = ["sha256-integration-test-0", "sha256-integration-test-1"] + + async with HaikuRAG(db_path=db_path, config=config, create=True) as client: + for i, doc_hash in enumerate(test_doc_hashes): + # Create chunks with pre-computed embeddings + chunks = [create_embedded_chunk(f"Chunk {j} of doc {i}", j) for j in range(3)] + + doc = await client.import_document( + chunks=chunks, + title=f"Test Document {i}", + uri=f"/path/to/test-doc-{i}.pdf", + metadata={"doc_id": doc_hash, "md5": f"md5-test-{i}", "source": "integration_test"}, + docling_document=docling_doc, + ) + rag_ids.append(doc.id) + + # Create corresponding DocumentDB records + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + for doc_hash, rag_id in zip(test_doc_hashes, rag_ids, strict=True): + await rag.create_document_db_record( + doc_hash=doc_hash, + source="integration_test", + step_config=step_config, + rag_id=rag_id, + chunk_count=3, + ) + + # Verify LanceDB files were created + assert db_path.exists(), f"LanceDB path {db_path} should exist" + + # Now test check_rag_db_consistency - should show all matched + with patch("soliplex.ingester.lib.operations.get_settings", return_value=mock_settings): + result = await doc_ops.check_rag_db_consistency(db_name, lancedb_dir=temp_lancedb_dir) + + # Verify results + assert "error" not in result, f"Unexpected error: {result.get('error')}" + assert result["db_name"] == db_name + assert result["lancedb_dir"] == temp_lancedb_dir + assert result["documentdb_count"] == 2 + assert result["lancedb_count"] == 2 + assert result["matched"] == 2 + assert len(result["in_documentdb_only"]) == 0, f"Unexpected in_documentdb_only: {result['in_documentdb_only']}" + assert len(result["in_lancedb_only"]) == 0, f"Unexpected in_lancedb_only: {result['in_lancedb_only']}" + + +@pytest.mark.asyncio +async def test_check_db_detects_missing_lancedb_entry(db, temp_lancedb_dir): + """ + Test that check_db detects when a DocumentDB record exists but LanceDB entry is missing. + This simulates a case where the LanceDB was corrupted or manually modified. + """ + from docling_core.types.doc.document import DoclingDocument + from haiku.rag.client import HaikuRAG + from haiku.rag.config import get_config + + db_name = "test-missing-lance-db" + + mock_settings = MagicMock() + mock_settings.lancedb_dir = temp_lancedb_dir + + step_config = MagicMock(spec=StepConfig) + step_config.config_json = {"data_dir": db_name} + + docling_doc = DoclingDocument(name="test-doc") + db_path = Path(temp_lancedb_dir) / db_name + config = get_config() + + # Create one document in LanceDB + async with HaikuRAG(db_path=db_path, config=config, create=True) as client: + chunks = [create_embedded_chunk("Test chunk content", 0)] + doc = await client.import_document( + chunks=chunks, + title="Saved Document", + uri="/path/to/saved.pdf", + metadata={"doc_id": "sha256-saved-to-lance", "md5": "md5-saved", "source": "test_source"}, + docling_document=docling_doc, + ) + rag_id1 = doc.id + + # Create DocumentDB records - one for the saved doc, one for a non-existent doc + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + # Create DocumentDB record for saved document + await rag.create_document_db_record( + doc_hash="sha256-saved-to-lance", + source="test_source", + step_config=step_config, + rag_id=rag_id1, + chunk_count=1, + ) + + # Create DocumentDB record for a document that was NEVER saved to LanceDB + # (simulating a failed save or data inconsistency) + await rag.create_document_db_record( + doc_hash="sha256-never-saved", + source="test_source", + step_config=step_config, + rag_id="fake-rag-id-never-existed", + chunk_count=5, + ) + + # Check consistency - should detect the missing LanceDB entry + with patch("soliplex.ingester.lib.operations.get_settings", return_value=mock_settings): + result = await doc_ops.check_rag_db_consistency(db_name, lancedb_dir=temp_lancedb_dir) + + assert "error" not in result + assert result["documentdb_count"] == 2 + assert result["lancedb_count"] == 1 + assert result["matched"] == 1 + + # Should have one document in DocumentDB but not in LanceDB + assert len(result["in_documentdb_only"]) == 1 + assert result["in_documentdb_only"][0]["rag_id"] == "fake-rag-id-never-existed" + assert result["in_documentdb_only"][0]["doc_hash"] == "sha256-never-saved" + + # No orphaned LanceDB entries + assert len(result["in_lancedb_only"]) == 0 + + +@pytest.mark.asyncio +async def test_check_db_detects_orphaned_lancedb_entry(db, temp_lancedb_dir): + """ + Test that check_db detects when a LanceDB entry exists but DocumentDB record is missing. + This simulates a case where DocumentDB was cleared but LanceDB wasn't. + """ + from docling_core.types.doc.document import DoclingDocument + from haiku.rag.client import HaikuRAG + from haiku.rag.config import get_config + + db_name = "test-orphaned-lance-db" + + mock_settings = MagicMock() + mock_settings.lancedb_dir = temp_lancedb_dir + + step_config = MagicMock(spec=StepConfig) + step_config.config_json = {"data_dir": db_name} + + docling_doc = DoclingDocument(name="test-doc") + db_path = Path(temp_lancedb_dir) / db_name + config = get_config() + + # Create two documents in LanceDB + async with HaikuRAG(db_path=db_path, config=config, create=True) as client: + # First document - will be tracked + chunks1 = [create_embedded_chunk("Tracked chunk content", 0)] + doc1 = await client.import_document( + chunks=chunks1, + title="Tracked Document", + uri="/path/to/tracked.pdf", + metadata={"doc_id": "sha256-tracked", "md5": "md5-tracked", "source": "test_source"}, + docling_document=docling_doc, + ) + rag_id1 = doc1.id + + # Second document - will be orphaned (no DocumentDB record) + chunks2 = [create_embedded_chunk("Orphaned chunk content", 0)] + doc2 = await client.import_document( + chunks=chunks2, + title="Orphaned Document", + uri="/path/to/orphaned.pdf", + metadata={"doc_id": "sha256-orphaned", "source": "unknown"}, + docling_document=docling_doc, + ) + orphan_rag_id = doc2.id + + # Only create DocumentDB record for the first document + with patch("soliplex.ingester.lib.rag.get_settings", return_value=mock_settings): + await rag.create_document_db_record( + doc_hash="sha256-tracked", + source="test_source", + step_config=step_config, + rag_id=rag_id1, + chunk_count=1, + ) + + # Check consistency - should detect the orphaned LanceDB entry + with patch("soliplex.ingester.lib.operations.get_settings", return_value=mock_settings): + result = await doc_ops.check_rag_db_consistency(db_name, lancedb_dir=temp_lancedb_dir) + + assert "error" not in result + assert result["documentdb_count"] == 1 + assert result["lancedb_count"] == 2 + assert result["matched"] == 1 + + # No missing LanceDB entries + assert len(result["in_documentdb_only"]) == 0 + + # Should have one orphaned LanceDB entry + assert len(result["in_lancedb_only"]) == 1 + assert result["in_lancedb_only"][0]["rag_id"] == orphan_rag_id + assert result["in_lancedb_only"][0]["uri"] == "/path/to/orphaned.pdf" diff --git a/tests/unit/test_endpoint_coverage.py b/tests/unit/test_endpoint_coverage.py new file mode 100644 index 0000000..ab3833f --- /dev/null +++ b/tests/unit/test_endpoint_coverage.py @@ -0,0 +1,219 @@ +""" +Unit tests for newly added and fixed endpoints. + +Tests cover: +- GET /api/v1/stats/durations (bug fix verification) +- POST /api/v1/document/cleanup-orphans (new endpoint) +- GET /api/v1/batch/{batch_id}/steps (new endpoint) +""" + +from unittest.mock import patch + +import pytest +from fastapi.testclient import TestClient + +from soliplex.ingester.lib import operations as doc_ops +from soliplex.ingester.lib.wf import operations as wf_ops +from soliplex.ingester.server import app + + +@pytest.fixture +def client(): + """Create test client with auth disabled.""" + with patch("soliplex.ingester.lib.auth.get_settings") as mock_settings: + settings = mock_settings.return_value + settings.auth_enabled = False + yield TestClient(app) + + +class TestGetRunGroupDurationsBugFix: + """Tests for GET /api/v1/stats/durations bug fix.""" + + def test_passes_correct_parameter(self, client): + """Test that endpoint passes run_group_id instead of status.""" + with patch("soliplex.ingester.server.routes.stats.wf_ops.get_run_group_durations") as mock_func: + mock_func.return_value = [{"doc_id": "test", "duration": 100}] + + response = client.get("/api/v1/stats/durations?run_group_id=123") + + # Verify function was called with correct parameter + mock_func.assert_called_once_with(123) + assert response.status_code == 200 + + def test_error_handling(self, client): + """Test error handling when function raises exception.""" + with patch("soliplex.ingester.server.routes.stats.wf_ops.get_run_group_durations") as mock_func: + mock_func.side_effect = RuntimeError("PostgreSQL required") + + response = client.get("/api/v1/stats/durations?run_group_id=123") + + assert response.status_code == 500 + assert "error" in response.json() + + +class TestCleanupOrphansEndpoint: + """Tests for POST /api/v1/document/cleanup-orphans.""" + + def test_cleanup_orphans_success(self, client): + """Test successful cleanup of orphaned documents.""" + with patch("soliplex.ingester.server.routes.document.operations.delete_orphaned_documents") as mock_func: + mock_func.return_value = {"deleted_documents": 5, "deleted_history": 3} + + response = client.post("/api/v1/document/cleanup-orphans") + + assert response.status_code == 200 + data = response.json() + assert data["message"] == "Orphaned documents cleaned up" + assert data["statistics"]["deleted_documents"] == 5 + assert data["statistics"]["deleted_history"] == 3 + mock_func.assert_called_once() + + def test_cleanup_orphans_no_orphans(self, client): + """Test cleanup when no orphans exist.""" + with patch("soliplex.ingester.server.routes.document.operations.delete_orphaned_documents") as mock_func: + mock_func.return_value = {"deleted_documents": 0, "deleted_history": 0} + + response = client.post("/api/v1/document/cleanup-orphans") + + assert response.status_code == 200 + data = response.json() + assert data["statistics"]["deleted_documents"] == 0 + assert data["statistics"]["deleted_history"] == 0 + + def test_cleanup_orphans_error(self, client): + """Test error handling during cleanup.""" + with patch("soliplex.ingester.server.routes.document.operations.delete_orphaned_documents") as mock_func: + mock_func.side_effect = Exception("Database error") + + response = client.post("/api/v1/document/cleanup-orphans") + + assert response.status_code == 500 + assert "error" in response.json() + + +class TestGetBatchStepsEndpoint: + """Tests for GET /api/v1/batch/{batch_id}/steps.""" + + def test_get_batch_steps_success(self, client): + """Test retrieving steps for a batch.""" + from soliplex.ingester.lib.models import RunStatus + from soliplex.ingester.lib.models import RunStep + from soliplex.ingester.lib.models import WorkflowStepType + + # Mock step data + mock_step = RunStep( + id=1, + workflow_run_id=1, + step_config_id=1, + step_type=WorkflowStepType.PARSE, + workflow_step_number=1, + workflow_step_name="parse", + status=RunStatus.COMPLETED, + priority=0, + retry=0, + retries=3, + ) + + with patch("soliplex.ingester.server.routes.batch.wf_ops.get_steps_for_batch") as mock_func: + mock_func.return_value = [mock_step] + + response = client.get("/api/v1/batch/123/steps") + + assert response.status_code == 200 + data = response.json() + assert isinstance(data, list) + assert len(data) == 1 + assert data[0]["id"] == 1 + assert data[0]["status"] == "COMPLETED" + mock_func.assert_called_once_with(123) + + def test_get_batch_steps_empty(self, client): + """Test retrieving steps when batch has no workflows.""" + with patch("soliplex.ingester.server.routes.batch.wf_ops.get_steps_for_batch") as mock_func: + mock_func.return_value = [] + + response = client.get("/api/v1/batch/999/steps") + + assert response.status_code == 200 + data = response.json() + assert isinstance(data, list) + assert len(data) == 0 + + def test_get_batch_steps_error(self, client): + """Test error handling when getting batch steps.""" + with patch("soliplex.ingester.server.routes.batch.wf_ops.get_steps_for_batch") as mock_func: + mock_func.side_effect = Exception("Database error") + + response = client.get("/api/v1/batch/123/steps") + + assert response.status_code == 500 + assert "error" in response.json() + + +class TestEndpointIntegration: + """Integration tests with database.""" + + @pytest.mark.asyncio + async def test_cleanup_orphans_integration(self, db, client): + """Test cleanup-orphans endpoint with real database.""" + from soliplex.ingester.lib.models import get_session + + # Create orphaned document + orphan_hash = "sha256-orphan-endpoint-test" + async with get_session() as session: + from soliplex.ingester.lib.models import Document + + orphan_doc = Document( + hash=orphan_hash, + mime_type="application/pdf", + file_size=100, + ) + session.add(orphan_doc) + await session.commit() + + # Call endpoint + response = client.post("/api/v1/document/cleanup-orphans") + + # Verify response + assert response.status_code == 200 + data = response.json() + assert data["statistics"]["deleted_documents"] == 1 + assert data["statistics"]["deleted_history"] == 0 + + # Verify document was deleted + async with get_session() as session: + from sqlmodel import select + + from soliplex.ingester.lib.models import Document + + q = select(Document).where(Document.hash == orphan_hash) + result = await session.exec(q) + doc = result.first() + assert doc is None + + @pytest.mark.asyncio + async def test_get_batch_steps_integration(self, db, client): + """Test get-batch-steps endpoint with real database.""" + # Create batch and workflow + batch_id = await doc_ops.new_batch("test_source", "Test Batch") + uri, doc = await doc_ops.create_document_from_uri( + "/tmp/test.pdf", "test_source", "application/pdf", b"content", batch_id=batch_id + ) + + run_group = await wf_ops.create_run_group(workflow_definition_id="batch", batch_id=batch_id, param_id="test_base") + workflow_run, steps = await wf_ops.create_workflow_run(run_group=run_group, doc_id=doc.hash) + + # Call endpoint + response = client.get(f"/api/v1/batch/{batch_id}/steps") + + # Verify response + assert response.status_code == 200 + data = response.json() + assert isinstance(data, list) + assert len(data) > 0 + + # Verify step data + step = data[0] + assert "id" in step + assert "status" in step + assert "step_type" in step diff --git a/uv.lock b/uv.lock index eb15537..ed59462 100644 --- a/uv.lock +++ b/uv.lock @@ -4025,7 +4025,6 @@ dependencies = [ { name = "alembic" }, { name = "async-lru" }, { name = "cryptography" }, - { name = "docling-core" }, { name = "fastapi", extra = ["standard"] }, { name = "fsspec" }, { name = "haiku-rag-slim" }, @@ -4067,7 +4066,6 @@ requires-dist = [ { name = "alembic", specifier = ">=1.17.2" }, { name = "async-lru", specifier = ">=2.0.5" }, { name = "cryptography", specifier = ">=44.0.0" }, - { name = "docling-core", specifier = ">=2.51.1" }, { name = "fastapi", extras = ["standard"], specifier = ">=0.125.0" }, { name = "fsspec", specifier = ">=2025.12.0" }, { name = "haiku-rag-slim", specifier = ">=0.30.2" },