ad

agentgraph/methods/production/openai_structured_extractor.py

@@ -28,8 +28,8 @@ class KnowledgeGraphReasoning(BaseModel):
     steps: List[Step]
     final_answer: KnowledgeGraph
 
-# Load environment variables from root directory
-load_dotenv('/Users/zekunwu/Desktop/agent_monitoring/.env')
+# Load environment variables (searches current directory and parents)
+load_dotenv()
 
 # Configure logging
 logger = logging.getLogger(__name__)
@@ -40,18 +40,14 @@ class OpenAIStructuredExtractor:
     Simple knowledge graph extractor using OpenAI's structured outputs.
     """
     
-    def __init__(self, model: str = "gpt-5-mini", *, compatibility_mode: bool = False, allow_required_by: bool = False):
+    def __init__(self, model: str = "gpt-5-mini"):
         """
         Initialize the extractor.
         
         Args:
             model: OpenAI model to use (must support structured outputs)
-            compatibility_mode: Enable light post-processing to better align with sample/reference KGs
-            allow_required_by: When in compatibility mode, also add REQUIRED_BY Tool→Task relation when appropriate
         """
         self.model = model
-        self.compatibility_mode = compatibility_mode
-        self.allow_required_by = allow_required_by
         self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
         logger.info(f"OpenAI Structured Extractor initialized with model: {model}")
     
@@ -83,8 +79,10 @@ ENTITY TYPES & PRIORITIES:
 RELATION PRIORITIES (ULTRA-SIMPLIFIED):
 - PERFORMS (Agent→Task): ONLY agent-task relation needed
 - Input→Agent→Task→Output→Human: Essential workflow chain
-- NO COMPLEX RELATIONS: Avoid ASSIGNED_TO, INTERVENES, REQUIRED_BY
-- TARGET: 6-8 total relations maximum (keep it simple!)
+- USES (Agent→Tool): Include exactly one if a Tool exists and is used
+- REQUIRED_BY (Tool→Task): Conditionally allowed when the task explicitly depends on a Tool (e.g., execution on Computer Terminal). Use at most one.
+- FORBIDDEN: ASSIGNED_TO, INTERVENES
+- TARGET: 6-10 total relations maximum (simple but complete)
 
 WORKFLOW PATTERNS:
 - Simple (1-2 agents): Single consolidated task, basic relations
@@ -101,6 +99,13 @@ OUTPUT REQUIREMENTS:
 - Empty raw_prompt/interaction_prompt fields
 - Include 1-2 failures and optimizations
 
+NAMING STYLE & CANONICALIZATION:
+- Title Case with spaces (avoid underscores in names)
+- Prefer canonical names where applicable: "Computer Terminal" (not Computer_terminal)
+- Domain templates (use when content clearly matches):
+  * Location-based discovery: Tasks → Geographic Proximity Analysis → Restaurant Data Collection → Operating Hours Validation; Output → Restaurant Recommendations; System → Location-Based Restaurant Discovery System
+  * Wildlife data analysis: Tasks → Government Dataset Processing → Population Statistics Calculation → Ecological Data Validation; Human → Wildlife Researcher; System → Wildlife Data Analysis and Ecological Monitoring System
+
 == FEW-SHOT EXAMPLE (COMPLEX DISCOVERY WORKFLOW) ==
 
 Input Query: "What is the closest eatery to Harkness Memorial State Park open at 11pm on Wednesdays?"
@@ -195,7 +200,7 @@ ANALYSIS STEPS:
    * Task sharing = IMMEDIATE FAILURE - completely unacceptable
    * ALWAYS decompose into independent subtasks for each agent
 
-4. RELATION MAPPING (FORCED 1:1 MAPPING):
+4. RELATION MAPPING (STRICT BUT COMPLETE):
    - PERFORMS: Each agent performs EXACTLY ONE UNIQUE task
      * 3 agents = 3 different tasks = 3 PERFORMS relations
      * agent_001 → task_001, agent_002 → task_002, agent_003 → task_003
@@ -207,10 +212,11 @@ ANALYSIS STEPS:
      * Task→Task (NEXT): For sequential workflows only
      * Last Task→Output (PRODUCES): 1 relation only  
      * Output→Human (DELIVERS_TO): 1 relation only
-     * Agent→Tool (USES): Only if tools exist
+     * Agent→Tool (USES): Include exactly one when a Tool exists
+     * Tool→Task (REQUIRED_BY): Include one only when the task explicitly depends on a Tool (execution/platform dependency). Otherwise omit.
    
-   - FORBIDDEN: ASSIGNED_TO, INTERVENES, REQUIRED_BY, complex multi-connections
-   - TARGET: Maximum 9 total relations for 3-agent workflows
+   - FORBIDDEN: ASSIGNED_TO, INTERVENES
+   - TARGET: Maximum 9–10 total relations for 3-agent workflows
 
 5. QUALITY CHECK (Contextual Graph Enhanced):
    - Verify all relation IDs reference existing entities
@@ -219,6 +225,11 @@ ANALYSIS STEPS:
    - CONTEXT VALIDATION: Ensure relationships align with trace context and domain semantics
    - EDGE CASE HANDLING: Apply simplification for simple workflows, complexity for multi-domain scenarios
 
+NAMING & CANONICALIZATION:
+- Title Case + spaces only (no underscores in names)
+- Prefer canonical domain templates (see system rules) when content/tags match
+- Prefer concise outputs like "Restaurant Recommendations" for discovery workflows
+
 FORMATTING:
 - IDs: agent_001, task_001, tool_001, etc.
 - Empty raw_prompt and interaction_prompt fields
@@ -243,17 +254,6 @@ TRACE DATA:
             logger.info(f"Response parsed successfully. Type: {type(reasoning_result)}")
             
             knowledge_graph = reasoning_result.final_answer
-            
-            # Optional post-alignment for reference/sample compatibility
-            if self.compatibility_mode:
-                try:
-                    knowledge_graph = self._align_for_reference_compatibility(
-                        knowledge_graph, input_data
-                    )
-                    # Ensure saved result contains aligned KG
-                    reasoning_result.final_answer = knowledge_graph
-                except Exception as align_err:
-                    logger.warning(f"Compatibility alignment skipped due to error: {align_err}")
             logger.info(f"Knowledge graph extracted. Type: {type(knowledge_graph)}")
             
             # Save reasoning steps and knowledge graph for analysis
@@ -273,87 +273,7 @@ TRACE DATA:
             logger.error(f"Extraction failed: {e}")
             raise
     
-    def _align_for_reference_compatibility(self, kg: KnowledgeGraph, trace_text: str) -> KnowledgeGraph:
-        """Apply light, deterministic adjustments to better match sample/reference KGs.
-
-        Adjustments:
-        - Ensure exactly one Agent→Tool USES relation if a Tool exists (choose the agent performing task_001 or the first Agent)
-        - Optionally add Tool→Task REQUIRED_BY for the first Task when allow_required_by is enabled
-        - Normalize common display names (e.g., "Computer_terminal" → "Computer Terminal")
-        - For discovery workflows, normalize Output name to "Restaurant Recommendations" when semantically identical
-        """
-        # Build indices
-        entities_by_id = {e.id: e for e in kg.entities}
-        agents = [e for e in kg.entities if e.type == "Agent"]
-        tools = [e for e in kg.entities if e.type == "Tool"]
-        tasks = [e for e in kg.entities if e.type == "Task"]
-        outputs = [e for e in kg.entities if e.type == "Output"]
-
-        # Normalize tool naming
-        for tool in tools:
-            if tool.name == "Computer_terminal":
-                tool.name = "Computer Terminal"
-
-        # Detect if a USES relation already exists
-        def has_uses():
-            return any(r.type == "USES" for r in kg.relations)
-
-        # Pick primary agent (PERFORMS task_001 or first Agent)
-        def pick_primary_agent_id() -> str:
-            # Find task_001
-            task1 = next((t for t in tasks if t.id == "task_001"), None)
-            if task1 is not None:
-                for r in kg.relations:
-                    if r.type == "PERFORMS" and r.target == task1.id:
-                        return r.source
-            return agents[0].id if agents else ""
-
-        # Add USES if missing and a Tool exists
-        if tools and agents and not has_uses():
-            try:
-                primary_agent_id = pick_primary_agent_id()
-                tool_id = tools[0].id
-                from agentgraph.shared.models.reference_based import Relation, ContentReference
-                uses_rel = Relation(
-                    source=primary_agent_id,
-                    target=tool_id,
-                    type="USES",
-                    importance="MEDIUM",
-                    interaction_prompt="",
-                    interaction_prompt_ref=[ContentReference(line_start=None, line_end=None)]
-                )
-                kg.relations.append(uses_rel)
-            except Exception as e:
-                logger.debug(f"Failed to add USES relation: {e}")
-
-        # Optionally add REQUIRED_BY Tool→first Task
-        if self.allow_required_by and tools and tasks:
-            try:
-                has_required_by = any(r.type == "REQUIRED_BY" for r in kg.relations)
-                if not has_required_by:
-                    from agentgraph.shared.models.reference_based import Relation, ContentReference
-                    req_rel = Relation(
-                        source=tools[0].id,
-                        target=tasks[0].id,
-                        type="REQUIRED_BY",
-                        importance="MEDIUM",
-                        interaction_prompt="",
-                        interaction_prompt_ref=[ContentReference(line_start=None, line_end=None)]
-                    )
-                    kg.relations.append(req_rel)
-            except Exception as e:
-                logger.debug(f"Failed to add REQUIRED_BY relation: {e}")
-
-        # Normalize Output naming for discovery workflows
-        system_name = (kg.system_name or "").lower()
-        is_discovery = "restaurant" in system_name or "location" in system_name
-        if is_discovery and outputs:
-            out = outputs[0]
-            # If name already contains Restaurant Recommendations, standardize to concise form
-            if "restaurant recommendations" in (out.name or "").lower() and out.name.strip() != "Restaurant Recommendations":
-                out.name = "Restaurant Recommendations"
-
-        return kg
+    
 
     def _save_extraction_results(self, reasoning_result, trace_text: str):
         """Save reasoning steps and knowledge graph for analysis"""