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

gnosisllm_knowledge/backends/opensearch/agentic.py

@@ -2,6 +2,12 @@
 
 Uses OpenSearch ML agents for AI-powered search with reasoning capabilities.
 Supports flow agents (fast RAG) and conversational agents (multi-turn with memory).
+
+Note:
+    This module is **tenant-agnostic**. Multi-tenancy is achieved through index isolation:
+    each tenant's data resides in a separate OpenSearch index. The caller (e.g., gnosisllm-api)
+    is responsible for constructing the appropriate index name (e.g., `knowledge-{account_id}`).
+    The library operates on the provided index without any tenant-specific filtering logic.
 """
 
 from __future__ import annotations
@@ -9,7 +15,6 @@ from __future__ import annotations
 import asyncio
 import json
 import logging
-import uuid
 from datetime import UTC, datetime
 from typing import TYPE_CHECKING, Any
 
@@ -96,11 +101,12 @@ class OpenSearchAgenticSearcher:
     ) -> AgenticSearchResult:
         """Execute agentic search with agent orchestration.
 
-        The flow:
+        The flow with RAGTool:
         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
+        4. RAGTool searches the index and generates an AI answer
+        5. Parse response for answer, reasoning, and source documents
 
         Args:
             query: Agentic search query with agent type and context.
@@ -121,7 +127,7 @@ class OpenSearchAgenticSearcher:
             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."},
+                details={"hint": "Run 'gnosisllm-knowledge agentic setup --force' to configure agents."},
             )
 
         # Build execution request
@@ -133,15 +139,114 @@ class OpenSearchAgenticSearcher:
                 "agent_id": agent_id,
                 "agent_type": query.agent_type.value,
                 "query": query.text[:100],
+                "index_name": index_name,
             },
         )
 
-        # Execute agent
-        response = await self._execute_agent(agent_id, execute_body)
+        # Execute agent - RAGTool handles search AND answer generation
+        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)
+        return self._parse_rag_response(query, agent_response, duration_ms)
+
+    def _parse_rag_response(
+        self,
+        query: AgenticSearchQuery,
+        response: dict[str, Any],
+        duration_ms: float,
+    ) -> AgenticSearchResult:
+        """Parse RAGTool response into AgenticSearchResult.
+
+        RAGTool returns both an AI-generated answer and source documents.
+
+        Response format:
+        {
+            "inference_results": [
+                {
+                    "output": [
+                        {"name": "knowledge_search", "result": "<LLM answer>"}
+                    ]
+                }
+            ]
+        }
+
+        Args:
+            query: The original query.
+            response: Agent execution response.
+            duration_ms: Total execution duration.
+
+        Returns:
+            Parsed AgenticSearchResult with answer and sources.
+        """
+        answer: str | None = None
+        reasoning_steps: list[ReasoningStep] = []
+        items: list[SearchResultItem] = []
+        conversation_id = response.get("memory_id")
+
+        # 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", "")
+                result = output.get("result", "")
+
+                # Handle dataAsMap structure
+                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":
+                    conversation_id = str(result) if result else None
+                elif name in ("knowledge_search", "RAGTool", "response"):
+                    # RAGTool returns the LLM-generated answer
+                    answer = self._extract_answer_from_result(result)
+
+                    if query.include_reasoning:
+                        reasoning_steps.append(
+                            ReasoningStep(
+                                tool="RAGTool",
+                                action="rag_search",
+                                input=query.text,
+                                output=answer[:200] if answer else None,
+                                duration_ms=duration_ms,
+                            )
+                        )
+
+            # Extract source documents if available in the response
+            for output in outputs:
+                name = output.get("name", "")
+                result = output.get("result", "")
+
+                # Try to extract source documents from additional_info or similar
+                additional_info = output.get("additional_info", {})
+                if additional_info:
+                    hits = additional_info.get("hits", {})
+                    if hits:
+                        items.extend(self._parse_opensearch_hits(hits))
+
+        # If no answer from structured output, try 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
+        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]],
+        )
 
     async def get_conversation(
         self,
@@ -197,13 +302,15 @@ class OpenSearchAgenticSearcher:
 
     async def list_conversations(
         self,
-        account_id: str | None = None,
         limit: int = 100,
     ) -> list[dict[str, Any]]:
         """List active conversations.
 
+        Note:
+            This library is tenant-agnostic. Multi-tenancy is achieved through
+            index isolation (separate index per account).
+
         Args:
-            account_id: Filter by account (multi-tenant).
             limit: Maximum number of conversations.
 
         Returns:
@@ -211,8 +318,6 @@ class OpenSearchAgenticSearcher:
         """
         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",
@@ -265,16 +370,18 @@ class OpenSearchAgenticSearcher:
     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).
 
+        Note:
+            This library is tenant-agnostic. Multi-tenancy is achieved through
+            index isolation (separate index per account).
+
         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.
@@ -282,8 +389,6 @@ class OpenSearchAgenticSearcher:
         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+)
@@ -317,17 +422,18 @@ class OpenSearchAgenticSearcher:
     ) -> dict[str, Any]:
         """Build agent execution request.
 
-        Only includes parameters that the agent actually uses:
+        RAGTool requires:
         - 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.
+        The index is configured at agent creation time, not at execution time.
+        RAGTool searches the configured index and generates an AI answer.
+
+        Conversational agents also support:
+        - memory_id: For conversation continuity
 
         Args:
             query: The agentic search query.
-            index_name: Target index name (not used - agent has hardcoded index).
+            index_name: Target index name (for logging, not used by RAGTool).
 
         Returns:
             Request body for agent execution.
@@ -339,7 +445,6 @@ class OpenSearchAgenticSearcher:
         }
 
         # 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
 
@@ -386,32 +491,155 @@ class OpenSearchAgenticSearcher:
                 cause=e,
             )
 
+    def _extract_dsl_from_agent_response(
+        self,
+        response: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Extract generated DSL query from agent response.
+
+        The flow agent with QueryPlanningTool returns the DSL in the output.
+        Format: {"inference_results": [{"output": [{"name": "response", "result": "<DSL JSON>"}]}]}
+
+        Args:
+            response: Agent execution response.
+
+        Returns:
+            Parsed DSL query dict, or None if not found.
+        """
+        try:
+            inference_results = response.get("inference_results", [])
+            if not inference_results:
+                return None
+
+            outputs = inference_results[0].get("output", [])
+            for output in outputs:
+                name = output.get("name", "")
+                result = output.get("result", "")
+
+                # QueryPlanningTool outputs come as "response" or "query_planner"
+                if name in ("response", "query_planner", "QueryPlanningTool"):
+                    if isinstance(result, dict):
+                        return result
+                    if isinstance(result, str) and result.strip():
+                        # Try to parse as JSON
+                        return self._parse_dsl_string(result)
+
+            return None
+        except Exception as e:
+            self._logger.warning(f"Failed to extract DSL from agent response: {e}")
+            return None
+
+    def _parse_dsl_string(self, dsl_string: str) -> dict[str, Any] | None:
+        """Parse a DSL query string into a dictionary.
+
+        Handles various formats:
+        - Raw JSON
+        - Markdown code blocks
+        - JSON with surrounding text
+
+        Args:
+            dsl_string: The DSL query as a string.
+
+        Returns:
+            Parsed DSL query dict, or None if parsing fails.
+        """
+        dsl_string = dsl_string.strip()
+
+        # Remove markdown code blocks if present
+        if dsl_string.startswith("```"):
+            lines = dsl_string.split("\n")
+            # Remove first line (```json or ```)
+            lines = lines[1:] if lines else []
+            # Remove last line (```)
+            if lines and lines[-1].strip() == "```":
+                lines = lines[:-1]
+            dsl_string = "\n".join(lines).strip()
+
+        # Try to find and parse JSON
+        try:
+            # Find the first { and last }
+            start = dsl_string.find("{")
+            end = dsl_string.rfind("}") + 1
+            if start >= 0 and end > start:
+                json_str = dsl_string[start:end]
+                return json.loads(json_str)
+        except json.JSONDecodeError as e:
+            self._logger.debug(f"Failed to parse DSL JSON: {e}")
+
+        # Try parsing the whole string
+        try:
+            return json.loads(dsl_string)
+        except json.JSONDecodeError:
+            pass
+
+        return None
+
+    async def _execute_dsl_query(
+        self,
+        index_name: str,
+        dsl_query: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Execute a DSL query against the index.
+
+        Args:
+            index_name: Target index name.
+            dsl_query: OpenSearch DSL query to execute.
+
+        Returns:
+            Search response with hits.
+
+        Raises:
+            AgenticSearchError: If query execution fails.
+        """
+        try:
+            response = await self._client.search(
+                index=index_name,
+                body=dsl_query,
+            )
+            return response
+        except Exception as e:
+            self._logger.error(f"DSL query execution failed: {e}")
+            raise AgenticSearchError(
+                message=f"Failed to execute generated DSL query: {e}",
+                details={"index_name": index_name, "query": str(dsl_query)[:200]},
+                cause=e,
+            )
+
     def _parse_agentic_response(
         self,
         query: AgenticSearchQuery,
-        response: dict[str, Any],
+        agent_response: dict[str, Any],
         duration_ms: float,
+        search_response: dict[str, Any] | None = None,
+        generated_dsl: dict[str, Any] | None = None,
     ) -> 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": "...",
-            ...
-        }
+        Supports two response formats:
+
+        1. QueryPlanningTool (OpenSearch 3.2+):
+           The agent generates DSL queries which we then execute.
+           Agent response: {"inference_results": [{"output": [{"name": "response", "result": "<DSL JSON>"}]}]}
+           Search response: Standard OpenSearch search response with hits.
+
+        2. Legacy VectorDBTool + MLModelTool:
+           {
+               "inference_results": [
+                   {
+                       "output": [
+                           {"name": "knowledge_search", "result": {...}},
+                           {"name": "answer_generator", "result": "..."}
+                       ]
+                   }
+               ]
+           }
 
         Args:
             query: The original query.
-            response: Agent execution response.
+            agent_response: Agent execution response.
             duration_ms: Total execution duration.
+            search_response: Search results from executing the generated DSL (optional).
+            generated_dsl: The DSL query generated by the agent (optional).
 
         Returns:
             Parsed AgenticSearchResult.
@@ -419,13 +647,40 @@ class OpenSearchAgenticSearcher:
         answer: str | None = None
         reasoning_steps: list[ReasoningStep] = []
         items: list[SearchResultItem] = []
-        conversation_id = response.get("memory_id")
+        conversation_id = agent_response.get("memory_id")
+        dsl_string: str | None = None
         total_tokens = 0
         prompt_tokens = 0
         completion_tokens = 0
 
-        # Parse inference results
-        inference_results = response.get("inference_results", [])
+        # Parse search results from executed DSL query first (QueryPlanningTool flow)
+        if search_response:
+            hits_data = search_response.get("hits", {})
+            items.extend(self._parse_opensearch_hits(hits_data))
+
+            if query.include_reasoning:
+                dsl_string = json.dumps(generated_dsl) if generated_dsl else None
+                reasoning_steps.append(
+                    ReasoningStep(
+                        tool="QueryPlanningTool",
+                        action="query_generation",
+                        input=query.text,
+                        output=dsl_string[:200] if dsl_string else None,
+                        duration_ms=0,
+                    )
+                )
+                reasoning_steps.append(
+                    ReasoningStep(
+                        tool="QueryPlanningTool",
+                        action="search_execution",
+                        input=dsl_string[:100] if dsl_string else query.text,
+                        output=f"Found {len(items)} documents",
+                        duration_ms=0,
+                    )
+                )
+
+        # Parse inference results for additional outputs (legacy or conversational)
+        inference_results = agent_response.get("inference_results", [])
         if inference_results:
             outputs = inference_results[0].get("output", [])
 
@@ -443,8 +698,8 @@ class OpenSearchAgenticSearcher:
                 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
+                elif name in ("answer_generator", "MLModelTool"):
+                    # Parse answer from output (legacy format)
                     answer = self._extract_answer_from_result(result)
 
                     # Add reasoning step for answer generation
@@ -459,7 +714,7 @@ class OpenSearchAgenticSearcher:
                             )
                         )
                 elif name in ("knowledge_search", "VectorDBTool"):
-                    # Parse search results from tool output
+                    # Parse search results from legacy VectorDBTool output
                     items.extend(self._parse_tool_search_results(result))
 
                     # Add reasoning step
@@ -470,9 +725,10 @@ class OpenSearchAgenticSearcher:
                                 action="search",
                                 input=query.text,
                                 output=f"Found {len(items)} documents",
-                                duration_ms=0,  # Not tracked per-step
+                                duration_ms=0,
                             )
                         )
+                # Skip "response" and "query_planner" here - they're handled via generated_dsl parameter
 
             # Parse token usage if available
             usage = inference_results[0].get("usage", {})
@@ -480,8 +736,8 @@ class OpenSearchAgenticSearcher:
             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", {})
+        # Parse agentic context for reasoning traces (if present)
+        agentic_context = agent_response.get("agentic_context", {})
         traces = agentic_context.get("traces", [])
         for trace in traces:
             if query.include_reasoning:
@@ -497,8 +753,8 @@ class OpenSearchAgenticSearcher:
                 )
 
         # If no answer from structured output, try to get from raw response
-        if not answer and "response" in response:
-            answer = response.get("response")
+        if not answer and "response" in agent_response:
+            answer = agent_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
@@ -519,6 +775,7 @@ class OpenSearchAgenticSearcher:
             total_tokens=total_tokens,
             prompt_tokens=prompt_tokens,
             completion_tokens=completion_tokens,
+            generated_query=dsl_string,  # Include the generated DSL for debugging
         )
 
     def _extract_answer_from_result(
@@ -622,6 +879,56 @@ class OpenSearchAgenticSearcher:
 
         return items
 
+    def _parse_opensearch_hits(
+        self,
+        hits_data: dict[str, Any],
+    ) -> list[SearchResultItem]:
+        """Parse OpenSearch hits structure into SearchResultItems.
+
+        Standard OpenSearch response format:
+        {
+            "total": {"value": 10},
+            "max_score": 1.0,
+            "hits": [
+                {"_id": "...", "_score": 0.9, "_source": {...}}
+            ]
+        }
+
+        Args:
+            hits_data: OpenSearch hits object.
+
+        Returns:
+            List of SearchResultItem.
+        """
+        items: list[SearchResultItem] = []
+
+        hits = hits_data.get("hits", [])
+        for hit in hits:
+            if not isinstance(hit, dict):
+                continue
+
+            source = hit.get("_source", {})
+            if not source:
+                continue
+
+            items.append(
+                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"),
+                    source_id=source.get("source_id"),
+                    chunk_index=source.get("chunk_index"),
+                    total_chunks=source.get("total_chunks"),
+                    metadata=source.get("metadata"),
+                )
+            )
+
+        return items
+
 
 class AgenticSearchFallback:
     """Fallback handler for when agentic search fails.