"""End-to-end ingest pipeline: parse text -> chunk -> embed -> persist to Postgres. The script reads ``data/seed_sources/manifest.json`` (or a custom path) and, for each listed source, requires a pre-registered row in ``sources``/``source_versions`` (so governance state remains the source of truth). It then: 1. Chunks the source text via the hardened ScientificChunker (sentence-aware sub-splitting, configurable overlap, tiny-fragment merging). 2. Optionally embeds each chunk via the Ollama embedding model from .env. 3. Replaces existing rows for ``(source_id, version_id)`` in ``chunks`` and ``chunk_embeddings`` so the script is idempotent. Usage: python3 scripts/run_ingest_pipeline.py \ --manifest data/seed_sources/manifest.json \ --database-url postgresql+psycopg://mobcoderid-296@localhost/ai_knowledge_spine Skip embeddings (chunks only, much faster): python3 scripts/run_ingest_pipeline.py --no-embed """ from __future__ import annotations import argparse import json import os import sys import uuid from dataclasses import dataclass from datetime import UTC, datetime from pathlib import Path from typing import Any REPO_ROOT = Path(__file__).resolve().parent.parent CHUNKING_ROOT = REPO_ROOT / "services" / "chunking-service" SHARED_ROOT = REPO_ROOT / "shared" # Sources used by integration tests — ingest must not replace these unless --force-fixtures. PROTECTED_SOURCE_IDS = frozenset( { "LBL-NSCLC-RET-EMA-2026", "DOC-CSR-NSCLC-RET-2026", "SOP-MED-NSCLC-RET-2026", "LBL-NSCLC-TEST-EMA-2026", "DOC-CSR-NSCLC-TEST-2026", } ) for path in (str(CHUNKING_ROOT), str(SHARED_ROOT)): if path not in sys.path: sys.path.insert(0, path) # Load .env so AKS_DATABASE_URL / OLLAMA_* are visible. _env_file = REPO_ROOT / ".env" if _env_file.exists(): for raw_line in _env_file.read_text().splitlines(): line = raw_line.strip() if not line or line.startswith("#") or "=" not in line: continue key, value = line.split("=", 1) os.environ.setdefault(key.strip(), value.strip()) from app.schemas.chunking import ChunkPreviewRequest # noqa: E402 from app.services.heuristics import ScientificChunker # noqa: E402 import psycopg # noqa: E402 @dataclass class SourceManifestEntry: source_id: str version_id: str source_class: str therapy_area: str geography: str audience: list[str] text: str def _load_manifest(manifest_path: Path) -> list[SourceManifestEntry]: data = json.loads(manifest_path.read_text()) base_dir = manifest_path.parent entries: list[SourceManifestEntry] = [] for raw in data.get("sources", []): text_file = raw.get("text_file") if text_file: text = (base_dir / text_file).read_text() elif "text" in raw: text = raw["text"] else: raise ValueError(f"Manifest entry {raw.get('source_id')!r} has no text_file or text") entries.append( SourceManifestEntry( source_id=raw["source_id"], version_id=raw["version_id"], source_class=raw["source_class"], therapy_area=raw["therapy_area"], geography=raw["geography"], audience=list(raw.get("audience", ["HCP", "Internal"])), text=text, ) ) return entries def _dsn_from_sqlalchemy_url(url: str) -> str: # Accept either SQLAlchemy URL (postgresql+psycopg://user@host/db) or raw libpq DSN. if url.startswith("postgresql+psycopg://"): return url.replace("postgresql+psycopg://", "postgresql://", 1) return url def _ensure_source_registered(cursor: "psycopg.Cursor", entry: SourceManifestEntry) -> None: cursor.execute( "SELECT 1 FROM sources WHERE source_id = %s", (entry.source_id,), ) if cursor.fetchone() is None: raise RuntimeError( f"Source {entry.source_id!r} is not registered in `sources`. " f"Register the source and version via ingestion-service (or seed scripts) before running ingest." ) cursor.execute( "SELECT 1 FROM source_versions WHERE source_id = %s AND version_id = %s", (entry.source_id, entry.version_id), ) if cursor.fetchone() is None: raise RuntimeError( f"Source version {entry.version_id!r} for {entry.source_id!r} is not registered in `source_versions`." ) def _delete_existing_chunks(cursor: "psycopg.Cursor", source_id: str, version_id: str) -> int: cursor.execute( "DELETE FROM chunks WHERE source_id = %s AND version_id = %s", (source_id, version_id), ) return cursor.rowcount or 0 def _insert_chunk(cursor: "psycopg.Cursor", row: dict[str, Any]) -> None: cursor.execute( """ INSERT INTO chunks ( chunk_id, source_id, version_id, text, claim_type, section_path, page_start, page_end, token_count, audience_fit, geography_fit, therapy_area, embedding_id, created_at ) VALUES ( %(chunk_id)s, %(source_id)s, %(version_id)s, %(text)s, %(claim_type)s, %(section_path)s, %(page_start)s, %(page_end)s, %(token_count)s, %(audience_fit)s::json, %(geography_fit)s, %(therapy_area)s, %(embedding_id)s, %(created_at)s ) """, row, ) def _relink_claims_for_source(cursor: "psycopg.Cursor", source_id: str) -> int: """Restore claim_evidence_links for any claim whose ``primary_source_id`` is this source. Replacing chunks cascades-deletes prior evidence links, so we re-anchor each claim to the chunk in the same source whose ``claim_type`` matches and whose text shares the most distinctive tokens with the claim's canonical text. """ cursor.execute( "SELECT claim_id, claim_type, canonical_text FROM claims WHERE primary_source_id = %s", (source_id,), ) claims = cursor.fetchall() if not claims: return 0 cursor.execute( "SELECT chunk_id, claim_type, text FROM chunks WHERE source_id = %s", (source_id,), ) chunks = cursor.fetchall() if not chunks: return 0 links_created = 0 for claim_id, claim_type, canonical_text in claims: claim_tokens = { token for token in (word.strip(".,;:()").lower() for word in str(canonical_text).split()) if len(token) > 3 } best_chunk_id = None best_score = -1 claim_type_value = str(claim_type) # Postgres returns enum as plain string; strip module-qualified prefix if any. normalized_claim_type = claim_type_value.lower().split(".")[-1] for chunk_id, chunk_claim_type, chunk_text in chunks: chunk_tokens = { token for token in (word.strip(".,;:()").lower() for word in str(chunk_text).split()) if len(token) > 3 } overlap = len(claim_tokens & chunk_tokens) if str(chunk_claim_type).lower() == normalized_claim_type: overlap += 5 if overlap > best_score: best_score = overlap best_chunk_id = chunk_id if best_chunk_id is None: continue cursor.execute( """ INSERT INTO claim_evidence_links ( claim_id, chunk_id, source_id, support_type, extraction_confidence, is_primary_support ) VALUES (%s, %s, %s, %s, %s, %s) ON CONFLICT DO NOTHING """, (claim_id, best_chunk_id, source_id, "PRIMARY", 0.7, True), ) links_created += cursor.rowcount or 0 return links_created def _insert_embedding(cursor: "psycopg.Cursor", row: dict[str, Any]) -> None: cursor.execute( """ INSERT INTO chunk_embeddings ( chunk_id, embedding_id, embedding_vector, embedding_model, created_at ) VALUES ( %(chunk_id)s, %(embedding_id)s, %(embedding_vector)s::json, %(embedding_model)s, %(created_at)s ) """, row, ) def _build_ollama_client(): from ollama_client.client import OllamaClient, OllamaSettings enabled = os.getenv("OLLAMA_ENABLED", "true").lower() in {"1", "true", "yes"} return OllamaClient( OllamaSettings( enabled=enabled, base_url=os.getenv("OLLAMA_BASE_URL", "http://127.0.0.1:11434"), embedding_model=os.getenv("OLLAMA_EMBEDDING_MODEL", "qwen3-embedding:8b"), generation_model=os.getenv("OLLAMA_GENERATION_MODEL", "qwen3.5:9b"), request_timeout_seconds=float(os.getenv("OLLAMA_TIMEOUT_SECONDS", "120")), ) ) def ingest( manifest_path: Path, dsn: str, *, embed: bool, max_tokens: int, overlap_tokens: int, min_tokens: int, force_fixtures: bool, ) -> dict[str, Any]: entries = _load_manifest(manifest_path) if not force_fixtures: blocked = [e for e in entries if e.source_id in PROTECTED_SOURCE_IDS] if blocked: names = ", ".join(e.source_id for e in blocked) raise SystemExit( f"Refusing to ingest protected integration-test sources: {names}\n" "Use a separate manifest for new corpus work, or pass --force-fixtures " "(then run: python3 scripts/restore_postgres_fixtures.py to undo)." ) chunker = ScientificChunker() ollama = _build_ollama_client() if embed else None embedding_model = None if embed: assert ollama is not None if not ollama.settings.enabled: raise RuntimeError("Embedding requested but OLLAMA_ENABLED is false in .env.") embedding_model = ollama.resolved_embedding_model() if not embedding_model: raise RuntimeError( f"Required Ollama embedding model {ollama.settings.embedding_model!r} is not " f"registered on {ollama.settings.base_url}. Pull it with: " f"ollama pull {ollama.settings.embedding_model}" ) summary: dict[str, Any] = { "manifest": str(manifest_path), "sources": [], "embedding_model": embedding_model, } conn = psycopg.connect(dsn) try: with conn: with conn.cursor() as cursor: for entry in entries: _ensure_source_registered(cursor, entry) deleted = _delete_existing_chunks(cursor, entry.source_id, entry.version_id) chunks = chunker.preview( ChunkPreviewRequest( source_id=entry.source_id, version_id=entry.version_id, source_class=entry.source_class, therapy_area=entry.therapy_area, geography_scope=entry.geography, audience_suitability=entry.audience, text=entry.text, max_tokens=max_tokens, overlap_tokens=overlap_tokens, min_tokens=min_tokens, ) ) now = datetime.now(UTC) inserted_chunks = 0 inserted_embeddings = 0 for chunk in chunks: embedding_id = None embedding_vector: list[float] | None = None if embed: assert ollama is not None embedding_vector = ollama.embed(chunk.text) embedding_id = f"emb-{uuid.uuid4()}" _insert_chunk( cursor, { "chunk_id": chunk.chunk_id, "source_id": chunk.source_id, "version_id": chunk.version_id, "text": chunk.text, # SQLAlchemy SqlEnum persists Python enum NAME (e.g. DOSE), not value (dose). "claim_type": chunk.claim_type.name if hasattr(chunk.claim_type, "name") else str(chunk.claim_type), "section_path": chunk.section_path, "page_start": chunk.page_start, "page_end": chunk.page_end, "token_count": chunk.token_count, "audience_fit": json.dumps(chunk.audience_fit), "geography_fit": chunk.geography_fit, "therapy_area": chunk.therapy_area, "embedding_id": embedding_id, "created_at": now, }, ) inserted_chunks += 1 if embed and embedding_vector is not None: _insert_embedding( cursor, { "chunk_id": chunk.chunk_id, "embedding_id": embedding_id, "embedding_vector": json.dumps(embedding_vector), "embedding_model": embedding_model, "created_at": now, }, ) inserted_embeddings += 1 relinked = _relink_claims_for_source(cursor, entry.source_id) summary["sources"].append( { "source_id": entry.source_id, "version_id": entry.version_id, "chunks_deleted": deleted, "chunks_inserted": inserted_chunks, "embeddings_inserted": inserted_embeddings, "claim_evidence_links_created": relinked, } ) finally: conn.close() return summary def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--manifest", type=Path, default=REPO_ROOT / "data" / "seed_sources" / "manifest.json", help="Path to manifest JSON (default: data/seed_sources/manifest.json)", ) parser.add_argument( "--database-url", default=os.getenv( "AKS_DATABASE_URL", "postgresql+psycopg://mobcoderid-296@localhost/ai_knowledge_spine", ), help="SQLAlchemy or libpq URL for Postgres", ) parser.add_argument("--no-embed", action="store_true", help="Skip Ollama embeddings (chunks only)") parser.add_argument("--max-tokens", type=int, default=220) parser.add_argument("--overlap-tokens", type=int, default=30) parser.add_argument("--min-tokens", type=int, default=20) parser.add_argument( "--force-fixtures", action="store_true", help="Allow ingest on integration-test source IDs (destructive; restore fixtures after).", ) args = parser.parse_args() dsn = _dsn_from_sqlalchemy_url(args.database_url) summary = ingest( args.manifest, dsn, embed=not args.no_embed, max_tokens=args.max_tokens, overlap_tokens=args.overlap_tokens, min_tokens=args.min_tokens, force_fixtures=args.force_fixtures, ) print(json.dumps(summary, indent=2, default=str)) return 0 if __name__ == "__main__": raise SystemExit(main())