gnosisllm-knowledge 0.2.0__py3-none-any.whl

gnosisllm_knowledge/backends/opensearch/searcher.py

@@ -0,0 +1,383 @@
+"""OpenSearch knowledge searcher implementation.
+
+Uses OpenSearch neural search - embeddings are generated automatically
+by the deployed ML model. No Python-side embedding generation needed.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import TYPE_CHECKING, Any
+
+from gnosisllm_knowledge.backends.opensearch.config import OpenSearchConfig
+from gnosisllm_knowledge.backends.opensearch.queries import QueryBuilder
+from gnosisllm_knowledge.core.domain.search import (
+    SearchMode,
+    SearchQuery,
+    SearchResult,
+    SearchResultItem,
+)
+from gnosisllm_knowledge.core.exceptions import SearchError
+
+if TYPE_CHECKING:
+    from opensearchpy import AsyncOpenSearch
+
+logger = logging.getLogger(__name__)
+
+
+class OpenSearchKnowledgeSearcher:
+    """OpenSearch knowledge searcher.
+
+    Implements the IKnowledgeSearcher protocol for semantic, keyword,
+    and hybrid search over knowledge documents.
+
+    Uses OpenSearch neural search for semantic queries - the deployed
+    ML model handles embedding generation automatically.
+
+    Example:
+        ```python
+        config = OpenSearchConfig.from_env()
+        client = AsyncOpenSearch(hosts=[config.url])
+        searcher = OpenSearchKnowledgeSearcher(client, config)
+
+        query = SearchQuery(text="how to configure", mode=SearchMode.HYBRID)
+        results = await searcher.search(query, "my-index")
+        ```
+    """
+
+    def __init__(
+        self,
+        client: AsyncOpenSearch,
+        config: OpenSearchConfig,
+    ) -> None:
+        """Initialize the searcher.
+
+        Args:
+            client: OpenSearch async client.
+            config: OpenSearch configuration (includes model_id).
+        """
+        self._client = client
+        self._config = config
+
+    async def search(
+        self,
+        query: SearchQuery,
+        index_name: str,
+        **options: Any,
+    ) -> SearchResult:
+        """Execute search query using the specified mode.
+
+        Args:
+            query: Search query with mode and parameters.
+            index_name: Index to search.
+            **options: Additional search options.
+
+        Returns:
+            Search results.
+        """
+        mode = query.mode
+
+        if mode == SearchMode.SEMANTIC:
+            return await self.semantic_search(query, index_name, **options)
+        elif mode == SearchMode.KEYWORD:
+            return await self.keyword_search(query, index_name, **options)
+        elif mode == SearchMode.HYBRID:
+            return await self.hybrid_search(query, index_name, **options)
+        else:
+            # Default to hybrid
+            return await self.hybrid_search(query, index_name, **options)
+
+    async def semantic_search(
+        self,
+        query: SearchQuery,
+        index_name: str,
+        **options: Any,
+    ) -> SearchResult:
+        """Execute semantic (neural) search.
+
+        OpenSearch handles embedding generation via the deployed model.
+
+        Args:
+            query: Search query.
+            index_name: Index to search.
+            **options: Additional options.
+
+        Returns:
+            Search results.
+        """
+        start_time = time.perf_counter()
+
+        model_id = options.get("model_id", self._config.model_id)
+        if not model_id:
+            raise SearchError(
+                message="model_id required for semantic search",
+                details={"query": query.text[:100]},
+            )
+
+        # Build and execute query
+        builder = QueryBuilder(
+            query,
+            model_id=model_id,
+            embedding_field=self._config.embedding_field,
+        )
+        os_query = builder.build_semantic_query()
+
+        result = await self._execute_search(
+            query, os_query, index_name, start_time
+        )
+        return result
+
+    async def keyword_search(
+        self,
+        query: SearchQuery,
+        index_name: str,
+        **options: Any,
+    ) -> SearchResult:
+        """Execute keyword (BM25) search.
+
+        Args:
+            query: Search query.
+            index_name: Index to search.
+            **options: Additional options.
+
+        Returns:
+            Search results.
+        """
+        start_time = time.perf_counter()
+
+        builder = QueryBuilder(query)
+        os_query = builder.build_keyword_query()
+
+        result = await self._execute_search(
+            query, os_query, index_name, start_time
+        )
+        return result
+
+    async def hybrid_search(
+        self,
+        query: SearchQuery,
+        index_name: str,
+        **options: Any,
+    ) -> SearchResult:
+        """Execute hybrid search (semantic + keyword).
+
+        Uses OpenSearch hybrid query with search pipeline for
+        score normalization.
+
+        Args:
+            query: Search query.
+            index_name: Index to search.
+            **options: semantic_weight, keyword_weight, model_id.
+
+        Returns:
+            Search results.
+        """
+        start_time = time.perf_counter()
+
+        model_id = options.get("model_id", self._config.model_id)
+
+        # Build hybrid query
+        semantic_weight = options.get("semantic_weight", 0.7)
+        keyword_weight = options.get("keyword_weight", 0.3)
+
+        builder = QueryBuilder(
+            query,
+            model_id=model_id,
+            embedding_field=self._config.embedding_field,
+        )
+        os_query = builder.build_hybrid_query(
+            semantic_weight=semantic_weight,
+            keyword_weight=keyword_weight,
+        )
+
+        # Use search pipeline for hybrid if configured
+        search_params = {}
+        if self._config.search_pipeline_name:
+            search_params["search_pipeline"] = self._config.search_pipeline_name
+
+        result = await self._execute_search(
+            query, os_query, index_name, start_time, **search_params
+        )
+        return result
+
+    async def get_similar_documents(
+        self,
+        doc_id: str,
+        index_name: str,
+        limit: int = 10,
+    ) -> SearchResult:
+        """Find documents similar to a given document.
+
+        Args:
+            doc_id: Document ID to find similar documents for.
+            index_name: Index to search.
+            limit: Maximum results.
+
+        Returns:
+            Search results with similar documents.
+        """
+        start_time = time.perf_counter()
+
+        query = SearchQuery(text="", limit=limit)
+        builder = QueryBuilder(query)
+        os_query = builder.build_more_like_this_query(doc_id)
+
+        result = await self._execute_search(
+            query, os_query, index_name, start_time
+        )
+        return result
+
+    async def multi_search(
+        self,
+        queries: list[SearchQuery],
+        index_name: str,
+    ) -> list[SearchResult]:
+        """Execute multiple searches in a single request.
+
+        Args:
+            queries: List of search queries.
+            index_name: Index to search.
+
+        Returns:
+            List of search results.
+        """
+        if not queries:
+            return []
+
+        start_time = time.perf_counter()
+        msearch_body = []
+
+        for query in queries:
+            builder = QueryBuilder(
+                query,
+                model_id=self._config.model_id,
+                embedding_field=self._config.embedding_field,
+            )
+
+            if query.mode == SearchMode.SEMANTIC:
+                os_query = builder.build_semantic_query()
+            elif query.mode == SearchMode.KEYWORD:
+                os_query = builder.build_keyword_query()
+            else:
+                os_query = builder.build_hybrid_query()
+
+            msearch_body.append({"index": index_name})
+            msearch_body.append(os_query)
+
+        response = await self._client.msearch(body=msearch_body)
+
+        results = []
+        for i, resp in enumerate(response.get("responses", [])):
+            duration_ms = (time.perf_counter() - start_time) * 1000
+            result = self._parse_response(queries[i], resp, duration_ms)
+            results.append(result)
+
+        return results
+
+    async def _execute_search(
+        self,
+        query: SearchQuery,
+        os_query: dict[str, Any],
+        index_name: str,
+        start_time: float,
+        **params: Any,
+    ) -> SearchResult:
+        """Execute search and parse results.
+
+        Args:
+            query: Original search query.
+            os_query: OpenSearch query dictionary.
+            index_name: Index to search.
+            start_time: Search start time.
+            **params: Additional search parameters.
+
+        Returns:
+            Parsed search results.
+        """
+        try:
+            response = await self._client.search(
+                index=index_name,
+                body=os_query,
+                **params,
+            )
+
+            duration_ms = (time.perf_counter() - start_time) * 1000
+            return self._parse_response(query, response, duration_ms)
+
+        except Exception as e:
+            logger.error(f"Search failed: {e}")
+            raise SearchError(
+                message=f"Search failed: {e}",
+                details={"query": query.text[:100], "index": index_name},
+                cause=e,
+            ) from e
+
+    def _parse_response(
+        self,
+        query: SearchQuery,
+        response: dict[str, Any],
+        duration_ms: float,
+    ) -> SearchResult:
+        """Parse OpenSearch response into SearchResult.
+
+        Args:
+            query: Original query.
+            response: OpenSearch response.
+            duration_ms: Search duration in milliseconds.
+
+        Returns:
+            Parsed search result.
+        """
+        hits = response.get("hits", {})
+        total = hits.get("total", {})
+        total_hits = total.get("value", 0) if isinstance(total, dict) else total
+        max_score = hits.get("max_score")
+
+        items = []
+        search_after_token = None
+
+        for hit in hits.get("hits", []):
+            source = hit.get("_source", {})
+
+            # Extract highlights
+            highlights = None
+            highlighted_title = None
+            if "highlight" in hit:
+                highlight_data = hit["highlight"]
+                highlights = highlight_data.get("content", [])
+                title_highlights = highlight_data.get("title", [])
+                if title_highlights:
+                    highlighted_title = title_highlights[0]
+
+            item = SearchResultItem(
+                doc_id=hit.get("_id", ""),
+                content=source.get("content", ""),
+                score=hit.get("_score", 0.0),
+                title=source.get("title"),
+                url=source.get("url"),
+                source=source.get("source"),
+                collection_id=source.get("collection_id"),
+                highlights=highlights,
+                highlighted_title=highlighted_title,
+                metadata=source.get("metadata"),
+                chunk_index=source.get("chunk_index"),
+                total_chunks=source.get("total_chunks"),
+                explanation=hit.get("_explanation") if query.explain else None,
+            )
+            items.append(item)
+
+            # Track search_after for cursor pagination
+            if "sort" in hit:
+                search_after_token = hit["sort"]
+
+        return SearchResult(
+            query=query.text,
+            mode=query.mode,
+            items=items,
+            total_hits=total_hits,
+            duration_ms=duration_ms,
+            max_score=max_score,
+            search_after_token=search_after_token,
+            has_more=len(items) == query.limit and total_hits > query.offset + len(items),
+        )