Spaces:
Sleeping
Sleeping
File size: 5,131 Bytes
e5e35a3 db2df31 cc8beab db2df31 e5e35a3 db2df31 e5e35a3 cc8beab db2df31 e5e35a3 db2df31 e5e35a3 db2df31 e5e35a3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | 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
|