# retoor from __future__ import annotations import logging import math import re from collections import Counter from dataclasses import dataclass, field from urllib.parse import urlparse from devplacepy.config import DEEPSEARCH_CHROMA_DIR logger = logging.getLogger(__name__) TOKEN_PATTERN = re.compile(r"[a-z0-9]+") BM25_K1 = 1.5 BM25_B = 0.75 RRF_K = 60.0 DEFAULT_TOP_K = 8 CANDIDATE_MULTIPLIER = 4 @dataclass class Chunk: uid: str text: str url: str title: str depth: int = 0 source: str = "" position: int = 0 score: float = 0.0 metadata: dict = field(default_factory=dict) def _tokenize(text: str) -> list[str]: return TOKEN_PATTERN.findall((text or "").lower()) class VectorStore: def __init__(self, collection_name: str) -> None: self.collection_name = collection_name self._client = None self._collection = None self._dims: int | None = None def _ensure(self): if self._collection is not None: return self._collection import chromadb DEEPSEARCH_CHROMA_DIR.mkdir(parents=True, exist_ok=True) self._client = chromadb.PersistentClient(path=str(DEEPSEARCH_CHROMA_DIR)) self._collection = self._client.get_or_create_collection( name=self.collection_name, metadata={"hnsw:space": "cosine"} ) return self._collection @property def dims(self) -> int | None: if self._dims is not None: return self._dims try: collection = self._ensure() data = collection.get(include=["embeddings"], limit=1) rows = data.get("embeddings") or [] if rows and rows[0]: self._dims = len(rows[0]) except Exception: return self._dims return self._dims def add(self, chunks: list[Chunk], vectors: list[list[float]]) -> None: if not chunks: return keep_chunks: list[Chunk] = [] keep_vectors: list[list[float]] = [] for chunk, vector in zip(chunks, vectors): if not vector: continue if self._dims is None: self._dims = len(vector) if len(vector) != self._dims: logger.warning( "deepsearch dropping chunk with mismatched embedding dim %d != %d", len(vector), self._dims, ) continue keep_chunks.append(chunk) keep_vectors.append(vector) if not keep_chunks: return chunks = keep_chunks vectors = keep_vectors collection = self._ensure() collection.add( ids=[chunk.uid for chunk in chunks], embeddings=vectors, documents=[chunk.text for chunk in chunks], metadatas=[ { "url": chunk.url, "title": chunk.title, "depth": chunk.depth, "source": chunk.source, "position": chunk.position, } for chunk in chunks ], ) def all_chunks(self) -> list[Chunk]: collection = self._ensure() data = collection.get(include=["documents", "metadatas"]) chunks: list[Chunk] = [] ids = data.get("ids") or [] documents = data.get("documents") or [] metadatas = data.get("metadatas") or [] for index, uid in enumerate(ids): meta = metadatas[index] if index < len(metadatas) else {} chunks.append( Chunk( uid=uid, text=documents[index] if index < len(documents) else "", url=str(meta.get("url", "")), title=str(meta.get("title", "")), depth=int(meta.get("depth", 0) or 0), source=str(meta.get("source", "")), position=int(meta.get("position", 0) or 0), metadata=dict(meta), ) ) return chunks def count(self) -> int: try: return self._ensure().count() except Exception: return 0 def vector_search( self, query_vector: list[float], top_k: int, where: dict | None = None ) -> list[Chunk]: collection = self._ensure() result = collection.query( query_embeddings=[query_vector], n_results=top_k, where=where or None, include=["documents", "metadatas", "distances"], ) ids = (result.get("ids") or [[]])[0] documents = (result.get("documents") or [[]])[0] metadatas = (result.get("metadatas") or [[]])[0] distances = (result.get("distances") or [[]])[0] chunks: list[Chunk] = [] for index, uid in enumerate(ids): meta = metadatas[index] if index < len(metadatas) else {} distance = distances[index] if index < len(distances) else 1.0 chunks.append( Chunk( uid=uid, text=documents[index] if index < len(documents) else "", url=str(meta.get("url", "")), title=str(meta.get("title", "")), depth=int(meta.get("depth", 0) or 0), source=str(meta.get("source", "")), position=int(meta.get("position", 0) or 0), score=1.0 - float(distance), metadata=dict(meta), ) ) return chunks def keyword_scores(self, query: str, chunks: list[Chunk]) -> dict[str, float]: terms = _tokenize(query) if not terms or not chunks: return {} docs = [_tokenize(chunk.text) for chunk in chunks] lengths = [len(doc) for doc in docs] avg_len = (sum(lengths) / len(lengths)) if lengths else 0.0 doc_freq: Counter = Counter() for doc in docs: for term in set(doc): if term in terms: doc_freq[term] += 1 total_docs = len(docs) scores: dict[str, float] = {} for index, chunk in enumerate(chunks): counts = Counter(docs[index]) length = lengths[index] or 1 score = 0.0 for term in terms: freq = counts.get(term, 0) if freq == 0: continue idf = math.log( 1 + (total_docs - doc_freq[term] + 0.5) / (doc_freq[term] + 0.5) ) denom = freq + BM25_K1 * ( 1 - BM25_B + BM25_B * (length / (avg_len or 1)) ) score += idf * (freq * (BM25_K1 + 1)) / (denom or 1) scores[chunk.uid] = score return scores def hybrid_search( self, query: str, query_vector: list[float], top_k: int = DEFAULT_TOP_K, where: dict | None = None, ) -> list[Chunk]: candidates = self.vector_search( query_vector, top_k * CANDIDATE_MULTIPLIER, where ) if not candidates: return [] vector_rank = sorted( candidates, key=lambda chunk: chunk.score, reverse=True ) keyword = self.keyword_scores(query, candidates) keyword_rank = sorted( candidates, key=lambda chunk: keyword.get(chunk.uid, 0.0), reverse=True ) fused: dict[str, float] = {} for rank, chunk in enumerate(vector_rank, start=1): fused[chunk.uid] = fused.get(chunk.uid, 0.0) + 1.0 / (RRF_K + rank) for rank, chunk in enumerate(keyword_rank, start=1): fused[chunk.uid] = fused.get(chunk.uid, 0.0) + 1.0 / (RRF_K + rank) for chunk in candidates: chunk.score = fused.get(chunk.uid, 0.0) candidates.sort(key=lambda chunk: chunk.score, reverse=True) return candidates[:top_k] def coverage_analytics(self) -> dict: chunks = self.all_chunks() if not chunks: return {"chunks": 0, "domains": 0, "sources": 0, "avg_chunk_chars": 0} domains = { urlparse(chunk.metadata.get("url", "")).netloc for chunk in chunks } domains.discard("") sources = {chunk.source for chunk in chunks} sources.discard("") avg_chars = int(sum(len(chunk.text) for chunk in chunks) / len(chunks)) return { "chunks": len(chunks), "domains": len(domains), "sources": len(sources), "avg_chunk_chars": avg_chars, } def drop(self) -> None: try: import chromadb DEEPSEARCH_CHROMA_DIR.mkdir(parents=True, exist_ok=True) client = self._client or chromadb.PersistentClient( path=str(DEEPSEARCH_CHROMA_DIR) ) client.delete_collection(self.collection_name) except Exception as exc: logger.info("deepsearch collection drop skipped for %s: %s", self.collection_name, exc) finally: self._collection = None