""" Agent-Specific GraphRAG Integration Service. Wraps GraphRAGEngine with agent-specific context retrieval, fail-fast validation, and instance-level relationship checking to prevent agent hallucination. """ import logging from typing import Dict, Any, List, Optional from sqlalchemy.orm import Session from core.graphrag_engine import GraphRAGEngine from core.models import GraphNode, GraphEdge logger = logging.getLogger(__name__) class AgentGraphRAGService: """ Agent-specific GraphRAG integration service. Provides: - Context retrieval with agent-specific logging - Instance-level relationship validation - Fail-fast validation for empty results - Performance-optimized result limits """ def __init__(self, db: Session, workspace_id: str, agent_id: str): """ Initialize agent GraphRAG service. Args: db: SQLAlchemy database session workspace_id: Workspace ID for isolation (Upstream primary key) agent_id: Agent ID for audit logging """ self.db = db self.workspace_id = workspace_id self.agent_id = agent_id self.graphrag = GraphRAGEngine() async def get_agent_context( self, query: str, mode: str = 'auto', max_entities: int = 15, max_relationships: int = 25 ) -> Dict[str, Any]: """Get GraphRAG context for agent reasoning.""" logger.info(f"Agent {self.agent_id} requesting GraphRAG context: {query[:50]}...") # Query GraphRAG (workspace-isolated) result = self.graphrag.query( workspace_id=self.workspace_id, query=query, mode=mode ) if result.get("mode") == "local": entities = result.get("entities", []) relationships = result.get("relationships", []) if not entities and not relationships: raise ValueError(f"GraphRAG validation failed: No entities found for query '{query}'") result["entities"] = entities[:max_entities] result["relationships"] = relationships[:max_relationships] elif result.get("mode") == "global": answer = result.get("answer", "") if not answer or answer.strip() == "": raise ValueError(f"GraphRAG global search failed: No community summaries for '{query}'") result["agent_id"] = self.agent_id result["has_results"] = True result["context"] = self._format_context(result) return result async def validate_entity_relationship( self, entity_name_a: str, entity_name_b: str, relationship_type: Optional[str] = None ) -> Dict[str, Any]: """Validate that a relationship exists between two entity instances.""" logger.info(f"Agent {self.agent_id} validating relationship: {entity_name_a} -> {entity_name_b}") node_a = self.db.query(GraphNode).filter( GraphNode.workspace_id == self.workspace_id, GraphNode.name == entity_name_a ).first() node_b = self.db.query(GraphNode).filter( GraphNode.workspace_id == self.workspace_id, GraphNode.name == entity_name_b ).first() if not node_a or not node_b: raise ValueError(f"Entities not found in GraphRAG for validation.") query = self.db.query(GraphEdge).filter( GraphEdge.workspace_id == self.workspace_id, GraphEdge.source_node_id == node_a.id, GraphEdge.target_node_id == node_b.id ) if relationship_type: query = query.filter(GraphEdge.relationship_type == relationship_type) edge = query.first() if not edge: raise ValueError(f"No relationship found between '{entity_name_a}' and '{entity_name_b}'.") properties = edge.properties or {} return { "exists": True, "relationship_type": edge.relationship_type, "description": properties.get('description', f"{entity_name_a} -> {entity_name_b}"), "weight": edge.weight, "metadata": properties } def _format_context(self, result: Dict[str, Any]) -> str: """Format GraphRAG result as context string.""" if result.get("mode") == "global": return f"Global Context: {result.get('answer', '')}" entities = result.get("entities", []) relationships = result.get("relationships", []) id_to_name = {e['id']: e['name'] for e in entities} lines = [f"Found {len(entities)} relevant entities:"] for e in entities[:15]: lines.append(f"- {e['name']} ({e['type']}): {e.get('description', '')}") lines.append(f"\n{len(relationships)} relationships:") for r in relationships[:25]: from_name = id_to_name.get(r['from'], r['from']) to_name = id_to_name.get(r['to'], r['to']) lines.append(f"- {from_name} -> {to_name} ({r.get('type', 'related')})") return "\n".join(lines)