ad

agentgraph/methods/production/openai_structured_extractor.py

@@ -255,9 +255,9 @@ TRACE DATA:
 {input_data}"""
         
         try:
-            response = self.client.beta.chat.completions.parse(
+            response = self.responses.parse(
                 model=self.model,
-                messages=[
+                input=[
                     {"role": "system", "content": system_prompt},
                     {"role": "user", "content": user_prompt}
                 ],
@@ -266,10 +266,10 @@ TRACE DATA:
             
             # Get the parsed response with reasoning steps
             logger.info("Parsing OpenAI response...")
-            reasoning_result = response.choices[0].message.parsed
+            reasoning_result = response.output_parsed
             logger.info(f"Response parsed successfully. Type: {type(reasoning_result)}")
             
-            knowledge_graph = reasoning_result.final_answer
+            knowledge_graph = reasoning_result["final_answer"]
             logger.info(f"Knowledge graph extracted. Type: {type(knowledge_graph)}")
             
             # Save reasoning steps and knowledge graph for analysis
@@ -279,7 +279,7 @@ TRACE DATA:
             
             # Log reasoning steps for debugging
             logger.info("CoT Reasoning Steps:")
-            for step in reasoning_result.steps:
+            for step in reasoning_result["steps"]:
                 logger.info(f"Step: {step.explanation}")
             
             logger.info(f"Extraction complete: {len(knowledge_graph.entities)} entities, {len(knowledge_graph.relations)} relations")
@@ -309,11 +309,11 @@ TRACE DATA:
             
             # Prepare data to save using model_dump for safe serialization
             logger.info("Serializing reasoning steps...")
-            reasoning_steps_data = [step.model_dump() for step in reasoning_result.steps]
+            reasoning_steps_data = [step.model_dump() for step in reasoning_result["steps"]]
             logger.info(f"Reasoning steps serialized: {len(reasoning_steps_data)} steps")
             
             logger.info("Serializing knowledge graph...")
-            kg_data = reasoning_result.final_answer.model_dump()
+            kg_data = reasoning_result['final_answer'].model_dump()
             logger.info(f"Knowledge graph serialized: {kg_data.get('system_name', 'Unknown')}")
             
             save_data = {

agentgraph/methods/production/openai_structured_extractor_backup_20250907_161109.py

@@ -1,548 +0,0 @@
-#!/usr/bin/env python3
-"""
-OpenAI Structured Outputs Knowledge Extractor
-
-A simple, direct approach using OpenAI's structured outputs API to extract
-knowledge graphs in one step using Pydantic models.
-"""
-
-import os
-import json
-import logging
-from typing import Optional, List, Dict, Any
-import uuid
-from datetime import datetime
-from dotenv import load_dotenv
-from openai import OpenAI
-from pydantic import BaseModel
-
-# Import Pydantic models
-from agentgraph.shared.models.reference_based import KnowledgeGraph, Entity, Relation, ContentReference
-
-# CoT reasoning models
-class Step(BaseModel):
-    explanation: str
-    output: str
-
-class KnowledgeGraphReasoning(BaseModel):
-    steps: List[Step]
-    final_answer: KnowledgeGraph
-
-# Load environment variables from root directory
-load_dotenv('/Users/zekunwu/Desktop/agent_monitoring/.env')
-
-# Configure logging
-logger = logging.getLogger(__name__)
-
-
-class OpenAIStructuredExtractor:
-    """
-    Simple knowledge graph extractor using OpenAI's structured outputs.
-    """
-    
-    def __init__(self, model: str = "gpt-5-mini"):
-        """
-        Initialize the extractor.
-        
-        Args:
-            model: OpenAI model to use (must support structured outputs)
-        """
-        self.model = model
-        self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
-        logger.info(f"OpenAI Structured Extractor initialized with model: {model}")
-    
-    def extract_knowledge_graph(self, input_data: str, context_documents: Optional[List[Dict[str, Any]]] = None) -> KnowledgeGraph:
-        """
-        Extract knowledge graph from input data using OpenAI structured outputs.
-        
-        Args:
-            input_data: The trace data to analyze
-            context_documents: Optional context documents (unused in this simple version)
-            
-        Returns:
-            KnowledgeGraph: Extracted knowledge graph
-        """
-        logger.info(f"Starting knowledge graph extraction for {len(input_data)} characters of input")
-        
-        # System prompt - focus on your role and methodology
-        system_prompt = """You are an expert knowledge graph analyst specializing in agent system traces.
-
-Extract comprehensive knowledge graphs capturing all entities and their precise relationships. Focus on workflow accuracy and relationship completeness.
-
-CORE PRINCIPLES:
-1. Capture ALL participants: agents, tools, tasks, inputs, outputs, humans
-2. Use professional naming (spaces, not underscores)
-3. Map complete workflows: Input → Agents perform Tasks → Output → Human
-4. Focus on MAIN workflow tasks (avoid over-decomposition into subtasks)
-5. Use NEXT only for clear sequential task dependencies
-6. Use REQUIRED_BY for essential tool dependencies (not USES for general support)
-7. Identify failures: errors, mistakes, broken processes, incorrect outputs
-8. Suggest optimizations: improvements, efficiency gains, better approaches
-
-ID ASSIGNMENT GUIDELINES:
-- Use sequential numbering starting from 001 for each entity type
-- agent_001, agent_002, agent_003... (in order of appearance/importance)
-- task_001, task_002, task_003... (in workflow order)
-- input_001, output_001, human_001 (typically one of each)
-- tool_001 for Computer Terminal or similar tools
-- Assign based on trace content, not predetermined roles
-
-RELATIONSHIP TYPES (use exactly these):
-- CONSUMED_BY: Input consumed by agent (input_001 → agent_001)
-- PERFORMS: Agent performs task (agent_001 → task_001)  
-- NEXT: Task A leads to Task B in sequence (task_001 → task_002)
-- PRODUCES: Task produces output (task_001 → output_001)
-- DELIVERS_TO: Output delivered to human (output_001 → human_001)
-- REQUIRED_BY: Tool required by task (agent_004 → task_001)
-- USES: Agent uses tool (agent_001 → agent_004)
-- ASSIGNED_TO: Task assigned to agent (task_001 → agent_002)
-- INTERVENES: Agent coordinates task (agent_002 → task_001)
-
-Return complete, accurate knowledge graphs with proper workflow sequences."""
-
-        # User prompt - specific instructions with few-shot example and data
-        user_prompt = f"""Analyze this agent system trace step by step and extract a knowledge graph.
-
-Think through this systematically. Use as many reasoning steps as needed to thoroughly analyze the trace.
-Each step should have a clear explanation and specific output findings.
-
-IMPORTANT FORMATTING RULES:
-- Set raw_prompt to empty string "" for ALL entities (content should only be in raw_prompt_ref)
-- Set interaction_prompt to empty string "" for ALL relations (content should only be in interaction_prompt_ref)
-- Use only line number references in raw_prompt_ref and interaction_prompt_ref, never include actual text content
-- Determine your own number of reasoning steps based on trace complexity
-
-EXAMPLE OUTPUT STRUCTURE:
-Here's the expected knowledge graph structure for multi-agent collaboration traces:
-{{
-  "system_name": "California Great America Ticket Analysis System",
-  "system_summary": "This system helps analyze the cost-saving potential of purchasing season passes versus individual daily tickets at California's Great America in San Jose. The process starts with an inquiry regarding savings from the Inquiry about Savings from Season Pass vs Daily Tickets (input_001), which is consumed by the Verification Expert (agent_002), who performs the Verify Cost of Daily Ticket and Season Pass in 2024 (task_001).",
-  "entities": [
-    {{
-      "id": "agent_001",
-      "type": "Agent",
-      "name": "Problem Solving Expert",
-      "importance": "HIGH",
-      "raw_prompt": "",
-      "raw_prompt_ref": [{{"line_start": 17, "line_end": 17}}, {{"line_start": 34, "line_end": 34}}, {{"line_start": 45, "line_end": 45}}]
-    }},
-    {{
-      "id": "agent_002",
-      "type": "Agent",
-      "name": "Verification Expert",
-      "importance": "HIGH",
-      "raw_prompt": "",
-      "raw_prompt_ref": [{{"line_start": 66, "line_end": 66}}, {{"line_start": 112, "line_end": 112}}, {{"line_start": 149, "line_end": 149}}, {{"line_start": 164, "line_end": 164}}]
-    }},
-    {{
-      "id": "agent_003",
-      "type": "Agent",
-      "name": "Arithmetic Progressions Expert",
-      "importance": "MEDIUM",
-      "raw_prompt": "",
-      "raw_prompt_ref": [{{"line_start": 172, "line_end": 172}}, {{"line_start": 181, "line_end": 181}}]
-    }},
-    {{
-      "id": "agent_004",
-      "type": "Tool",
-      "name": "Computer Terminal",
-      "importance": "LOW",
-      "raw_prompt": "",
-      "raw_prompt_ref": [{{"line_start": 21, "line_end": 21}}, {{"line_start": 32, "line_end": 32}}]
-    }},
-    {{
-      "id": "task_001",
-      "type": "Task",
-      "name": "Verify Cost of Daily Ticket and Season Pass in 2024",
-      "importance": "HIGH",
-      "raw_prompt": "",
-      "raw_prompt_ref": [{{"line_start": 8, "line_end": 8}}, {{"line_start": 10, "line_end": 10}}, {{"line_start": 11, "line_end": 12}}]
-    }},
-    {{
-      "id": "task_002",
-      "type": "Task",
-      "name": "Calculate Savings Amount",
-      "importance": "HIGH",
-      "raw_prompt": "",
-      "raw_prompt_ref": [{{"line_start": 119, "line_end": 126}}]
-    }},
-    {{
-      "id": "input_001",
-      "type": "Input",
-      "name": "Inquiry about Savings from Season Pass vs Daily Tickets",
-      "importance": "HIGH",
-      "raw_prompt": "",
-      "raw_prompt_ref": [{{"line_start": 6, "line_end": 6}}]
-    }},
-    {{
-      "id": "output_001",
-      "type": "Output",
-      "name": "Saved Amount from Season Pass Purchase",
-      "importance": "HIGH",
-      "raw_prompt": "",
-      "raw_prompt_ref": [{{"line_start": 119, "line_end": 119}}, {{"line_start": 126, "line_end": 126}}]
-    }},
-    {{
-      "id": "human_001",
-      "type": "Human",
-      "name": "Park Visitor",
-      "importance": "HIGH",
-      "raw_prompt": "",
-      "raw_prompt_ref": [{{"line_start": 1, "line_end": 1}}]
-    }}
-  ],
-  "relations": [
-    {{
-      "id": "relation_001",
-      "source": "input_001",
-      "target": "agent_002",
-      "type": "CONSUMED_BY",
-      "importance": "HIGH",
-      "interaction_prompt": "",
-      "interaction_prompt_ref": [{{"line_start": 6, "line_end": 6}}]
-    }},
-    {{
-      "id": "relation_002",
-      "source": "agent_002",
-      "target": "task_001",
-      "type": "PERFORMS",
-      "importance": "HIGH",
-      "interaction_prompt": "",
-      "interaction_prompt_ref": [{{"line_start": 112, "line_end": 112}}, {{"line_start": 164, "line_end": 164}}]
-    }},
-    {{
-      "id": "relation_003",
-      "source": "task_001",
-      "target": "agent_002",
-      "type": "ASSIGNED_TO",
-      "importance": "HIGH",
-      "interaction_prompt": "",
-      "interaction_prompt_ref": [{{"line_start": 8, "line_end": 8}}]
-    }},
-    {{
-      "id": "relation_004",
-      "source": "task_001",
-      "target": "output_001",
-      "type": "PRODUCES",
-      "importance": "HIGH",
-      "interaction_prompt": "",
-      "interaction_prompt_ref": [{{"line_start": 119, "line_end": 119}}]
-    }},
-    {{
-      "id": "relation_005",
-      "source": "output_001",
-      "target": "human_001",
-      "type": "DELIVERS_TO",
-      "importance": "HIGH",
-      "interaction_prompt": "",
-      "interaction_prompt_ref": [{{"line_start": 126, "line_end": 126}}]
-    }},
-    {{
-      "id": "relation_006",
-      "source": "task_001",
-      "target": "task_002",
-      "type": "NEXT",
-      "importance": "HIGH",
-      "interaction_prompt": "",
-      "interaction_prompt_ref": [{{"line_start": 164, "line_end": 164}}]
-    }},
-    {{
-      "id": "relation_007",
-      "source": "task_002",
-      "target": "agent_004",
-      "type": "REQUIRED_BY",
-      "importance": "MEDIUM",
-      "interaction_prompt": "",
-      "interaction_prompt_ref": [{{"line_start": 50, "line_end": 55}}]
-    }}
-  ],
-  "failures": [
-    {{
-      "id": "failure_001",
-      "description": "Verification Expert failed to access real-time pricing data, relying on potentially outdated cost estimates",
-      "raw_text": "However, since I am currently unable to access external websites, I will use the provided cost",
-      "raw_text_ref": [],
-      "affected_id": "agent_002",
-      "risk_type": "RETRIEVAL_ERROR"
-    }}
-  ],
-  "optimizations": [
-    {{
-      "id": "opt_001",
-      "description": "Implement automated price verification system to reduce manual verification overhead and improve accuracy",
-      "raw_text": "Enhanced price verification with real-time data access",
-      "raw_text_ref": [],
-      "affected_ids": ["agent_002", "task_001"],
-      "recommendation_type": "TOOL_ENHANCEMENT"
-    }}
-  ]
-}}
-
-Now analyze the following trace data:
-
-ENTITY TYPES:
-- Agent: AI agents with specific roles
-- Task: Specific tasks or objectives  
-- Tool: Tools or functions used
-- Input: Data inputs to the system
-- Output: Data outputs from the system
-- Human: Human users or stakeholders
-
-RELATIONSHIP TYPES:
-- CONSUMED_BY: Input→Agent
-- PERFORMS: Agent→Task  
-- ASSIGNED_TO: Task→Agent
-- USES: Agent→Tool
-- REQUIRED_BY: Tool→Task
-- SUBTASK_OF: Task→Task
-- NEXT: Task→Task (sequence)
-- PRODUCES: Task→Output
-- DELIVERS_TO: Output→Human
-- INTERVENES: Agent/Human→Task
-
-REQUIREMENTS:
-For each entity provide:
-- id: unique identifier
-- 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 pointing to relevant lines in the input trace where this entity is mentioned or defined
-
-For each relation provide:
-- id: unique identifier
-- source/target: entity IDs
-- 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 pointing to relevant lines in the input trace where this relationship evidence is found
-
-IMPORTANT: Create appropriate line number references based on where entities and relationships are mentioned in the trace. Use reasonable line estimates for JSON structure.
-
-Also provide system_name and system_summary for the overall system.
-
-EXTRACTION FOCUS:
-1. Identify MAIN participants: primary agents, core tasks, key inputs/outputs, humans
-2. Focus on the PRIMARY workflow (avoid creating too many subtasks)
-3. Connect entities based on ACTUAL interactions in the trace
-4. Use clean professional naming (no underscores)
-5. DETECT FAILURES: Look for errors, exceptions, incorrect results, failed executions
-6. SUGGEST OPTIMIZATIONS: Identify inefficiencies, improvement opportunities
-
-CRITICAL RELATIONSHIP PATTERNS (follow exactly):
-- input_001 --CONSUMED_BY--> agent_002 (main processing agent)
-- agent_002 --PERFORMS--> task_001 (primary task execution)
-- task_001 --ASSIGNED_TO--> agent_002 (task ownership)
-- task_001 --PRODUCES--> output_001 (direct output generation)
-- output_001 --DELIVERS_TO--> human_001 (final delivery)
-- agent_002 --INTERVENES--> task_001 (coordination/oversight)
-- agent_001 --USES--> agent_004 (tool utilization)
-
-ENTITY ID RULES:
-- agent_001: First mentioned expert, agent_002: Main performer, agent_003: Supporting expert
-- agent_004: Always Computer Terminal (tool type)
-- task_001: Primary workflow task (avoid creating task_002 unless truly sequential)
-- Use exactly these IDs, don't create tool_001 or extra tasks unnecessarily
-
-FAILURE DETECTION (look for):
-- Execution errors, exceptions, failed operations
-- Incorrect outputs, wrong calculations, invalid results  
-- Missing validations, incomplete processes
-- Tool limitations, access restrictions
-- Agent coordination problems
-
-OPTIMIZATION OPPORTUNITIES (suggest):
-- Tool enhancements, automation possibilities
-- Workflow simplifications, redundancy removal
-- Agent merging or specialization
-- Performance improvements, efficiency gains
-- Missing capabilities or better approaches
-
-IMPORTANT: Always provide at least 1-2 failures and optimizations based on trace analysis.
-
-TRACE DATA:
-{input_data}"""
-        
-        try:
-            response = self.client.responses.parse(
-                model=self.model,
-                input=[
-                    {"role": "system", "content": system_prompt},
-                    {"role": "user", "content": user_prompt}
-                ],
-                text_format=KnowledgeGraphReasoning,
-            )
-            
-            # Get the parsed response with reasoning steps
-            logger.info("Parsing OpenAI response...")
-            reasoning_result = response.output_parsed
-            logger.info(f"Response parsed successfully. Type: {type(reasoning_result)}")
-            
-            knowledge_graph = reasoning_result.final_answer
-            logger.info(f"Knowledge graph extracted. Type: {type(knowledge_graph)}")
-            
-            # Save reasoning steps and knowledge graph for analysis
-            logger.info("About to save extraction results...")
-            self._save_extraction_results(reasoning_result, input_data)
-            logger.info("Save extraction results completed.")
-            
-            # Log reasoning steps for debugging
-            logger.info("CoT Reasoning Steps:")
-            for step in reasoning_result.steps:
-                logger.info(f"Step: {step.explanation}")
-            
-            logger.info(f"Extraction complete: {len(knowledge_graph.entities)} entities, {len(knowledge_graph.relations)} relations")
-            return knowledge_graph
-            
-        except Exception as e:
-            logger.error(f"Extraction failed: {e}")
-            raise
-    
-    def _save_extraction_results(self, reasoning_result, trace_text: str):
-        """Save reasoning steps and knowledge graph for analysis"""
-        try:
-            logger.info("Starting to save extraction results...")
-            
-            # Create results directory if it doesn't exist
-            results_dir = "extraction_analysis"
-            os.makedirs(results_dir, exist_ok=True)
-            logger.info(f"Results directory ready: {results_dir}")
-            
-            # Generate timestamp and unique ID
-            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-            extraction_id = str(uuid.uuid4())[:8]
-            filename = f"cot_extraction_{timestamp}_{extraction_id}.json"
-            logger.info(f"Generated filename: {filename}")
-            
-            # Prepare data to save using model_dump for safe serialization
-            logger.info("Serializing reasoning steps...")
-            reasoning_steps_data = [step.model_dump() for step in reasoning_result.steps]
-            logger.info(f"Reasoning steps serialized: {len(reasoning_steps_data)} steps")
-            
-            logger.info("Serializing knowledge graph...")
-            kg_data = reasoning_result.final_answer.model_dump()
-            logger.info(f"Knowledge graph serialized: {kg_data.get('system_name', 'Unknown')}")
-            
-            save_data = {
-                "timestamp": timestamp,
-                "extraction_id": extraction_id,
-                "model": self.model,
-                "reasoning_steps": reasoning_steps_data,
-                "knowledge_graph": kg_data,
-                "input_trace_length": len(trace_text),
-                "input_trace_preview": trace_text[:500] + "..." if len(trace_text) > 500 else trace_text
-            }
-            
-            # Save to file
-            filepath = os.path.join(results_dir, filename)
-            logger.info(f"Writing to file: {filepath}")
-            
-            with open(filepath, 'w', encoding='utf-8') as f:
-                json.dump(save_data, f, indent=2, ensure_ascii=False)
-            
-            # Verify file was written
-            file_size = os.path.getsize(filepath)
-            logger.info(f"Extraction results saved successfully: {filepath} ({file_size} bytes)")
-            
-        except Exception as e:
-            logger.error(f"Failed to save extraction results: {e}")
-            import traceback
-            logger.error(f"Traceback: {traceback.format_exc()}")
-
-    def process_text(self, input_data: str) -> Dict[str, Any]:
-        """
-        Process text and return structured response (for compatibility with extraction factory).
-        
-        Args:
-            input_data: The trace data to analyze
-            
-        Returns:
-            Dict with success status and kg_data
-        """
-        try:
-            kg = self.extract_knowledge_graph(input_data)
-            return {
-                "success": True,
-                "kg_data": kg.model_dump()
-            }
-        except Exception as e:
-            return {
-                "success": False,
-                "error": str(e),
-                "kg_data": {}
-            }
-
-def extract_knowledge_graph_with_context(
-    input_data: str,
-    context_documents: Optional[List[Dict[str, Any]]] = None,
-    model: str = "gpt-5-mini"
-) -> KnowledgeGraph:
-    """
-    Main entry point for knowledge graph extraction.
-    
-    Args:
-        input_data: The trace data to analyze
-        context_documents: Optional context documents
-        model: OpenAI model to use
-        
-    Returns:
-        KnowledgeGraph: Extracted knowledge graph
-    """
-    extractor = OpenAIStructuredExtractor(model=model)
-    return extractor.extract_knowledge_graph(input_data, context_documents)
-
-def extract_knowledge_graph(input_data: str) -> KnowledgeGraph:
-    """
-    Simple entry point without context (for backward compatibility).
-    
-    Args:
-        input_data: The trace data to analyze
-        
-    Returns:
-        KnowledgeGraph: Extracted knowledge graph
-    """
-    return extract_knowledge_graph_with_context(input_data)
-
-# Factory class for integration
-class OpenAIStructuredFactory:
-    """Factory class for OpenAI structured extraction method."""
-    
-    def __init__(self, model: str = "gpt-5-mini"):
-        self.model = model
-        self.extractor = OpenAIStructuredExtractor(model)
-    
-    def set_model(self, model: str):
-        """Set the model for this factory."""
-        self.model = model
-        self.extractor = OpenAIStructuredExtractor(model)
-    
-    def process_text(self, input_data: str) -> Dict[str, Any]:
-        """Process text using the extractor."""
-        return self.extractor.process_text(input_data)
-    
-    def extract_knowledge_graph(self, input_data: str, context_documents: Optional[List[Dict[str, Any]]] = None) -> KnowledgeGraph:
-        """Extract knowledge graph using the current model."""
-        return self.extractor.extract_knowledge_graph(input_data, context_documents)
-
-# Export factory instance
-openai_structured_factory = OpenAIStructuredFactory()
-
-if __name__ == "__main__":
-    # Simple test
-    test_input = """
-    Agent: DataAnalyzer
-    Task: Analyze customer data and generate insights
-    Tool: pandas_analyzer
-    
-    The DataAnalyzer agent processes customer data using pandas_analyzer tool
-    to generate business insights for the marketing team.
-    """
-    
-    print("Testing OpenAI Structured Extractor...")
-    try:
-        kg = extract_knowledge_graph(test_input)
-        print(f"✅ Success! Extracted {len(kg.entities)} entities and {len(kg.relations)} relations")
-        print(f"System: {kg.system_name}")
-        print(f"Summary: {kg.system_summary}")
-    except Exception as e:
-        print(f"❌ Error: {e}")