src/data/UpdateKBStore.py

import json
import hashlib
import logging
import os
from typing import Dict, List, Tuple

from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings

from src.data.chroma_config import COLLECTION_NAME, PERSIST_DIRECTORY, ensure_persist_dir

# CONFIG
PERSIST_DIR = PERSIST_DIRECTORY
BATCH_SIZE = 100

logger = logging.getLogger(__name__)


class UpdateKBStore:

    def __init__(self, persist_dir: str = PERSIST_DIR, batch_size: int = BATCH_SIZE):
        self.persist_dir = persist_dir
        self.batch_size = batch_size
        self.embedding = OpenAIEmbeddings(model="text-embedding-3-large")
        ensure_persist_dir(self.persist_dir)

    def build_text(self, doc: Dict) -> str:
        return f"""
Domain: {doc.get('domain')}
Category: {doc.get('category')}
Persona: {doc.get('persona')}
Market Context: {doc.get('market_context')}
Source: {doc.get('source')}
KeyWords: {", ".join(doc.get('keywords', []))}

Q: {doc.get('question')}
A: {doc.get('answer')}
"""

    def clean_metadata(self, doc):
        return {
            "id": str(doc.get("id", "")),
            "domain": str(doc.get("domain", "")),
            "category": str(doc.get("category", "")),
            "topic": str(doc.get("topic", "")),
            "persona": str(doc.get("persona", "")),
            "market_context": str(doc.get("market_context", "")),
            "source": str(doc.get("source", "")),
            "regulatory_scope": ",".join(doc.get("regulatory_scope", [])) if doc.get("regulatory_scope") else "",
            "tax_year": int(doc["tax_year"]) if doc.get("tax_year") is not None else 0,
            "difficulty": str(doc.get("difficulty", ""))
        }

    def load_json(self, file_path: str):
        with open(file_path, "r") as f:
            return json.load(f)

    def _doc_id(self, doc: Dict) -> str:
        """
        Stable ID for de-duplication inside the single shared Chroma collection.
        Prefers the explicit `id` field; otherwise hashes the Q/A text.
        """
        raw_id = str(doc.get("id") or "").strip()
        if raw_id:
            return raw_id
        text = (
            (str(doc.get("question") or "") + "\n" + str(doc.get("answer") or ""))
            .strip()
        )
        h = hashlib.sha256()
        h.update(text.encode("utf-8"))
        return h.hexdigest()

    def ingest_file(self, file_path: str, collection_name: str) -> Dict:
        data = self.load_json(file_path)

        # This project uses a single shared Chroma collection. `collection_name` is
        # treated as a logical label only (kept in metadata for traceability).
        logical_collection = collection_name

        print(
            f"Ingesting {len(data)} records into {COLLECTION_NAME} (logical={logical_collection}), "
            f"file name = {file_path}"
        )
        db = Chroma(
            collection_name=COLLECTION_NAME,
            embedding_function=self.embedding,
            persist_directory=self.persist_dir
        )

        texts: List[str] = []
        metadatas: List[Dict] = []
        ids: List[str] = []
        inserted = 0

        def flush_batch(batch: List[Tuple[str, Dict, str]]) -> int:
            if not batch:
                return 0
            batch_texts = [t for t, _, _ in batch]
            batch_metas = [m for _, m, _ in batch]
            batch_ids = [i for _, _, i in batch]
            try:
                existing = set(db.get(ids=batch_ids).get("ids", []))
            except Exception:
                existing = set()
            to_add = [
                (t, m, i)
                for t, m, i in zip(batch_texts, batch_metas, batch_ids)
                if i not in existing
            ]
            if not to_add:
                return 0
            db.add_texts(
                texts=[t for t, _, _ in to_add],
                metadatas=[m for _, m, _ in to_add],
                ids=[i for _, _, i in to_add],
            )
            return len(to_add)

        for doc in data:
            text = self.build_text(doc)
            metadata = self.clean_metadata(doc)
            doc_id = self._doc_id(doc)

            metadata = {
                **metadata,
                "namespace": "kb",
                "logical_collection": str(logical_collection),
                "source_file": os.path.basename(file_path),
            }

            texts.append(text)
            metadatas.append(metadata)
            ids.append(doc_id)

            if len(texts) >= self.batch_size:
                inserted += flush_batch(list(zip(texts, metadatas, ids)))
                texts, metadatas = [], []
                ids = []

        if texts and metadatas and ids:
            inserted += flush_batch(list(zip(texts, metadatas, ids)))

        # Persist best-effort (older wrappers require explicit persist).
        try:
            db.persist()
        except Exception:
            logger.debug("Chroma persist() unavailable or failed", exc_info=True)

        return {
            "collection": COLLECTION_NAME,
            "logical_collection": logical_collection,
            "records_inserted": inserted
        }

    def ingest_all_shards(self) -> dict:
        """
        Auto-ingest all JSON shard files under finance_kb_shards directory.

        Returns:
            summary dict for UI
        """

        DATA_DIR = "src/data/finance_kb_shards"

        if not os.path.exists(DATA_DIR):
            return {
                "status": "error",
                "message": f"Directory not found: {DATA_DIR}"
            }

        results = []
        total_records = 0
        files_processed = 0

        for file in os.listdir(DATA_DIR):
            if not file.endswith(".json"):
                continue

            file_path = os.path.join(DATA_DIR, file)
            collection_name = "kb_shards"

            try:
                result = self.ingest_file(file_path, collection_name)

                total_records += result["records_inserted"]
                results.append(result)
                files_processed += 1

            except Exception as e:
                results.append({
                    "collection": collection_name,
                    "error": str(e)
                })

        return {
            "status": "success",
            "files_processed": files_processed,
            "collections": results,
            "total_records": total_records
        }