aiecs 1.5.1__py3-none-any.whl

aiecs/domain/agent/knowledge_aware_agent.py

@@ -0,0 +1,646 @@
+"""
+Knowledge-Aware Agent
+
+Enhanced hybrid agent with knowledge graph integration.
+Extends the standard HybridAgent with graph reasoning capabilities.
+"""
+
+import logging
+from typing import Dict, List, Any, Optional
+from datetime import datetime
+
+from aiecs.llm import BaseLLMClient
+from aiecs.infrastructure.graph_storage.base import GraphStore
+from aiecs.tools.knowledge_graph import GraphReasoningTool
+from aiecs.domain.knowledge_graph.models.entity import Entity
+
+from .hybrid_agent import HybridAgent
+from .models import AgentConfiguration
+
+logger = logging.getLogger(__name__)
+
+
+class KnowledgeAwareAgent(HybridAgent):
+    """
+    Knowledge-Aware Agent with integrated knowledge graph reasoning.
+
+    Extends HybridAgent with:
+    - Knowledge graph consultation during reasoning
+    - Graph-aware tool selection
+    - Knowledge-augmented prompt construction
+    - Automatic access to graph reasoning capabilities
+
+    Example:
+        ```python
+        from aiecs.domain.agent import KnowledgeAwareAgent
+        from aiecs.infrastructure.graph_storage import InMemoryGraphStore
+
+        # Initialize with knowledge graph
+        graph_store = InMemoryGraphStore()
+        await graph_store.initialize()
+
+        agent = KnowledgeAwareAgent(
+            agent_id="kg_agent_001",
+            name="Knowledge Assistant",
+            llm_client=llm_client,
+            tools=["web_search", "calculator"],
+            config=config,
+            graph_store=graph_store
+        )
+
+        await agent.initialize()
+        result = await agent.execute_task("How is Alice connected to Company X?")
+        ```
+    """
+
+    def __init__(
+        self,
+        agent_id: str,
+        name: str,
+        llm_client: BaseLLMClient,
+        tools: List[str],
+        config: AgentConfiguration,
+        graph_store: Optional[GraphStore] = None,
+        description: Optional[str] = None,
+        version: str = "1.0.0",
+        max_iterations: int = 10,
+        enable_graph_reasoning: bool = True,
+    ):
+        """
+        Initialize Knowledge-Aware agent.
+
+        Args:
+            agent_id: Unique agent identifier
+            name: Agent name
+            llm_client: LLM client for reasoning
+            tools: List of tool names (graph_reasoning auto-added if graph_store provided)
+            config: Agent configuration
+            graph_store: Optional knowledge graph store
+            description: Optional description
+            version: Agent version
+            max_iterations: Maximum ReAct iterations
+            enable_graph_reasoning: Whether to enable graph reasoning capabilities
+        """
+        # Auto-add graph_reasoning tool if graph_store is provided
+        if graph_store is not None and enable_graph_reasoning:
+            if "graph_reasoning" not in tools:
+                tools = tools + ["graph_reasoning"]
+
+        super().__init__(
+            agent_id=agent_id,
+            name=name,
+            llm_client=llm_client,
+            tools=tools,
+            config=config,
+            description=description or "Knowledge-aware agent with integrated graph reasoning",
+            version=version,
+            max_iterations=max_iterations,
+        )
+
+        self.graph_store = graph_store
+        self.enable_graph_reasoning = enable_graph_reasoning
+        self._graph_reasoning_tool: Optional[GraphReasoningTool] = None
+        self._knowledge_context: Dict[str, Any] = {}
+
+        logger.info(
+            f"KnowledgeAwareAgent initialized: {agent_id} "
+            f"with graph_store={'enabled' if graph_store else 'disabled'}"
+        )
+
+    async def _initialize(self) -> None:
+        """Initialize Knowledge-Aware agent - setup graph tools and augmented prompts."""
+        # Call parent initialization
+        await super()._initialize()
+
+        # Initialize graph reasoning tool if graph store is available
+        if self.graph_store is not None and self.enable_graph_reasoning:
+            try:
+                self._graph_reasoning_tool = GraphReasoningTool(self.graph_store)
+                logger.info(f"KnowledgeAwareAgent {self.agent_id} initialized graph reasoning")
+            except Exception as e:
+                logger.warning(f"Failed to initialize graph reasoning tool: {e}")
+
+        # Rebuild system prompt with knowledge graph capabilities
+        if self.graph_store is not None:
+            self._system_prompt = self._build_kg_augmented_system_prompt()
+
+        logger.info(f"KnowledgeAwareAgent {self.agent_id} initialized with enhanced capabilities")
+
+    async def _shutdown(self) -> None:
+        """Shutdown Knowledge-Aware agent."""
+        # Clear knowledge context
+        self._knowledge_context.clear()
+
+        # Shutdown graph store if needed
+        if self.graph_store is not None:
+            try:
+                await self.graph_store.close()
+            except Exception as e:
+                logger.warning(f"Error closing graph store: {e}")
+
+        # Call parent shutdown
+        await super()._shutdown()
+
+        logger.info(f"KnowledgeAwareAgent {self.agent_id} shut down")
+
+    def _build_kg_augmented_system_prompt(self) -> str:
+        """
+        Build knowledge graph-augmented system prompt.
+
+        Returns:
+            Enhanced system prompt with KG capabilities
+        """
+        base_prompt = super()._build_system_prompt()
+
+        # Add knowledge graph capabilities section
+        kg_section = """
+
+KNOWLEDGE GRAPH CAPABILITIES:
+You have access to an integrated knowledge graph that can help answer complex questions.
+
+REASONING WITH KNOWLEDGE:
+Your reasoning process now includes an automatic RETRIEVE phase:
+1. RETRIEVE: Relevant knowledge is automatically fetched from the graph before each reasoning step
+2. THOUGHT: You analyze the task considering retrieved knowledge
+3. ACTION: Use tools or provide final answer
+4. OBSERVATION: Review results and continue
+
+Retrieved knowledge will be provided as:
+RETRIEVED KNOWLEDGE:
+- Entity: id (properties)
+- Entity: id (properties)
+...
+
+When to use the 'graph_reasoning' tool:
+- Multi-hop questions (e.g., "How is X connected to Y?")
+- Relationship discovery (e.g., "Who knows people at Company Z?")
+- Knowledge completion (e.g., "What do we know about Person A?")
+- Evidence-based reasoning (multiple sources needed)
+
+The 'graph_reasoning' tool supports these modes:
+- query_plan: Plan complex query execution
+- multi_hop: Find connections between entities
+- inference: Apply logical inference rules
+- full_reasoning: Complete reasoning pipeline with evidence synthesis
+
+Use graph reasoning proactively when questions involve:
+- Connections, relationships, or paths
+- Multiple entities or complex queries
+- Need for evidence from multiple sources
+"""
+
+        return base_prompt + kg_section
+
+    async def _reason_with_graph(
+        self, query: str, context: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """
+        Consult knowledge graph during reasoning.
+
+        Args:
+            query: Query to reason about
+            context: Optional context for reasoning
+
+        Returns:
+            Reasoning results from knowledge graph
+        """
+        if self._graph_reasoning_tool is None:
+            logger.warning("Graph reasoning tool not available")
+            return {"error": "Graph reasoning not available"}
+
+        try:
+            # Use multi_hop mode by default for general queries
+            from aiecs.tools.knowledge_graph.graph_reasoning_tool import (
+                GraphReasoningInput,
+            )
+
+            # Extract entity IDs from context if available
+            start_entity_id = None
+            target_entity_id = None
+            if context:
+                start_entity_id = context.get("start_entity_id")
+                target_entity_id = context.get("target_entity_id")
+
+            input_data = GraphReasoningInput(
+                mode="multi_hop",
+                query=query,
+                start_entity_id=start_entity_id,
+                target_entity_id=target_entity_id,
+                max_hops=3,
+                synthesize_evidence=True,
+                confidence_threshold=0.6,
+            )
+
+            result = await self._graph_reasoning_tool._execute(input_data)
+
+            # Store knowledge context for later use
+            self._knowledge_context[query] = {
+                "answer": result.get("answer"),
+                "confidence": result.get("confidence"),
+                "evidence_count": result.get("evidence_count"),
+                "timestamp": datetime.utcnow().isoformat(),
+            }
+
+            return result
+
+        except Exception as e:
+            logger.error(f"Error in graph reasoning: {e}")
+            return {"error": str(e)}
+
+    async def _select_tools_with_graph_awareness(
+        self, task: str, available_tools: List[str]
+    ) -> List[str]:
+        """
+        Select tools with graph awareness.
+
+        Prioritizes graph reasoning tool for knowledge-related queries.
+
+        Args:
+            task: Task description
+            available_tools: Available tool names
+
+        Returns:
+            Selected tool names
+        """
+        # Keywords that suggest graph reasoning might be useful
+        graph_keywords = [
+            "connected",
+            "connection",
+            "relationship",
+            "related",
+            "knows",
+            "works",
+            "friend",
+            "colleague",
+            "partner",
+            "how",
+            "why",
+            "who",
+            "what",
+            "which",
+            "find",
+            "discover",
+            "explore",
+            "trace",
+        ]
+
+        task_lower = task.lower()
+
+        # Check if task involves knowledge graph queries
+        uses_graph_keywords = any(keyword in task_lower for keyword in graph_keywords)
+
+        # If graph reasoning is available and task seems graph-related,
+        # prioritize it
+        if uses_graph_keywords and "graph_reasoning" in available_tools:
+            # Put graph_reasoning first
+            selected = ["graph_reasoning"]
+            # Add other tools
+            selected.extend([t for t in available_tools if t != "graph_reasoning"])
+            return selected
+
+        return available_tools
+
+    async def _augment_prompt_with_knowledge(
+        self, task: str, context: Optional[Dict[str, Any]] = None
+    ) -> str:
+        """
+        Augment prompt with relevant knowledge from graph.
+
+        Args:
+            task: Original task
+            context: Optional context
+
+        Returns:
+            Augmented task with knowledge context
+        """
+        if self.graph_store is None or not self.enable_graph_reasoning:
+            return task
+
+        # Check if we have cached knowledge for similar queries
+        relevant_knowledge = []
+        for query, kg_context in self._knowledge_context.items():
+            # Simple keyword matching (could be enhanced with embeddings)
+            if any(word in task.lower() for word in query.lower().split()):
+                relevant_knowledge.append(
+                    f"- {query}: {kg_context['answer']} (confidence: {kg_context['confidence']:.2f})"
+                )
+
+        if relevant_knowledge:
+            knowledge_section = "\n\nRELEVANT KNOWLEDGE FROM GRAPH:\n" + "\n".join(
+                relevant_knowledge[:3]
+            )
+            return task + knowledge_section
+
+        return task
+
+    async def execute_task(self, task: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Execute task with knowledge graph augmentation.
+
+        Uses knowledge-augmented ReAct loop that includes a RETRIEVE phase.
+
+        Args:
+            task: Task specification with 'description' or 'prompt'
+            context: Execution context
+
+        Returns:
+            Task execution result
+        """
+        # Extract task description
+        task_description = task.get("description") or task.get("prompt") or task.get("task")
+        if not task_description:
+            return await super().execute_task(task, context)
+
+        # Augment task with knowledge if available
+        augmented_task_desc = await self._augment_prompt_with_knowledge(task_description, context)
+
+        # If task seems graph-related, consult graph first
+        if self.graph_store is not None and self.enable_graph_reasoning:
+            # Check if this is a direct graph query
+            graph_keywords = [
+                "connected",
+                "connection",
+                "relationship",
+                "knows",
+                "works at",
+            ]
+            if any(keyword in task_description.lower() for keyword in graph_keywords):
+                logger.info(f"Consulting knowledge graph for task: {task_description}")
+
+                # Try graph reasoning
+                graph_result = await self._reason_with_graph(augmented_task_desc, context)
+
+                # If we got a good answer from the graph, use it
+                if "answer" in graph_result and graph_result.get("confidence", 0) > 0.7:
+                    return {
+                        "success": True,
+                        "output": graph_result["answer"],
+                        "confidence": graph_result["confidence"],
+                        "source": "knowledge_graph",
+                        "evidence_count": graph_result.get("evidence_count", 0),
+                        "reasoning_trace": graph_result.get("reasoning_trace", []),
+                        "timestamp": datetime.utcnow().isoformat(),
+                    }
+
+        # Fall back to standard hybrid agent execution
+        # This will use the overridden _react_loop with knowledge retrieval
+        # Create modified task dict with augmented description
+        augmented_task = task.copy()
+        if "description" in task:
+            augmented_task["description"] = augmented_task_desc
+        elif "prompt" in task:
+            augmented_task["prompt"] = augmented_task_desc
+        elif "task" in task:
+            augmented_task["task"] = augmented_task_desc
+
+        return await super().execute_task(augmented_task, context)
+
+    async def _react_loop(self, task: str, context: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Execute knowledge-augmented ReAct loop: Retrieve → Reason → Act → Observe.
+
+        Extends the standard ReAct loop with a RETRIEVE phase that fetches
+        relevant knowledge from the graph before each reasoning step.
+
+        Args:
+            task: Task description
+            context: Context dictionary
+
+        Returns:
+            Result dictionary with 'final_answer', 'steps', 'iterations'
+        """
+        steps = []
+        tool_calls_count = 0
+        total_tokens = 0
+        knowledge_retrievals = 0
+
+        # Build initial messages
+        from aiecs.llm import LLMMessage
+
+        messages = self._build_initial_messages(task, context)
+
+        for iteration in range(self._max_iterations):
+            logger.debug(f"KnowledgeAwareAgent {self.agent_id} - ReAct iteration {iteration + 1}")
+
+            # RETRIEVE: Get relevant knowledge from graph (if enabled)
+            retrieved_knowledge = []
+            if self.graph_store is not None and self.enable_graph_reasoning:
+                try:
+                    retrieved_knowledge = await self._retrieve_relevant_knowledge(
+                        task, context, iteration
+                    )
+
+                    if retrieved_knowledge:
+                        knowledge_retrievals += 1
+                        knowledge_str = self._format_retrieved_knowledge(retrieved_knowledge)
+
+                        steps.append(
+                            {
+                                "type": "retrieve",
+                                "knowledge_count": len(retrieved_knowledge),
+                                "content": (
+                                    knowledge_str[:200] + "..."
+                                    if len(knowledge_str) > 200
+                                    else knowledge_str
+                                ),
+                                "iteration": iteration + 1,
+                            }
+                        )
+
+                        # Add knowledge to messages
+                        messages.append(
+                            LLMMessage(
+                                role="system",
+                                content=f"RETRIEVED KNOWLEDGE:\n{knowledge_str}",
+                            )
+                        )
+                except Exception as e:
+                    logger.warning(f"Knowledge retrieval failed: {e}")
+
+            # THINK: LLM reasons about next action
+            response = await self.llm_client.generate_text(
+                messages=messages,
+                model=self._config.llm_model,
+                temperature=self._config.temperature,
+                max_tokens=self._config.max_tokens,
+            )
+
+            thought = response.content
+            total_tokens += getattr(response, "total_tokens", 0)
+
+            steps.append(
+                {
+                    "type": "thought",
+                    "content": thought,
+                    "iteration": iteration + 1,
+                }
+            )
+
+            # Check if final answer
+            if "FINAL ANSWER:" in thought:
+                final_answer = self._extract_final_answer(thought)
+                return {
+                    "final_answer": final_answer,
+                    "steps": steps,
+                    "iterations": iteration + 1,
+                    "tool_calls_count": tool_calls_count,
+                    "knowledge_retrievals": knowledge_retrievals,
+                    "total_tokens": total_tokens,
+                }
+
+            # Check if tool call
+            if "TOOL:" in thought:
+                # ACT: Execute tool
+                try:
+                    tool_info = self._parse_tool_call(thought)
+                    tool_result = await self._execute_tool(
+                        tool_info["tool"],
+                        tool_info.get("operation"),
+                        tool_info.get("parameters", {}),
+                    )
+                    tool_calls_count += 1
+
+                    steps.append(
+                        {
+                            "type": "action",
+                            "tool": tool_info["tool"],
+                            "operation": tool_info.get("operation"),
+                            "parameters": tool_info.get("parameters"),
+                            "iteration": iteration + 1,
+                        }
+                    )
+
+                    # OBSERVE: Add tool result to conversation
+                    observation = f"OBSERVATION: Tool '{tool_info['tool']}' returned: {tool_result}"
+                    steps.append(
+                        {
+                            "type": "observation",
+                            "content": observation,
+                            "iteration": iteration + 1,
+                        }
+                    )
+
+                    # Add to messages for next iteration
+                    messages.append(LLMMessage(role="assistant", content=thought))
+                    messages.append(LLMMessage(role="user", content=observation))
+
+                except Exception as e:
+                    error_msg = f"OBSERVATION: Tool execution failed: {str(e)}"
+                    steps.append(
+                        {
+                            "type": "observation",
+                            "content": error_msg,
+                            "iteration": iteration + 1,
+                            "error": True,
+                        }
+                    )
+                    messages.append(LLMMessage(role="assistant", content=thought))
+                    messages.append(LLMMessage(role="user", content=error_msg))
+
+            else:
+                # LLM didn't provide clear action - treat as final answer
+                return {
+                    "final_answer": thought,
+                    "steps": steps,
+                    "iterations": iteration + 1,
+                    "tool_calls_count": tool_calls_count,
+                    "knowledge_retrievals": knowledge_retrievals,
+                    "total_tokens": total_tokens,
+                }
+
+        # Max iterations reached
+        logger.warning(f"KnowledgeAwareAgent {self.agent_id} reached max iterations")
+        return {
+            "final_answer": "Max iterations reached. Unable to complete task fully.",
+            "steps": steps,
+            "iterations": self._max_iterations,
+            "tool_calls_count": tool_calls_count,
+            "knowledge_retrievals": knowledge_retrievals,
+            "total_tokens": total_tokens,
+            "max_iterations_reached": True,
+        }
+
+    async def _retrieve_relevant_knowledge(
+        self, task: str, context: Dict[str, Any], iteration: int
+    ) -> List[Entity]:
+        """
+        Retrieve relevant knowledge for the current reasoning step.
+
+        Args:
+            task: Task description
+            context: Context dictionary
+            iteration: Current iteration number
+
+        Returns:
+            List of relevant entities
+        """
+        # Extract entity IDs or types from context
+        context.get("entity_types")
+        context.get("session_id", f"temp_{self.agent_id}")
+
+        # Try to retrieve knowledge
+        # For now, use a simple approach - could be enhanced with embeddings
+        try:
+            # Use vector search if query is provided (simplified)
+            # In production, this would generate embeddings for the task
+
+            # For iteration 0, retrieve general context
+            # For later iterations, retrieve more specific knowledge
+            # limit = 5 if iteration == 0 else 3  # Reserved for future use
+
+            if hasattr(self, "_knowledge_context") and self._knowledge_context:
+                # Get entities mentioned in previous knowledge
+                for _, kg_ctx in self._knowledge_context.items():
+                    # This is simplified - in production would extract entity
+                    # IDs properly
+                    pass
+
+            # Placeholder: Return empty for now
+            # In a full implementation, this would:
+            # 1. Generate embedding for task
+            # 2. Use vector_search on graph_store
+            # 3. Filter by relevance
+            # 4. Return top-k results
+
+            return []
+
+        except Exception as e:
+            logger.error(f"Error retrieving knowledge: {e}")
+            return []
+
+    def _format_retrieved_knowledge(self, entities: List[Entity]) -> str:
+        """
+        Format retrieved knowledge entities for inclusion in prompt.
+
+        Args:
+            entities: List of entities retrieved from graph
+
+        Returns:
+            Formatted knowledge string
+        """
+        if not entities:
+            return ""
+
+        lines = []
+        for entity in entities:
+            entity_str = f"- {entity.entity_type}: {entity.id}"
+            if entity.properties:
+                props_str = ", ".join(f"{k}={v}" for k, v in entity.properties.items())
+                entity_str += f" ({props_str})"
+            lines.append(entity_str)
+
+        return "\n".join(lines)
+
+    def get_knowledge_context(self) -> Dict[str, Any]:
+        """
+        Get accumulated knowledge context.
+
+        Returns:
+            Dictionary of accumulated knowledge
+        """
+        return self._knowledge_context.copy()
+
+    def clear_knowledge_context(self) -> None:
+        """Clear accumulated knowledge context."""
+        self._knowledge_context.clear()
+        logger.debug(f"Cleared knowledge context for agent {self.agent_id}")