data_update.py

from __future__ import annotations

import argparse
import shutil
from pathlib import Path


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Clean raw semiconductor documents, chunk them, and sync embeddings into ChromaDB.",
    )
    parser.add_argument(
        "--rebuild",
        action="store_true",
        help="Clear processed outputs and rebuild the full vector index from data/raw.",
    )
    return parser


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()

    from rag_support import (
        MANIFEST_PATH,
        PROCESSED_DIR,
        RAW_DIR,
        build_chunk_records,
        collection_document_count,
        delete_chunk_ids,
        doc_id_for_path,
        embed_texts,
        ensure_data_directories,
        get_pipeline_fingerprint,
        get_chroma_client,
        get_or_create_collection,
        load_embedding_model,
        load_environment,
        list_raw_files,
        prepare_document,
        processed_path_for_doc_id,
        read_manifest,
        reset_collection,
        sha256_file,
        sha256_text,
        source_type_for_path,
        upsert_chunk_records,
        utc_now_iso,
        write_manifest,
    )

    load_environment()
    ensure_data_directories()

    raw_files = list_raw_files(RAW_DIR)
    if not raw_files:
        print("No supported files found under data/raw. Add .md, .txt, or .pdf files first.")
        return 0

    manifest = read_manifest()
    documents: dict[str, dict] = manifest.get("documents", {})
    pipeline_fingerprint = get_pipeline_fingerprint()

    client = get_chroma_client()
    if args.rebuild:
        if PROCESSED_DIR.exists():
            for child in PROCESSED_DIR.iterdir():
                if child.is_file():
                    child.unlink()
                elif child.is_dir():
                    shutil.rmtree(child)
        reset_collection(client)
        manifest = {"version": 2, "pipeline_fingerprint": pipeline_fingerprint, "documents": {}}
        documents = manifest["documents"]

    collection = get_or_create_collection(client)
    embedding_model = load_embedding_model()

    current_doc_ids: set[str] = set()
    added_count = 0
    updated_count = 0
    unchanged_count = 0

    for raw_path in raw_files:
        doc_id = doc_id_for_path(raw_path)
        current_doc_ids.add(doc_id)
        raw_hash = sha256_file(raw_path)
        existing = documents.get(doc_id)
        processed_path = processed_path_for_doc_id(doc_id, PROCESSED_DIR)

        unchanged = (
            existing
            and existing.get("raw_hash") == raw_hash
            and processed_path.exists()
            and existing.get("processed_path") == processed_path.as_posix()
            and existing.get("pipeline_fingerprint") == pipeline_fingerprint
        )
        if unchanged and not args.rebuild:
            unchanged_count += 1
            continue

        source_type = source_type_for_path(raw_path)
        prepared_document = prepare_document(raw_path, source_type)
        if not prepared_document.content_text.strip():
            print(f"Skipping empty document: {raw_path}")
            continue

        processed_path.write_text(prepared_document.processed_text, encoding="utf-8")
        chunk_records = build_chunk_records(
            prepared_document.content_text,
            raw_path,
            source_type,
            prepared_document.metadata,
        )
        embeddings = embed_texts(embedding_model, [record.text for record in chunk_records])

        if existing:
            delete_chunk_ids(collection, existing.get("chunk_ids", []))

        upsert_chunk_records(collection, chunk_records, embeddings)

        documents[doc_id] = {
            "doc_id": doc_id,
            "raw_path": raw_path.as_posix(),
            "raw_hash": raw_hash,
            "processed_path": processed_path.as_posix(),
            "processed_hash": sha256_text(prepared_document.processed_text),
            "updated_at": utc_now_iso(),
            "pipeline_fingerprint": pipeline_fingerprint,
            "source_type": source_type,
            "source_name": prepared_document.metadata.get("source_name", ""),
            "source_date": prepared_document.metadata.get("source_date", ""),
            "source_kind": prepared_document.metadata.get("source_kind", ""),
            "source_url": prepared_document.metadata.get("source_url", ""),
            "license_note": prepared_document.metadata.get("license_note", ""),
            "compliance_note": prepared_document.metadata.get("compliance_note", ""),
            "primary_references": prepared_document.metadata.get("primary_references", ""),
            "chunk_ids": [record.chunk_id for record in chunk_records],
            "chunk_count": len(chunk_records),
            "title": prepared_document.metadata.get("title", raw_path.stem),
        }

        if existing:
            updated_count += 1
            print(f"Updated: {raw_path.name} ({len(chunk_records)} chunks)")
        else:
            added_count += 1
            print(f"Added:   {raw_path.name} ({len(chunk_records)} chunks)")

    removed_count = 0
    for doc_id in sorted(set(documents) - current_doc_ids):
        existing = documents.pop(doc_id)
        delete_chunk_ids(collection, existing.get("chunk_ids", []))
        processed_path = Path(existing.get("processed_path", ""))
        if processed_path.exists():
            processed_path.unlink()
        removed_count += 1
        print(f"Removed: {doc_id}")

    manifest["version"] = 2
    manifest["pipeline_fingerprint"] = pipeline_fingerprint
    manifest["documents"] = documents
    write_manifest(manifest)

    print("")
    print("Sync complete")
    print(f"- Raw files scanned: {len(raw_files)}")
    print(f"- Added: {added_count}")
    print(f"- Updated: {updated_count}")
    print(f"- Unchanged: {unchanged_count}")
    print(f"- Removed: {removed_count}")
    print(f"- Indexed chunks in collection: {collection_document_count(collection)}")
    print(f"- Manifest: {MANIFEST_PATH}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())