Update

devplacepy/services/deepsearch/store.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import logging
 import math
 import re
@@ -19,6 +20,7 @@ BM25_B = 0.75
 RRF_K = 60.0
 DEFAULT_TOP_K = 8
 CANDIDATE_MULTIPLIER = 4
+CHROMADB_TIMEOUT = 30.0
 
 
 @dataclass
@@ -45,6 +47,20 @@ class VectorStore:
         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
@@ -71,7 +87,7 @@ class VectorStore:
             return self._dims
         return self._dims
 
-    def add(self, chunks: list[Chunk], vectors: list[list[float]]) -> None:
+    def _add_sync(self, chunks: list[Chunk], vectors: list[list[float]]) -> None:
         if not chunks:
             return
         keep_chunks: list[Chunk] = []
@@ -111,9 +127,17 @@ class VectorStore:
             ],
         )
 
-    def all_chunks(self) -> list[Chunk]:
+    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()
-        data = collection.get(include=["documents", "metadatas"])
+        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 []
@@ -134,13 +158,17 @@ class VectorStore:
             )
         return chunks
 
-    def count(self) -> int:
+    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:
-            return self._ensure().count()
+            collection = self._ensure()
+            return await self._run_sync(collection.count)
         except Exception:
             return 0
 
-    def vector_search(
+    def _vector_search_sync(
         self, query_vector: list[float], top_k: int, where: dict | None = None
     ) -> list[Chunk]:
         collection = self._ensure()
@@ -173,6 +201,13 @@ class VectorStore:
             )
         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:
@@ -205,14 +240,14 @@ class VectorStore:
             scores[chunk.uid] = score
         return scores
 
-    def hybrid_search(
+    async 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(
+        candidates = await self.vector_search(
             query_vector, top_k * CANDIDATE_MULTIPLIER, where
         )
         if not candidates:
@@ -234,10 +269,11 @@ class VectorStore:
         candidates.sort(key=lambda chunk: chunk.score, reverse=True)
         return candidates[:top_k]
 
-    def coverage_analytics(self) -> dict:
-        chunks = self.all_chunks()
+    async def coverage_analytics(self, sample_limit: int = 5000) -> dict:
+        chunks = await self.all_chunks(limit=sample_limit)
         if not chunks:
-            return {"chunks": 0, "domains": 0, "sources": 0, "avg_chunk_chars": 0}
+            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
         }