# retoor <retoor@molodetz.nl>
from __future__ import annotations
import asyncio
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
CHROMADB_TIMEOUT = 30.0
@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
async def _run_sync(self, func, *args, timeout: float = CHROMADB_TIMEOUT):
"""Run a synchronous ChromaDB call in a thread executor with a timeout."""
try:
return await asyncio.wait_for(
asyncio.to_thread(func, *args), timeout=timeout
)
except asyncio.TimeoutError:
logger.error(
"deepsearch ChromaDB operation timed out after %.1fs on collection %s",
timeout,
self.collection_name,
)
raise
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_sync(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
],
)
async def add(self, chunks: list[Chunk], vectors: list[list[float]]) -> None:
if not chunks:
return
await self._run_sync(self._add_sync, chunks, vectors)
def _all_chunks_sync(self, limit: int = 0) -> list[Chunk]:
collection = self._ensure()
kwargs: dict = {"include": ["documents", "metadatas"]}
if limit > 0:
kwargs["limit"] = limit
data = collection.get(**kwargs)
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
async def all_chunks(self, limit: int = 1000) -> list[Chunk]:
return await self._run_sync(self._all_chunks_sync, limit)
async def count(self) -> int:
try:
collection = self._ensure()
return await self._run_sync(collection.count)
except Exception:
return 0
def _vector_search_sync(
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
async def vector_search(
self, query_vector: list[float], top_k: int, where: dict | None = None
) -> list[Chunk]:
return await self._run_sync(
self._vector_search_sync, query_vector, top_k, where
)
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
async def hybrid_search(
self,
query: str,
query_vector: list[float],
top_k: int = DEFAULT_TOP_K,
where: dict | None = None,
) -> list[Chunk]:
candidates = await 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]
async def coverage_analytics(self, sample_limit: int = 5000) -> dict:
chunks = await self.all_chunks(limit=sample_limit)
if not chunks:
total = await self.count()
return {"chunks": total, "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
Reference in New Issue View Git Blame Copy Permalink