gnosisllm-knowledge 0.2.0__py3-none-any.whl

gnosisllm_knowledge/backends/opensearch/agentic.py

@@ -0,0 +1,738 @@
+"""OpenSearch agentic searcher implementation.
+
+Uses OpenSearch ML agents for AI-powered search with reasoning capabilities.
+Supports flow agents (fast RAG) and conversational agents (multi-turn with memory).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import uuid
+from datetime import UTC, datetime
+from typing import TYPE_CHECKING, Any
+
+from opensearchpy import AsyncOpenSearch
+
+from gnosisllm_knowledge.backends.opensearch.config import OpenSearchConfig
+from gnosisllm_knowledge.core.domain.search import (
+    AgentType,
+    AgenticSearchQuery,
+    AgenticSearchResult,
+    ReasoningStep,
+    SearchMode,
+    SearchResultItem,
+)
+from gnosisllm_knowledge.core.exceptions import AgenticSearchError
+
+if TYPE_CHECKING:
+    from gnosisllm_knowledge.core.interfaces.agentic import IAgenticSearcher
+
+
+class OpenSearchAgenticSearcher:
+    """Executes agentic search using OpenSearch ML agents.
+
+    Supports two agent types:
+    - Flow Agent: Fast, sequential RAG with minimal reasoning
+    - Conversational Agent: Multi-turn with memory and detailed reasoning
+
+    The searcher integrates with OpenSearch's ML Commons plugin to:
+    1. Execute agents via /_plugins/_ml/agents/{agent_id}/_execute
+    2. Parse agent responses including reasoning traces
+    3. Extract search results from VectorDBTool outputs
+    4. Manage conversation memory for multi-turn interactions
+
+    Example:
+        ```python
+        config = OpenSearchConfig.from_env()
+        client = AsyncOpenSearch(hosts=[config.url])
+        searcher = OpenSearchAgenticSearcher(client, config)
+
+        query = AgenticSearchQuery(
+            text="How do I configure OAuth2?",
+            agent_type=AgentType.FLOW,
+        )
+        result = await searcher.agentic_search(query, "knowledge")
+        print(result.answer)
+        ```
+    """
+
+    def __init__(
+        self,
+        client: AsyncOpenSearch,
+        config: OpenSearchConfig,
+    ) -> None:
+        """Initialize the agentic searcher.
+
+        Args:
+            client: Async OpenSearch client.
+            config: OpenSearch configuration with agent IDs.
+        """
+        self._client = client
+        self._config = config
+        self._logger = logging.getLogger(__name__)
+
+    @property
+    def is_configured(self) -> bool:
+        """Check if at least one agent is configured."""
+        return self.flow_agent_available or self.conversational_agent_available
+
+    @property
+    def flow_agent_available(self) -> bool:
+        """Check if flow agent is configured."""
+        return bool(self._config.flow_agent_id)
+
+    @property
+    def conversational_agent_available(self) -> bool:
+        """Check if conversational agent is configured."""
+        return bool(self._config.conversational_agent_id)
+
+    async def agentic_search(
+        self,
+        query: AgenticSearchQuery,
+        index_name: str,
+        **options: Any,
+    ) -> AgenticSearchResult:
+        """Execute agentic search with agent orchestration.
+
+        The flow:
+        1. Select agent based on query.agent_type
+        2. Build execution request with query and filters
+        3. Execute agent via OpenSearch ML API
+        4. Parse response for answer, reasoning, and results
+
+        Args:
+            query: Agentic search query with agent type and context.
+            index_name: Target index name.
+            **options: Additional agent options.
+
+        Returns:
+            AgenticSearchResult with answer, reasoning, and sources.
+
+        Raises:
+            AgenticSearchError: If agent execution fails.
+        """
+        start = datetime.now(UTC)
+
+        # Select agent based on type
+        agent_id = self._get_agent_id(query.agent_type)
+        if not agent_id:
+            raise AgenticSearchError(
+                message=f"Agent not configured for type: {query.agent_type.value}",
+                agent_type=query.agent_type.value,
+                details={"hint": "Run 'gnosisllm-knowledge agentic setup' to configure agents."},
+            )
+
+        # Build execution request
+        execute_body = self._build_execute_request(query, index_name)
+
+        self._logger.debug(
+            "Executing agentic search",
+            extra={
+                "agent_id": agent_id,
+                "agent_type": query.agent_type.value,
+                "query": query.text[:100],
+            },
+        )
+
+        # Execute agent
+        response = await self._execute_agent(agent_id, execute_body)
+
+        duration_ms = (datetime.now(UTC) - start).total_seconds() * 1000
+
+        return self._parse_agentic_response(query, response, duration_ms)
+
+    async def get_conversation(
+        self,
+        conversation_id: str,
+    ) -> list[dict[str, Any]]:
+        """Get conversation history for multi-turn searches.
+
+        Args:
+            conversation_id: Conversation identifier (memory_id).
+
+        Returns:
+            List of conversation messages with role and content.
+        """
+        try:
+            response = await self._client.transport.perform_request(
+                "GET",
+                f"/_plugins/_ml/memory/{conversation_id}/_messages",
+            )
+            messages = response.get("messages", [])
+            return [
+                {
+                    "role": msg.get("role", "unknown"),
+                    "content": msg.get("content", ""),
+                    "timestamp": msg.get("create_time"),
+                }
+                for msg in messages
+            ]
+        except Exception as e:
+            self._logger.warning(f"Failed to get conversation: {e}")
+            return []
+
+    async def clear_conversation(
+        self,
+        conversation_id: str,
+    ) -> bool:
+        """Clear conversation history.
+
+        Args:
+            conversation_id: Conversation to clear.
+
+        Returns:
+            True if cleared successfully, False if not found.
+        """
+        try:
+            await self._client.transport.perform_request(
+                "DELETE",
+                f"/_plugins/_ml/memory/{conversation_id}",
+            )
+            return True
+        except Exception as e:
+            self._logger.warning(f"Failed to clear conversation: {e}")
+            return False
+
+    async def list_conversations(
+        self,
+        account_id: str | None = None,
+        limit: int = 100,
+    ) -> list[dict[str, Any]]:
+        """List active conversations.
+
+        Args:
+            account_id: Filter by account (multi-tenant).
+            limit: Maximum number of conversations.
+
+        Returns:
+            List of conversation metadata dicts.
+        """
+        try:
+            body: dict[str, Any] = {"size": limit}
+            if account_id:
+                body["query"] = {"term": {"account_id": account_id}}
+
+            response = await self._client.transport.perform_request(
+                "POST",
+                "/_plugins/_ml/memory/_search",
+                body=body,
+            )
+            hits = response.get("hits", {}).get("hits", [])
+            return [
+                {
+                    "conversation_id": hit.get("_id"),
+                    "name": hit.get("_source", {}).get("name"),
+                    "created_at": hit.get("_source", {}).get("create_time"),
+                    "updated_at": hit.get("_source", {}).get("update_time"),
+                }
+                for hit in hits
+            ]
+        except Exception as e:
+            self._logger.warning(f"Failed to list conversations: {e}")
+            return []
+
+    async def get_agent_status(
+        self,
+        agent_id: str,
+    ) -> dict[str, Any] | None:
+        """Get status of an agent.
+
+        Args:
+            agent_id: Agent identifier.
+
+        Returns:
+            Agent status info or None if not found.
+        """
+        try:
+            response = await self._client.transport.perform_request(
+                "GET",
+                f"/_plugins/_ml/agents/{agent_id}",
+            )
+            return {
+                "agent_id": agent_id,
+                "name": response.get("name"),
+                "type": response.get("type"),
+                "description": response.get("description"),
+                "tools": [t.get("name") for t in response.get("tools", [])],
+                "created_at": response.get("created_time"),
+            }
+        except Exception as e:
+            self._logger.warning(f"Failed to get agent status: {e}")
+            return None
+
+    async def create_conversation(
+        self,
+        name: str | None = None,
+        account_id: str | None = None,
+    ) -> str | None:
+        """Create a new conversation memory.
+
+        Uses the OpenSearch Memory API to create a conversation memory.
+        The endpoint is POST /_plugins/_ml/memory (introduced in 2.12).
+
+        Args:
+            name: Optional name for the conversation.
+            account_id: Optional account ID for multi-tenancy.
+
+        Returns:
+            The new conversation/memory ID, or None if creation fails.
+        """
+        body: dict[str, Any] = {}
+        if name:
+            body["name"] = name
+        if account_id:
+            body["account_id"] = account_id
+
+        try:
+            # POST /_plugins/_ml/memory creates a new memory (OpenSearch 2.12+)
+            response = await self._client.transport.perform_request(
+                "POST",
+                "/_plugins/_ml/memory",
+                body=body if body else None,
+            )
+            memory_id = response.get("memory_id")
+            if memory_id:
+                self._logger.debug(f"Created conversation memory: {memory_id}")
+            return memory_id
+        except Exception as e:
+            self._logger.warning(f"Failed to create conversation: {e}")
+            # Return None - agent will work without pre-created memory
+            # (agent may create its own memory on first use)
+            return None
+
+    def _get_agent_id(self, agent_type: AgentType) -> str | None:
+        """Get agent ID for the specified type."""
+        if agent_type == AgentType.FLOW:
+            return self._config.flow_agent_id
+        elif agent_type == AgentType.CONVERSATIONAL:
+            return self._config.conversational_agent_id
+        return None
+
+    def _build_execute_request(
+        self,
+        query: AgenticSearchQuery,
+        index_name: str,
+    ) -> dict[str, Any]:
+        """Build agent execution request.
+
+        Only includes parameters that the agent actually uses:
+        - question: The user's query (required)
+        - memory_id: For conversation continuity (conversational agents)
+        - message_history_limit: Number of historical messages to include
+
+        Note: VectorDBTool's index and model_id are configured in the agent,
+        not passed at runtime. Extra parameters cause IllegalArgumentException.
+
+        Args:
+            query: The agentic search query.
+            index_name: Target index name (not used - agent has hardcoded index).
+
+        Returns:
+            Request body for agent execution.
+        """
+        request: dict[str, Any] = {
+            "parameters": {
+                "question": query.text,
+            }
+        }
+
+        # Add conversation context for conversational agents
+        # OpenSearch handles memory injection automatically with app_type=rag
+        if query.agent_type == AgentType.CONVERSATIONAL and query.conversation_id:
+            request["parameters"]["memory_id"] = query.conversation_id
+
+        return request
+
+    async def _execute_agent(
+        self,
+        agent_id: str,
+        body: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Execute agent and return response.
+
+        Args:
+            agent_id: The agent ID to execute.
+            body: Request body with parameters.
+
+        Returns:
+            Agent execution response.
+
+        Raises:
+            AgenticSearchError: If execution fails or times out.
+        """
+        try:
+            response = await asyncio.wait_for(
+                self._client.transport.perform_request(
+                    "POST",
+                    f"/_plugins/_ml/agents/{agent_id}/_execute",
+                    body=body,
+                ),
+                timeout=self._config.agentic_timeout_seconds,
+            )
+            return response
+        except asyncio.TimeoutError:
+            raise AgenticSearchError(
+                message="Agent execution timed out",
+                agent_id=agent_id,
+                details={"timeout_seconds": self._config.agentic_timeout_seconds},
+            )
+        except Exception as e:
+            self._logger.error(f"Agent execution failed: {e}")
+            raise AgenticSearchError(
+                message=f"Agent execution failed: {e}",
+                agent_id=agent_id,
+                cause=e,
+            )
+
+    def _parse_agentic_response(
+        self,
+        query: AgenticSearchQuery,
+        response: dict[str, Any],
+        duration_ms: float,
+    ) -> AgenticSearchResult:
+        """Parse agent response into AgenticSearchResult.
+
+        The response structure from OpenSearch ML agents:
+        {
+            "inference_results": [
+                {
+                    "output": [
+                        {"name": "response", "result": "The answer..."},
+                        {"name": "knowledge_search", "result": {...}}
+                    ]
+                }
+            ],
+            "memory_id": "...",
+            ...
+        }
+
+        Args:
+            query: The original query.
+            response: Agent execution response.
+            duration_ms: Total execution duration.
+
+        Returns:
+            Parsed AgenticSearchResult.
+        """
+        answer: str | None = None
+        reasoning_steps: list[ReasoningStep] = []
+        items: list[SearchResultItem] = []
+        conversation_id = response.get("memory_id")
+        total_tokens = 0
+        prompt_tokens = 0
+        completion_tokens = 0
+
+        # Parse inference results
+        inference_results = response.get("inference_results", [])
+        if inference_results:
+            outputs = inference_results[0].get("output", [])
+
+            for output in outputs:
+                name = output.get("name", "")
+                # Handle both direct result and dataAsMap structure (conversational agents)
+                result = output.get("result", "")
+                data_as_map = output.get("dataAsMap", {})
+                if data_as_map and "response" in data_as_map:
+                    result = data_as_map.get("response", result)
+
+                if name == "memory_id":
+                    # Extract conversation ID from conversational agent
+                    conversation_id = str(result) if result else None
+                elif name == "parent_message_id":
+                    # Track parent message ID for conversation threading
+                    pass  # Could store for future use
+                elif name in ("response", "answer_generator", "MLModelTool"):
+                    # Parse answer from output
+                    answer = self._extract_answer_from_result(result)
+
+                    # Add reasoning step for answer generation
+                    if query.include_reasoning:
+                        reasoning_steps.append(
+                            ReasoningStep(
+                                tool="MLModelTool",
+                                action="answer_generation",
+                                input=query.text,
+                                output=answer[:100] if answer else None,
+                                duration_ms=0,
+                            )
+                        )
+                elif name in ("knowledge_search", "VectorDBTool"):
+                    # Parse search results from tool output
+                    items.extend(self._parse_tool_search_results(result))
+
+                    # Add reasoning step
+                    if query.include_reasoning:
+                        reasoning_steps.append(
+                            ReasoningStep(
+                                tool="VectorDBTool",
+                                action="search",
+                                input=query.text,
+                                output=f"Found {len(items)} documents",
+                                duration_ms=0,  # Not tracked per-step
+                            )
+                        )
+
+            # Parse token usage if available
+            usage = inference_results[0].get("usage", {})
+            total_tokens = usage.get("total_tokens", 0)
+            prompt_tokens = usage.get("prompt_tokens", 0)
+            completion_tokens = usage.get("completion_tokens", 0)
+
+        # Parse agentic context for reasoning traces
+        agentic_context = response.get("agentic_context", {})
+        traces = agentic_context.get("traces", [])
+        for trace in traces:
+            if query.include_reasoning:
+                reasoning_steps.append(
+                    ReasoningStep(
+                        tool=trace.get("tool", "unknown"),
+                        action=trace.get("action", ""),
+                        input=trace.get("input"),
+                        output=trace.get("output"),
+                        duration_ms=trace.get("duration_ms", 0),
+                        tokens_used=trace.get("tokens", 0),
+                    )
+                )
+
+        # If no answer from structured output, try to get from raw response
+        if not answer and "response" in response:
+            answer = response.get("response")
+
+        # Preserve the query's conversation_id if agent didn't return one
+        # This allows multi-turn conversations when memory was created beforehand
+        final_conversation_id = conversation_id or query.conversation_id
+
+        return AgenticSearchResult(
+            query=query.text,
+            mode=SearchMode.AGENTIC,
+            items=items,
+            total_hits=len(items),
+            duration_ms=duration_ms,
+            max_score=items[0].score if items else None,
+            answer=answer,
+            reasoning_steps=reasoning_steps,
+            conversation_id=final_conversation_id,
+            agent_type=query.agent_type,
+            citations=[item.doc_id for item in items[:5]],  # Top 5 as citations
+            total_tokens=total_tokens,
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+        )
+
+    def _extract_answer_from_result(
+        self,
+        result: str | dict[str, Any],
+    ) -> str | None:
+        """Extract answer text from LLM tool result.
+
+        Handles raw OpenAI API response format:
+        {
+            "choices": [
+                {"message": {"content": "The answer..."}}
+            ]
+        }
+
+        Args:
+            result: Tool output (may be string JSON or dict).
+
+        Returns:
+            Extracted answer text or None.
+        """
+        # If it's already plain text, return it
+        if isinstance(result, str):
+            if not result.strip().startswith("{"):
+                return result
+
+            # Try to parse as JSON (OpenAI response format)
+            try:
+                result = json.loads(result)
+            except json.JSONDecodeError:
+                return result  # Return as-is if not valid JSON
+
+        if not isinstance(result, dict):
+            return str(result) if result else None
+
+        # OpenAI response format
+        choices = result.get("choices", [])
+        if choices:
+            message = choices[0].get("message", {})
+            content = message.get("content")
+            if content:
+                return content.strip()
+
+        # Fallback: look for common answer keys
+        for key in ("answer", "text", "content", "output", "result"):
+            if key in result:
+                return str(result[key]).strip()
+
+        return None
+
+    def _parse_tool_search_results(
+        self,
+        result: str | dict[str, Any] | list[Any],
+    ) -> list[SearchResultItem]:
+        """Parse search results from VectorDBTool output.
+
+        Args:
+            result: Tool output (may be string JSON or dict).
+
+        Returns:
+            List of SearchResultItem.
+        """
+        items: list[SearchResultItem] = []
+
+        # Parse if string
+        if isinstance(result, str):
+            try:
+                result = json.loads(result)
+            except json.JSONDecodeError:
+                return items
+
+        # Handle list of hits
+        if isinstance(result, list):
+            hits = result
+        elif isinstance(result, dict):
+            hits = result.get("hits", {}).get("hits", [])
+            if not hits:
+                hits = result.get("hits", [])
+        else:
+            return items
+
+        for hit in hits:
+            source = hit.get("_source", hit) if isinstance(hit, dict) else {}
+            if not source:
+                continue
+
+            items.append(
+                SearchResultItem(
+                    doc_id=hit.get("_id", source.get("id", "")),
+                    content=source.get("content", ""),
+                    score=hit.get("_score", source.get("score", 0.0)),
+                    title=source.get("title"),
+                    url=source.get("url"),
+                    source=source.get("source"),
+                    collection_id=source.get("collection_id"),
+                    source_id=source.get("source_id"),
+                    chunk_index=source.get("chunk_index"),
+                    metadata=source.get("metadata"),
+                )
+            )
+
+        return items
+
+
+class AgenticSearchFallback:
+    """Fallback handler for when agentic search fails.
+
+    Provides graceful degradation to standard hybrid search when:
+    - Agents are not configured
+    - Agent execution fails
+    - Timeout occurs
+
+    This ensures users always get results, even if the AI-powered
+    answer generation is unavailable.
+
+    Example:
+        ```python
+        agentic_searcher = OpenSearchAgenticSearcher(client, config)
+        standard_searcher = OpenSearchKnowledgeSearcher(client, config)
+        fallback = AgenticSearchFallback(agentic_searcher, standard_searcher)
+
+        # Always returns results, with or without AI answer
+        result = await fallback.search_with_fallback(query, "knowledge")
+        ```
+    """
+
+    def __init__(
+        self,
+        agentic_searcher: OpenSearchAgenticSearcher,
+        standard_searcher: Any,  # OpenSearchKnowledgeSearcher
+    ) -> None:
+        """Initialize the fallback handler.
+
+        Args:
+            agentic_searcher: Agentic search implementation.
+            standard_searcher: Standard knowledge searcher for fallback.
+        """
+        self._agentic = agentic_searcher
+        self._standard = standard_searcher
+        self._logger = logging.getLogger(__name__)
+
+    @property
+    def is_agentic_available(self) -> bool:
+        """Check if agentic search is available."""
+        return self._agentic.is_configured
+
+    async def search_with_fallback(
+        self,
+        query: AgenticSearchQuery,
+        index_name: str,
+        **options: Any,
+    ) -> AgenticSearchResult:
+        """Execute agentic search with fallback to standard search.
+
+        If agentic search is not configured or fails, automatically
+        falls back to hybrid search and wraps the results in an
+        AgenticSearchResult without an AI-generated answer.
+
+        Args:
+            query: Agentic search query.
+            index_name: Target index name.
+            **options: Additional options.
+
+        Returns:
+            AgenticSearchResult (may not have answer if in fallback mode).
+        """
+        if not self._agentic.is_configured:
+            self._logger.warning("Agentic search not configured, using fallback")
+            return await self._execute_fallback(query, index_name, "Agents not configured")
+
+        try:
+            return await self._agentic.agentic_search(query, index_name, **options)
+        except Exception as e:
+            self._logger.warning(f"Agentic search failed, falling back: {e}")
+            return await self._execute_fallback(query, index_name, str(e))
+
+    async def _execute_fallback(
+        self,
+        query: AgenticSearchQuery,
+        index_name: str,
+        reason: str,
+    ) -> AgenticSearchResult:
+        """Execute fallback search and convert to AgenticSearchResult.
+
+        Args:
+            query: Original agentic query.
+            index_name: Target index.
+            reason: Reason for fallback.
+
+        Returns:
+            AgenticSearchResult without AI answer.
+        """
+        # Convert to standard search query
+        standard_query = query.to_search_query()
+
+        # Execute standard hybrid search
+        result = await self._standard.search(standard_query, index_name)
+
+        # Convert to AgenticSearchResult
+        return AgenticSearchResult.from_search_result(
+            result,
+            answer=None,  # No AI answer in fallback
+            reasoning_steps=[
+                ReasoningStep(
+                    tool="FallbackSearch",
+                    action="hybrid_search",
+                    input=query.text,
+                    output=f"Fallback mode: {reason}. Found {result.total_hits} results.",
+                    duration_ms=result.duration_ms,
+                )
+            ],
+            agent_type=query.agent_type,
+        )
+
+
+# Type alias for protocol compliance
+AgenticSearcherImpl: type[IAgenticSearcher] = OpenSearchAgenticSearcher  # type: ignore[assignment]