feat: add deepsearch research system with CLI prune/clear and database schema

Implement a multi-agent deep web research subsystem including CLI commands for pruning expired jobs and clearing all artifacts, database tables for sessions/messages/URL cache with indexes, config paths for chroma storage, and internal embed URL for vector operations.
2026-06-14 01:34:21 +00:00
parent c7770ee21a
commit 4ae0b0db5d
53 changed files with 3956 additions and 18 deletions
@@ -0,0 +1,209 @@
+# retoor <retoor@molodetz.nl>
+
+from __future__ import annotations
+
+import logging
+import math
+import re
+from collections import Counter
+from dataclasses import dataclass, field
+
+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
+HYBRID_VECTOR_WEIGHT = 0.6
+HYBRID_KEYWORD_WEIGHT = 0.4
+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
+
+    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
+
+    def add(self, chunks: list[Chunk], vectors: list[list[float]]) -> None:
+        if not chunks:
+            return
+        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 []
+        keyword = self.keyword_scores(query, candidates)
+        vec_max = max((chunk.score for chunk in candidates), default=0.0) or 1.0
+        kw_max = max(keyword.values(), default=0.0) or 1.0
+        for chunk in candidates:
+            vec_norm = max(0.0, chunk.score) / vec_max
+            kw_norm = keyword.get(chunk.uid, 0.0) / kw_max
+            chunk.score = (
+                HYBRID_VECTOR_WEIGHT * vec_norm + HYBRID_KEYWORD_WEIGHT * kw_norm
+            )
+        candidates.sort(key=lambda chunk: chunk.score, reverse=True)
+        return candidates[:top_k]
+
+    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