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