src/rag/TavilyWebCache.py

import hashlib
import logging
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, 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
from src.core.settings import get_settings

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class CachedWebSource:
    title: str
    url: str
    excerpt: str
    score: float


class TavilyWebCache:
    """
    Stores Tavily search results in a local Chroma collection so future education queries
    can be answered from cached web snippets (with citations) instead of refetching.
    """

    def __init__(
        self,
        persist_directory: str = PERSIST_DIRECTORY,
        collection_name: str = COLLECTION_NAME,
    ):
        self.persist_directory = persist_directory
        settings = get_settings()
        self.embeddings = OpenAIEmbeddings(model=settings.models.embedding_model)
        ensure_persist_dir(self.persist_directory)
        self.vector_store = Chroma(
            collection_name=collection_name,
            embedding_function=self.embeddings,
            persist_directory=self.persist_directory,
        )

    def _doc_id(self, url: str, title: str) -> str:
        h = hashlib.sha256()
        h.update(url.strip().encode("utf-8"))
        h.update(b"\x00")
        h.update(title.strip().encode("utf-8"))
        return h.hexdigest()

    def save_results(
        self, query: str, results: List[Dict[str, Any]], intent: str = "education"
    ) -> int:
        """
        Saves Tavily results (list of dicts with title/url/content) into Chroma.
        Deduplicates by deterministic id (url+title).
        """
        if not results:
            return 0

        texts: List[str] = []
        metadatas: List[Dict[str, Any]] = []
        ids: List[str] = []
        now = datetime.now(timezone.utc).isoformat()

        for r in results:
            title = str(r.get("title") or "").strip()
            url = str(r.get("url") or "").strip()
            content = str(r.get("content") or "").strip()
            if not url and not content:
                continue
            doc_text = f"TITLE: {title}\nURL: {url}\nCONTENT:\n{content}".strip()
            doc_id = self._doc_id(url or content[:120], title or "untitled")

            texts.append(doc_text)
            metadatas.append(
                {
                    "namespace": "tavily",
                    "provider": "tavily",
                    "intent": intent,
                    "query": query,
                    "title": title,
                    "url": url,
                    "fetched_at": now,
                }
            )
            ids.append(doc_id)

        if not ids:
            return 0

        try:
            existing = set(self.vector_store.get(ids=ids).get("ids", []))
        except Exception:
            existing = set()

        to_add = [
            (t, m, i) for t, m, i in zip(texts, metadatas, ids) if i not in existing
        ]
        if not to_add:
            return 0

        add_texts = [t for t, _, _ in to_add]
        add_metas = [m for _, m, _ in to_add]
        add_ids = [i for _, _, i in to_add]

        try:
            self.vector_store.add_texts(
                texts=add_texts, metadatas=add_metas, ids=add_ids
            )
        except Exception:
            logger.exception("Failed saving Tavily results to Chroma web cache")
            return 0

        return len(to_add)

    def retrieve(
        self,
        query: str,
        k: int = 4,
        intent: str = "education",
        threshold: float = 0.65,
    ) -> List[CachedWebSource]:
        """
        Retrieves cached web snippets if a sufficiently similar match exists.
        """
        where: Optional[Dict[str, Any]] = {
            "namespace": "tavily",
            "provider": "tavily",
            "intent": intent,
        }

        try:
            results: List[Tuple[Any, float]] = (
                self.vector_store.similarity_search_with_relevance_scores(
                    query, k=k, filter=where
                )
            )
        except TypeError:
            results = self.vector_store.similarity_search_with_relevance_scores(
                query, k=k, where=where
            )  # type: ignore[call-arg]
        except Exception:
            logger.exception("Web cache retrieve failed")
            return []

        if not results:
            return []

        out: List[CachedWebSource] = []
        for doc, score in results:
            if score < threshold:
                continue
            meta = getattr(doc, "metadata", {}) or {}
            out.append(
                CachedWebSource(
                    title=str(meta.get("title") or "Web Source"),
                    url=str(meta.get("url") or ""),
                    excerpt=str(getattr(doc, "page_content", "") or "")[:900],
                    score=float(score),
                )
            )
        return out