add

agentgraph/methods/production/openai_structured_extractor.py

@@ -16,7 +16,7 @@ from openai import OpenAI
 from pydantic import BaseModel
 
 # Import Pydantic models
-from agentgraph.shared.models.reference_based import KnowledgeGraph, Entity, Relation
+from agentgraph.shared.models.reference_based import KnowledgeGraph, Entity, Relation, ContentReference
 
 # Load environment variables from root directory
 load_dotenv('/Users/zekunwu/Desktop/agent_monitoring/.env')
@@ -24,80 +24,6 @@ load_dotenv('/Users/zekunwu/Desktop/agent_monitoring/.env')
 # Configure logging
 logger = logging.getLogger(__name__)
 
-# Simplified models for OpenAI structured outputs
-class SimpleEntity(BaseModel):
-    id: str
-    type: str  # Agent, Task, Tool, Input, Output, Human
-    name: str
-    importance: str  # HIGH, MEDIUM, LOW
-
-class SimpleRelation(BaseModel):
-    id: str
-    source: str
-    target: str
-    type: str  # PERFORMS, USES, etc.
-    importance: str
-
-class SimpleKnowledgeGraph(BaseModel):
-    system_name: str
-    system_summary: str
-    entities: List[SimpleEntity]
-    relations: List[SimpleRelation]
-
-def normalize_importance(importance: str) -> str:
-    """Normalize importance values to HIGH/MEDIUM/LOW."""
-    importance_upper = importance.upper()
-    # Map common variations to standard values
-    mapping = {
-        "CRITICAL": "HIGH",
-        "VERY HIGH": "HIGH", 
-        "VERY LOW": "LOW",
-        "NORMAL": "MEDIUM",
-        "STANDARD": "MEDIUM"
-    }
-    return mapping.get(importance_upper, importance_upper)
-
-def convert_simple_to_full_kg(simple_kg: SimpleKnowledgeGraph) -> KnowledgeGraph:
-    """Convert simplified KG to full KnowledgeGraph model."""
-    
-    # Convert entities
-    entities = []
-    for se in simple_kg.entities:
-        entity = Entity(
-            id=se.id,
-            type=se.type,
-            name=se.name,
-            importance=normalize_importance(se.importance),  # Normalize importance
-            raw_prompt="",  # Empty as per requirements
-            raw_prompt_ref=[]  # Empty for now
-        )
-        entities.append(entity)
-    
-    # Convert relations
-    relations = []
-    for sr in simple_kg.relations:
-        relation = Relation(
-            id=sr.id,
-            source=sr.source,
-            target=sr.target,
-            type=sr.type,
-            importance=normalize_importance(sr.importance),  # Normalize importance
-            interaction_prompt="",  # Empty as per requirements
-            interaction_prompt_ref=[]  # Empty for now
-        )
-        relations.append(relation)
-    
-    # Create full KnowledgeGraph
-    kg = KnowledgeGraph(
-        system_name=simple_kg.system_name,
-        system_summary=simple_kg.system_summary,
-        entities=entities,
-        relations=relations,
-        failures=None,  # Not generated by this simple method
-        optimizations=None  # Not generated by this simple method
-    )
-    
-    return kg
 
 class OpenAIStructuredExtractor:
     """
@@ -128,8 +54,10 @@ class OpenAIStructuredExtractor:
         """
         logger.info(f"Starting knowledge graph extraction for {len(input_data)} characters of input")
         
-        # Simple system prompt - much shorter than the complex ones
-        system_prompt = """You are an expert at analyzing agent system traces and extracting knowledge graphs.
+        # System prompt for direct KnowledgeGraph extraction with content references
+        system_prompt = """You are an expert at analyzing agent system traces and extracting knowledge graphs with precise content references.
+
+The input may contain line markers like <L1>, <L2>, etc. Use these to create accurate content references when available.
 
 Extract a knowledge graph with these entity types:
 - Agent: AI agents with specific roles
@@ -151,13 +79,26 @@ Use these relationship types:
 - DELIVERS_TO: Output→Human
 - INTERVENES: Agent/Human→Task
 
-Create a complete knowledge graph with:
-1. Meaningful entities with descriptive names
-2. Logical relationships between entities
-3. A system name and summary
-4. At least 3-5 entities for any non-trivial workflow
+For each entity provide:
+- id: unique identifier (generate if needed)
+- type: one of the types above
+- name: descriptive name
+- importance: HIGH, MEDIUM, or LOW
+- raw_prompt: actual prompt/specification content that defines this entity
+- raw_prompt_ref: list of content references with line_start and line_end (if line markers available)
+
+For each relation provide:
+- id: unique identifier
+- source: source entity id
+- target: target entity id  
+- type: one of the types above
+- importance: HIGH, MEDIUM, or LOW
+- interaction_prompt: runtime evidence showing this relationship occurred
+- interaction_prompt_ref: list of content references (if line markers available)
+
+Provide system_name and system_summary for the overall system.
 
-Focus on identifying the actual workflow, not framework details."""
+Focus on extracting the actual workflow with meaningful entities and relationships."""
 
         user_prompt = f"Analyze this agent system trace and extract a knowledge graph:\n\n{input_data}"
         
@@ -168,12 +109,11 @@ Focus on identifying the actual workflow, not framework details."""
                     {"role": "system", "content": system_prompt},
                     {"role": "user", "content": user_prompt}
                 ],
-                text_format=SimpleKnowledgeGraph,
+                text_format=KnowledgeGraph,
             )
             
-            # Get the parsed response and convert to full model
-            simple_kg = response.output_parsed
-            knowledge_graph = convert_simple_to_full_kg(simple_kg)
+            # Get the parsed response directly as KnowledgeGraph
+            knowledge_graph = response.output_parsed
             
             logger.info(f"Extraction complete: {len(knowledge_graph.entities)} entities, {len(knowledge_graph.relations)} relations")
             return knowledge_graph