gnosisllm-knowledge 0.2.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

gnosisllm_knowledge/backends/opensearch/searcher.py

@@ -2,6 +2,10 @@
 
 Uses OpenSearch neural search - embeddings are generated automatically
 by the deployed ML model. No Python-side embedding generation needed.
+
+Note: This module is tenant-agnostic. Multi-tenancy should be handled
+at the API layer by using separate indices per account (e.g.,
+`knowledge-{account_id}`) rather than filtering by account_id.
 """
 
 from __future__ import annotations
@@ -296,12 +300,25 @@ class OpenSearchKnowledgeSearcher:
             Parsed search results.
         """
         try:
+            logger.debug(f"OpenSearch query: {os_query}")
+            logger.debug(f"Query include_highlights: {query.include_highlights}")
+
             response = await self._client.search(
                 index=index_name,
                 body=os_query,
                 **params,
             )
 
+            # Debug: Log first hit to see if highlights are present
+            hits = response.get("hits", {}).get("hits", [])
+            if hits:
+                first_hit = hits[0]
+                logger.debug(f"First hit keys: {first_hit.keys()}")
+                if "highlight" in first_hit:
+                    logger.debug(f"Highlight data: {first_hit['highlight']}")
+                else:
+                    logger.debug("No 'highlight' key in response hit")
+
             duration_ms = (time.perf_counter() - start_time) * 1000
             return self._parse_response(query, response, duration_ms)
 
@@ -381,3 +398,224 @@ class OpenSearchKnowledgeSearcher:
             search_after_token=search_after_token,
             has_more=len(items) == query.limit and total_hits > query.offset + len(items),
         )
+
+    async def get_collections(self, index_name: str) -> list[dict[str, Any]]:
+        """Get unique collections with document counts via aggregation.
+
+        Args:
+            index_name: Index to query.
+
+        Returns:
+            List of collections with id, name, and document_count.
+        """
+        try:
+            # Check if index exists
+            exists = await self._client.indices.exists(index=index_name)
+            if not exists:
+                logger.debug(f"Index {index_name} does not exist")
+                return []
+
+            # Aggregation query for unique collection_ids with counts
+            # Also aggregate collection_name for display
+            query = {
+                "size": 0,
+                "aggs": {
+                    "collections": {
+                        "terms": {
+                            "field": "collection_id",
+                            "size": 1000,  # Max collections to return
+                        },
+                        "aggs": {
+                            "collection_name": {
+                                "terms": {
+                                    "field": "collection_name",
+                                    "size": 1,
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+
+            response = await self._client.search(index=index_name, body=query)
+
+            collections = []
+            buckets = response.get("aggregations", {}).get("collections", {}).get("buckets", [])
+
+            for bucket in buckets:
+                collection_id = bucket.get("key")
+                if not collection_id:
+                    continue
+
+                doc_count = bucket.get("doc_count", 0)
+
+                # Get collection name from nested agg or use ID as fallback
+                name_buckets = bucket.get("collection_name", {}).get("buckets", [])
+                collection_name = name_buckets[0].get("key") if name_buckets else collection_id
+
+                collections.append({
+                    "id": collection_id,
+                    "name": collection_name or collection_id,
+                    "document_count": doc_count,
+                })
+
+            logger.debug(f"Found {len(collections)} collections in {index_name}")
+            return collections
+
+        except Exception as e:
+            logger.error(f"Failed to get collections from {index_name}: {e}")
+            return []
+
+    async def get_stats(self, index_name: str) -> dict[str, Any]:
+        """Get index statistics.
+
+        Args:
+            index_name: Index to query.
+
+        Returns:
+            Dictionary with document_count and index info.
+        """
+        try:
+            # Check if index exists
+            exists = await self._client.indices.exists(index=index_name)
+            if not exists:
+                return {
+                    "document_count": 0,
+                    "index_name": index_name,
+                    "exists": False,
+                }
+
+            # Get index stats
+            stats = await self._client.indices.stats(index=index_name)
+            index_stats = stats.get("indices", {}).get(index_name, {})
+            primaries = index_stats.get("primaries", {})
+            docs = primaries.get("docs", {})
+
+            return {
+                "document_count": docs.get("count", 0),
+                "index_name": index_name,
+                "exists": True,
+                "size_bytes": primaries.get("store", {}).get("size_in_bytes", 0),
+            }
+
+        except Exception as e:
+            logger.error(f"Failed to get stats for {index_name}: {e}")
+            return {
+                "document_count": 0,
+                "index_name": index_name,
+                "error": str(e),
+            }
+
+    async def list_documents(
+        self,
+        index_name: str,
+        *,
+        source_id: str | None = None,
+        collection_id: str | None = None,
+        limit: int = 50,
+        offset: int = 0,
+    ) -> dict[str, Any]:
+        """List documents with optional filters.
+
+        Note:
+            This method is tenant-agnostic. Multi-tenancy should be handled
+            at the API layer by using separate indices per account.
+
+        Args:
+            index_name: Index to query (use tenant-specific name for isolation).
+            source_id: Optional source ID filter.
+            collection_id: Optional collection ID filter.
+            limit: Maximum documents to return.
+            offset: Number of documents to skip.
+
+        Returns:
+            Dictionary with documents, total, limit, offset.
+        """
+        try:
+            # Check if index exists
+            exists = await self._client.indices.exists(index=index_name)
+            if not exists:
+                logger.debug(f"Index {index_name} does not exist")
+                return {
+                    "documents": [],
+                    "total": 0,
+                    "limit": limit,
+                    "offset": offset,
+                }
+
+            # Build filter clauses
+            filters: list[dict[str, Any]] = []
+
+            if source_id:
+                filters.append({"term": {"source_id": source_id}})
+
+            if collection_id:
+                filters.append({"term": {"collection_id": collection_id}})
+
+            # Build query with match_all and filters
+            query: dict[str, Any] = {
+                "size": limit,
+                "from": offset,
+                "sort": [
+                    {"created_at": {"order": "desc", "unmapped_type": "date"}},
+                    {"_id": {"order": "asc"}},
+                ],
+                "_source": {
+                    "excludes": ["content_embedding"],
+                },
+            }
+
+            if filters:
+                query["query"] = {
+                    "bool": {
+                        "must": [{"match_all": {}}],
+                        "filter": filters,
+                    }
+                }
+            else:
+                query["query"] = {"match_all": {}}
+
+            response = await self._client.search(index=index_name, body=query)
+
+            hits = response.get("hits", {})
+            total = hits.get("total", {})
+            total_hits = total.get("value", 0) if isinstance(total, dict) else total
+
+            documents = []
+            for hit in hits.get("hits", []):
+                source = hit.get("_source", {})
+                doc = {
+                    "id": hit.get("_id", ""),
+                    "title": source.get("title"),
+                    "url": source.get("url"),
+                    "content_preview": (source.get("content", ""))[:200],
+                    "content": source.get("content"),
+                    "chunk_index": source.get("chunk_index"),
+                    "total_chunks": source.get("total_chunks"),
+                    "source_id": source.get("source_id"),
+                    "collection_id": source.get("collection_id"),
+                    "created_at": source.get("created_at"),
+                    "metadata": source.get("metadata"),
+                }
+                documents.append(doc)
+
+            logger.debug(
+                f"Listed {len(documents)} documents from {index_name} "
+                f"(total: {total_hits}, source_id: {source_id})"
+            )
+
+            return {
+                "documents": documents,
+                "total": total_hits,
+                "limit": limit,
+                "offset": offset,
+            }
+
+        except Exception as e:
+            logger.error(f"Failed to list documents from {index_name}: {e}")
+            return {
+                "documents": [],
+                "total": 0,
+                "limit": limit,
+                "offset": offset,
+            }