from agents import Agent, Runner, set_trace_processors, TraceProcessor
import json
from datetime import datetime
import os

class LocalFileTraceProcessor(TraceProcessor):
    def __init__(self, output_dir="traces"):
        self.output_dir = output_dir
        os.makedirs(output_dir, exist_ok=True)
        
    def process_trace(self, trace):
        """Process and save each complete trace"""
        trace_data = {
            "trace_id": trace.trace_id,
            "workflow_name": trace.workflow_name,
            "timestamp": datetime.now().isoformat(),
            "spans": []
        }
        
        # Extract all spans from the trace
        for span in trace.spans:
            span_data = {
                "span_id": span.span_id,
                "parent_id": span.parent_id,
                "operation": span.operation_name,
                "start_time": span.start_time,
                "end_time": span.end_time,
                "attributes": span.attributes,
                "events": []
            }
            
            # Capture events within spans (LLM calls, tool calls, etc.)
            for event in span.events:
                event_data = {
                    "name": event.name,
                    "timestamp": event.timestamp,
                    "attributes": event.attributes
                }
                
                # Extract key data for evals
                if event.name == "llm.generation":
                    event_data["input"] = event.attributes.get("input", "")
                    event_data["output"] = event.attributes.get("output", "")
                    event_data["model"] = event.attributes.get("model", "")
                    event_data["usage"] = event.attributes.get("usage", {})
                
                span_data["events"].append(event_data)
            
            trace_data["spans"].append(span_data)
        
        # Save to file
        filename = f"{self.output_dir}/trace_{trace.trace_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        with open(filename, 'w') as f:
            json.dump(trace_data, f, indent=2, default=str)
    
    def process_span(self, span):
        """Process individual spans as they complete"""
        # You can optionally process spans in real-time
        pass

# Set up the custom processor
custom_processor = LocalFileTraceProcessor(output_dir="agent_traces")
set_trace_processors([custom_processor])

# Your agent code
agent = Agent(
    name="MyAgent",
    instructions="You are a helpful assistant.",
    model="gpt-4"
)

# Run the agent - traces will be automatically saved
result = Runner.run_sync(agent, "What's the weather like today?")

from agents import set_trace_processors, add_trace_processor
from agents.tracing.processors import OTLPHTTPTraceSpanProcessor
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
import json

# Configure OpenTelemetry
provider = TracerProvider()
trace.set_tracer_provider(provider)

# Export to local OTLP collector (which can write to files)
otlp_exporter = OTLPSpanExporter(
    endpoint="http://localhost:4318/v1/traces",
    headers={"Content-Type": "application/json"}
)
span_processor = BatchSpanProcessor(otlp_exporter)
provider.add_span_processor(span_processor)

# Add OTLP processor to agents SDK
add_trace_processor(
    OTLPHTTPTraceSpanProcessor(endpoint="http://localhost:4318/v1/traces")
)

# Also add a custom processor for immediate local access
class EvalDataProcessor(TraceProcessor):
    def __init__(self, eval_file="eval_data.jsonl"):
        self.eval_file = eval_file
        
    def process_trace(self, trace):
        # Extract data specifically for evaluations
        for span in trace.spans:
            for event in span.events:
                if event.name == "llm.generation":
                    eval_entry = {
                        "trace_id": trace.trace_id,
                        "span_id": span.span_id,
                        "timestamp": event.timestamp,
                        "input": event.attributes.get("messages", []),
                        "output": event.attributes.get("output", ""),
                        "model": event.attributes.get("model", ""),
                        "temperature": event.attributes.get("temperature", 1.0),
                        "usage": event.attributes.get("usage", {}),
                        "ideal_answer": None  # Placeholder for manual annotation
                    }
                    
                    # Append to JSONL file for easy streaming
                    with open(self.eval_file, 'a') as f:
                        f.write(json.dumps(eval_entry) + '\n')

add_trace_processor(EvalDataProcessor())

import pandas as pd
from typing import List, Dict, Any

class EvaluationDataCollector:
    def __init__(self, output_file="evaluation_dataset.csv"):
        self.output_file = output_file
        self.eval_data = []
        
    def extract_from_trace_files(self, trace_directory="agent_traces"):
        """Extract evaluation data from saved trace files"""
        import glob
        
        for trace_file in glob.glob(f"{trace_directory}/*.json"):
            with open(trace_file, 'r') as f:
                trace_data = json.load(f)
                
            # Extract LLM calls from trace
            for span in trace_data.get("spans", []):
                for event in span.get("events", []):
                    if event.get("name") == "llm.generation":
                        eval_entry = {
                            "trace_id": trace_data["trace_id"],
                            "workflow_name": trace_data.get("workflow_name", ""),
                            "timestamp": event["timestamp"],
                            "model": event["attributes"].get("model", ""),
                            "input_messages": json.dumps(event["attributes"].get("input", [])),
                            "output": event["attributes"].get("output", ""),
                            "prompt_tokens": event["attributes"].get("usage", {}).get("prompt_tokens", 0),
                            "completion_tokens": event["attributes"].get("usage", {}).get("completion_tokens", 0),
                            "total_tokens": event["attributes"].get("usage", {}).get("total_tokens", 0),
                            "ideal_answer": ""  # To be filled manually
                        }
                        self.eval_data.append(eval_entry)
        
        # Save to CSV for easy editing in spreadsheets
        df = pd.DataFrame(self.eval_data)
        df.to_csv(self.output_file, index=False)
        print(f"Saved {len(self.eval_data)} evaluation entries to {self.output_file}")
        
        return df

# Usage
collector = EvaluationDataCollector()
eval_df = collector.extract_from_trace_files()

class WorkflowTraceProcessor(TraceProcessor):
    def __init__(self, output_dir="workflow_traces"):
        self.output_dir = output_dir
        os.makedirs(output_dir, exist_ok=True)
        self.active_workflows = {}
        
    def process_trace(self, trace):
        workflow_id = trace.workflow_name or trace.trace_id
        
        if workflow_id not in self.active_workflows:
            self.active_workflows[workflow_id] = {
                "workflow_id": workflow_id,
                "start_time": datetime.now().isoformat(),
                "traces": [],
                "all_inputs": [],
                "all_outputs": []
            }
        
        workflow = self.active_workflows[workflow_id]
        
        # Extract all LLM interactions
        trace_summary = {
            "trace_id": trace.trace_id,
            "timestamp": datetime.now().isoformat(),
            "llm_calls": []
        }
        
        for span in trace.spans:
            for event in span.events:
                if event.name == "llm.generation":
                    llm_call = {
                        "input": event.attributes.get("input", ""),
                        "output": event.attributes.get("output", ""),
                        "model": event.attributes.get("model", "")
                    }
                    trace_summary["llm_calls"].append(llm_call)
                    workflow["all_inputs"].append(llm_call["input"])
                    workflow["all_outputs"].append(llm_call["output"])
        
        workflow["traces"].append(trace_summary)
        
        # Check if workflow is complete (you'll need to define this logic)
        if self.is_workflow_complete(trace):
            self.save_workflow(workflow_id)
            del self.active_workflows[workflow_id]
    
    def is_workflow_complete(self, trace):
        # Implement logic to determine if a workflow is done
        # Could be based on specific events, timeouts, or explicit markers
        return trace.attributes.get("workflow_complete", False)
    
    def save_workflow(self, workflow_id):
        workflow = self.active_workflows[workflow_id]
        workflow["end_time"] = datetime.now().isoformat()
        
        filename = f"{self.output_dir}/workflow_{workflow_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        with open(filename, 'w') as f:
            json.dump(workflow, f, indent=2, default=str)

def traces_to_analytics_format(trace_dir="agent_traces"):
    """Convert trace files to a format suitable for analytics tools"""
    all_events = []
    
    for trace_file in glob.glob(f"{trace_dir}/*.json"):
        with open(trace_file, 'r') as f:
            trace = json.load(f)
            
        base_info = {
            "trace_id": trace["trace_id"],
            "workflow_name": trace.get("workflow_name", ""),
            "trace_timestamp": trace["timestamp"]
        }
        
        for span in trace.get("spans", []):
            for event in span.get("events", []):
                event_record = {**base_info}
                event_record.update({
                    "span_id": span["span_id"],
                    "operation": span["operation"],
                    "event_name": event["name"],
                    "event_timestamp": event["timestamp"]
                })
                
                # Flatten attributes
                for key, value in event.get("attributes", {}).items():
                    if isinstance(value, (str, int, float, bool)):
                        event_record[f"attr_{key}"] = value
                    else:
                        event_record[f"attr_{key}"] = json.dumps(value)
                
                all_events.append(event_record)
    
    # Save as CSV for easy import into analytics tools
    df = pd.DataFrame(all_events)
    df.to_csv("traces_analytics.csv", index=False)
    
    # Also save as Parquet for efficient querying
    df.to_parquet("traces_analytics.parquet", index=False)
    
    return df

from agents import add_trace_processor, set_tracing_export_api_key
import os
set_tracing_export_api_key(os.environ["OPENAI_API_KEY"])   # keeps the dashboard view

from agents.tracing.processors import OTLPHTTPTraceSpanProcessor
add_trace_processor(
    OTLPHTTPTraceSpanProcessor(
        endpoint="http://localhost:4318/v1/traces"   # any OTLP-compatible collector
    )
)

import json, csv, glob
with open('eval_sets/today.csv','w', newline='') as f:
    w = csv.writer(f)
    w.writerow(['trace_id','input','model_output'])
    for fn in glob.glob('raw/2025-06-01*.jsonl'):
        for line in open(fn):
            span = json.loads(line)
            if span.get("attributes", {}).get("openai.event_type") == "LLM_RUN":
                w.writerow([
                    span["trace_id"],
                    span["attributes"]["openai.llm.input"],
                    span["attributes"]["openai.llm.output"]
                ])

response = client.chat.completions.create(...)
log = {
    "trace_id": response.id,
    "input": messages,
    "output": response.choices[0].message.content,
    "usage": response.usage
}
jsonl_writer.write(log)

import os
import json
from datetime import datetime

# OpenTelemetry imports
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor, ConsoleSpanExporter
from opentelemetry.sdk.resources import Resource
from opentelemetry.semconv.resource import ResourceAttributes

# We'll use a custom file exporter for simple local JSON saving.
# In a production setup, you'd send to a local OTel Collector or a database.
class CustomFileSpanExporter(ConsoleSpanExporter): # Inherit from ConsoleSpanExporter for structure, but override export
    def __init__(self, file_path="traces.jsonl"):
        self.file_path = file_path
        with open(self.file_path, 'w') as f:
            f.write("") # Clear file on init or ensure append mode

    def export(self, spans):
        # Convert OpenTelemetry spans to a dictionary format for JSON line output
        # This is a simplified conversion. A full OTel exporter would handle OTLP protobufs.
        # Here, we're just extracting key attributes for your specific needs.
        for span in spans:
            span_dict = {
                "trace_id": format(span.context.trace_id, '032x'),
                "span_id": format(span.context.span_id, '016x'),
                "parent_id": format(span.parent.span_id, '016x') if span.parent else None,
                "name": span.name,
                "start_time_unix_nano": span.start_time,
                "end_time_unix_nano": span.end_time,
                "attributes": dict(span.attributes), # Convert Attributes to dict
                "events": [{"name": event.name, "timestamp": event.timestamp, "attributes": dict(event.attributes)} for event in span.events],
                "status": {"status_code": span.status.status_code.name, "description": span.status.description},
                "resource_attributes": dict(span.resource.attributes)
            }

            with open(self.file_path, 'a') as f:
                f.write(json.dumps(span_dict) + "\n")
        return trace.export.SpanExportResult.SUCCESS


# Configure the TracerProvider
resource = Resource.create({
    ResourceAttributes.SERVICE_NAME: "my-ai-agent-app",
    ResourceAttributes.SERVICE_VERSION: "1.0.0",
})
provider = TracerProvider(resource=resource)
file_exporter = CustomFileSpanExporter()
span_processor = SimpleSpanProcessor(file_exporter)
provider.add_span_processor(span_processor)
trace.set_tracer_provider(provider)

# Get a tracer instance (though Agents SDK handles much of this automatically)
otel_tracer = trace.get_tracer(__name__)

from agents import set_trace_processors, set_tracing_export_api_key
from openai.lib.streaming._trace import TraceProcessor, TraceEvent, TraceType
import logging

# Set up basic logging to see process
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class OpenAIAgentOTELProcessor(TraceProcessor):
    """
    A custom trace processor for OpenAI Agents SDK that converts SDK trace events
    into OpenTelemetry spans and then exports them via the configured OTel provider.
    """
    def __init__(self, otel_tracer_instance):
        self.otel_tracer = otel_tracer_instance
        self.span_map = {} # To keep track of active OTel spans by OpenAI trace/span ID

    def process_trace(self, trace_event: TraceEvent) -> None:
        # Each TraceEvent from the OpenAI Agents SDK corresponds to a unit of work.
        # We map these to OpenTelemetry spans.

        otel_span_name = f"openai_agent.{trace_event.event_type.value.lower()}"
        
        # Determine parent span for OpenTelemetry hierarchy
        parent_span_context = None
        if trace_event.parent_id and trace_event.parent_id in self.span_map:
            parent_span_context = self.span_map[trace_event.parent_id].context
        elif trace_event.trace_id != trace_event.parent_id and trace_event.trace_id in self.span_map:
            # Handle cases where parent_id is for the root trace itself but not the immediate parent span
            parent_span_context = self.span_map[trace_event.trace_id].context

        with self.otel_tracer.start_as_current_span(
            otel_span_name,
            context=parent_span_context,
            kind=trace.SpanKind.INTERNAL # Most agent steps are internal
        ) as otel_span:
            # Store the OTel span for child events to reference
            self.span_map[trace_event.id] = otel_span

            # Set core attributes for the span
            otel_span.set_attribute("openai.sdk.trace_id", trace_event.trace_id)
            otel_span.set_attribute("openai.sdk.span_id", trace_event.id)
            if trace_event.parent_id:
                otel_span.set_attribute("openai.sdk.parent_id", trace_event.parent_id)
            otel_span.set_attribute("openai.sdk.event_type", trace_event.event_type.value)
            otel_span.set_attribute("openai.sdk.start_time", trace_event.start_time.isoformat())

            # Add detailed data based on event type
            if trace_event.data:
                # Convert Pydantic models in data to dicts for JSON serialization
                processed_data = {}
                for key, value in trace_event.data.items():
                    try:
                        processed_data[key] = value.model_dump() if hasattr(value, 'model_dump') else value
                    except Exception:
                        processed_data[key] = str(value) # Fallback to string if cannot dump

                otel_span.set_attribute("openai.sdk.data", json.dumps(processed_data))

                # --- CRUCIAL FOR YOUR EVALS: Extracting Input and LLM Output ---
                if trace_event.event_type == TraceType.LLM_RUN:
                    # Input (everything that went into the context window)
                    if 'input' in processed_data:
                        # For chat models, 'input' is typically a list of message dicts
                        otel_span.set_attribute("llm.input.messages_json", json.dumps(processed_data['input']))
                    # LLM Output
                    if 'output' in processed_data and 'content' in processed_data['output']:
                        otel_span.set_attribute("llm.output.content", processed_data['output']['content'])
                    if 'usage' in processed_data:
                        otel_span.set_attribute("llm.usage.prompt_tokens", processed_data['usage'].get('prompt_tokens'))
                        otel_span.set_attribute("llm.usage.completion_tokens", processed_data['usage'].get('completion_tokens'))
                        otel_span.set_attribute("llm.usage.total_tokens", processed_data['usage'].get('total_tokens'))
                    if 'model' in processed_data:
                        otel_span.set_attribute("llm.model", processed_data['model'])

                elif trace_event.event_type == TraceType.TOOL_CALL:
                    if 'tool_name' in processed_data:
                        otel_span.set_attribute("tool.name", processed_data['tool_name'])
                    if 'tool_args' in processed_data:
                        otel_span.set_attribute("tool.args_json", json.dumps(processed_data['tool_args']))

                elif trace_event.event_type == TraceType.TOOL_OUTPUT:
                    if 'output' in processed_data:
                        otel_span.set_attribute("tool.output", str(processed_data['output'])) # Tool output can be complex

                elif trace_event.event_type == TraceType.ERROR:
                    otel_span.set_status(trace.Status(trace.StatusCode.ERROR, description=str(trace_event.data)))
                    otel_span.record_exception(Exception(str(trace_event.data)))

            # Mark the end of the span if it's a terminal event for that trace_id/span_id
            # The SDK handles span closure implicitly, but this is for conceptual mapping.
            if trace_event.event_type in [TraceType.AGENT_END, TraceType.LLM_END, TraceType.TOOL_END, TraceType.ERROR]:
                 if trace_event.id in self.span_map:
                     self.span_map[trace_event.id].end()
                     del self.span_map[trace_event.id]
                 if trace_event.trace_id == trace_event.id and trace_event.trace_id in self.span_map:
                     # If it's the root trace ending, ensure its span ends
                     self.span_map[trace_event.trace_id].end()
                     del self.span_map[trace_event.trace_id]

# Instantiate and register your custom processor
agent_otel_processor = OpenAIAgentOTELProcessor(otel_tracer)

# IMPORTANT: This replaces OpenAI's default processors.
# If you still want traces in the OpenAI dashboard, you would need to also
# include OpenAI's default exporter in the list if the SDK supported multi-exporters.
# As per research, `set_trace_processors` replaces.
# If you want both, consider sending OTel traces to a collector that then forwards to OpenAI (if that's an option via OTel, which it often is).
# For now, this is for local capture ONLY.
set_trace_processors([agent_otel_processor])

# If you still want to send some basic usage data to OpenAI for cost tracking (optional),
# you might still need this line, but it won't be sending full traces to their dashboard
# if set_trace_processors replaced the default one.
# set_tracing_export_api_key(os.getenv("OPENAI_API_KEY"))

import os
from openai import OpenAI
from agents import Agent, Runner, trace, set_trace_processors
# Make sure to run the OTel and TraceProcessor setup code above first!

# Example Agent Definition and Run
client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

# Define a tool for the agent
def get_current_weather(location: str):
    """Get the current weather in a given location."""
    if "san francisco" in location.lower():
        return {"temperature": "72F", "conditions": "Sunny"}
    elif "london" in location.lower():
        return {"temperature": "55F", "conditions": "Cloudy"}
    else:
        return {"temperature": "N/A", "conditions": "Unknown"}

# Create an agent
weather_agent = Agent(
    client=client,
    model="gpt-4o",
    instructions="You are a helpful weather assistant. Use the available tools.",
    tools=[get_current_weather]
)

# Run the agent within a custom trace context for better identification
# The `with trace("your_workflow_name")` block helps categorize the top-level
# interaction in your traces.
with trace("WeatherQueryWorkflow"):
    print("Running Agent 1: What is the weather in San Francisco?")
    result1 = Runner.run_sync(weather_agent, "What is the weather in San Francisco?")
    print(f"Agent 1 Output: {result1.final_output}\n")

with trace("GeneralKnowledgeWorkflow"):
    print("Running Agent 2: What is the capital of France?")
    result2 = Runner.run_sync(weather_agent, "What is the capital of France?")
    print(f"Agent 2 Output: {result2.final_output}\n")

# You might want to explicitly flush the exporter's buffer at the end of your script
# if using a BatchSpanProcessor or similar. SimpleSpanProcessor flushes immediately.
print(f"Traces saved to traces.jsonl")

# 1. Install required packages
pip install openai-agents pandas

# 2. Set up your trace collector (add this to your main app)
from evaluation_trace_collector import setup_evaluation_collection

# Initialize evaluation data collection
collector = setup_evaluation_collection()

# 3. Run your agents normally - traces are automatically collected!
# Your existing agent code doesn't need to change

# Your existing agent code with trace collection added
collector = setup_evaluation_collection()

# Run your agents normally
result = await Runner.run(agent, "User question here")

# Data is automatically saved to CSV files in eval_data/ folder

# Define what the perfect response should be
ideal_answers = {
    "What is your return policy?": "Our return policy allows...",
    "How do I track my order?": "You can track your order by...",
    # Add more ideal responses
}

# Add them to your CSV
updated_csv = add_ideal_answers_to_csv("eval_data/your_file.csv", ideal_answers)

# Extend the EvaluationRecord dataclass to add your own fields
@dataclass 
class MyEvaluationRecord(EvaluationRecord):
    custom_score: float = 0.0
    category: str = ""
    complexity_level: int = 1

# Export all data at once
master_file = collector.export_all_data()

# Or export per trace automatically
# (happens automatically when each trace completes)