evaluate.py

import os
import json
import asyncio
import base64
import httpx
from typing import List, Dict, Any
from openai import AsyncOpenAI

# OpenTelemetry
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider

# App imports (triggers setup_telemetry)
from app import query_agent

tracer = trace.get_tracer("evaluation_script")

# Initialize OpenAI Client for Evaluation
aclient = AsyncOpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

class LangfuseClient:
    def __init__(self):
        self.secret_key = os.environ.get("LANGFUSE_SECRET_KEY")
        self.public_key = os.environ.get("LANGFUSE_PUBLIC_KEY")
        self.base_url = os.environ.get("LANGFUSE_BASE_URL", "https://cloud.langfuse.com").rstrip("/")
        
        if not (self.secret_key and self.public_key):
            print("⚠ Langfuse credentials missing. Scoring disabled.")
            self.enabled = False
            return
            
        self.enabled = True
        auth_str = f"{self.public_key}:{self.secret_key}"
        auth_bytes = auth_str.encode("ascii")
        base64_auth = base64.b64encode(auth_bytes).decode("ascii")
        self.headers = {
            "Authorization": f"Basic {base64_auth}",
            "Content-Type": "application/json"
        }

    def score_trace(self, trace_id: str, name: str, value: float, comment: str = None):
        if not self.enabled:
            return
            
        url = f"{self.base_url}/api/public/scores"
        payload = {
            "traceId": trace_id,
            "name": name,
            "value": value,
            "comment": comment
        }
        
        try:
            # Synchronous call for simplicity in this script
            resp = httpx.post(url, json=payload, headers=self.headers, timeout=10.0)
            if resp.status_code not in (200, 201):
                print(f"   ⚠ Failed to log score {name}: {resp.status_code} - {resp.text}")
        except Exception as e:
            print(f"   ⚠ Error logging score: {e}")

async def evaluate_helpfulness(question: str, answer: str) -> dict:
    """Evaluates if the answer is helpful using LLM."""
    prompt = f"""
    You are an expert evaluator. Rate the helpfulness of the AI response to the user question.
    
    Question: {question}
    Response: {answer}
    
    Score 0.0 to 1.0 (1.0 is most helpful).
    Provide reasoning.
    
    Output JSON: {{"score": float, "reason": "string"}}
    """
    try:
        response = await aclient.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": prompt}],
            response_format={"type": "json_object"},
            temperature=0
        )
        return json.loads(response.choices[0].message.content)
    except Exception as e:
        print(f"   ⚠ Eval Error: {e}")
        return {"score": 0.0, "reason": "Evaluation failed"}

async def evaluate_faithfulness(question: str, answer: str, context: str) -> dict:
    """Evaluates if the answer is faithful to the context."""
    if not context:
        return {"score": 0.5, "reason": "No context available for faithfulness check."}
        
    prompt = f"""
    You are an expert evaluator. Rate the faithfulness of the AI response to the retrieved context.
    
    Context:
    {context[:10000]}... (truncated)

    Question: {question}
    Response: {answer}
    
    Score 0.0 to 1.0 (1.0 is fully supported by context).
    If the response contains information NOT in the context (hallucination), penalize heavily.
    
    Output JSON: {{"score": float, "reason": "string"}}
    """
    try:
        response = await aclient.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": prompt}],
            response_format={"type": "json_object"},
            temperature=0
        )
        return json.loads(response.choices[0].message.content)
    except Exception as e:
        print(f"   ⚠ Faithfulness Eval Error: {e}")
        return {"score": 0.0, "reason": "Evaluation failed"}

async def evaluate_trajectory(question: str, tool_calls: List[str], rubric: str) -> dict:
    """Evaluates if the tool usage followed the rubric."""
    prompt = f"""
    You are an expert evaluator. Rate the agent's execution trajectory against the rubric.
    
    Rubric: {rubric}
    
    Question: {question}
    Tool Sequence: {json.dumps(tool_calls)}
    
    Score 0.0 to 1.0 (1.0 is perfect adherence).
    Did it skip required steps? Did it use irrelevant tools?
    
    Output JSON: {{"score": float, "reason": "string"}}
    """
    try:
        response = await aclient.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": prompt}],
            response_format={"type": "json_object"},
            temperature=0
        )
        return json.loads(response.choices[0].message.content)
    except Exception as e:
        print(f"   ⚠ Trajectory Eval Error: {e}")
        return {"score": 0.0, "reason": "Evaluation failed"}

def extract_context_and_tools(agent_result) -> tuple[str, List[str]]:
    """Extracts retrieved text and tool names from AgentResult."""
    context = []
    tool_calls = []
    
    if not hasattr(agent_result, 'trace') or not agent_result.trace:
        return "", []

    for span in agent_result.trace.spans:
        # Check for tool execution spans (simplified check)
        if hasattr(span, 'span_type') and str(span.span_type) == 'tool_execution':
             # Tool Name
             tool_name = span.tool_call.name
             tool_calls.append(tool_name)
             
             # Context from Search/Load Tools
             if 'confluence' in tool_name or 'get_application_summary' in tool_name or 'compare' in tool_name:
                 context.append(f"Source ({tool_name}): {span.tool_result.content}")

    return "\n\n".join(context), tool_calls

async def run_evaluation():
    print("🚀 Starting Evaluation Pipeline (Custom LLM Evals: Helpfulness, Faithfulness, Trajectory)...")
    
    # Initialize Client
    lf_client = LangfuseClient()
    
    # Load dataset
    try:
        with open("evaluation/dataset.json", "r") as f:
            dataset = json.load(f)
    except FileNotFoundError:
        print("❌ evaluation/dataset.json not found.")
        return
    
    print(f"📋 Loaded {len(dataset)} test cases.")

    results = []

    for case in dataset:
        case_id = case["id"]
        question = case["question"]
        expected_key_points = case["expected_answer_key_points"]
        
        print(f"\n🧪 Running Case: {case_id}")
        
        # Start a span for this evaluation case
        with tracer.start_as_current_span(f"Eval: {case_id}") as span:
            # Get Trace ID (OTel stores it as int, needs hex conversion)
            trace_id_int = span.get_span_context().trace_id
            trace_id_hex = "{:032x}".format(trace_id_int)
            
            # Add metadata
            span.set_attribute("evaluation.case_id", case_id)
            span.set_attribute("evaluation.question", question)
            
            # 1. Run Agent (Request Full Result)
            # The agent's internal spans will be nested under this span automatically
            result_obj = query_agent(question, return_full_result=True)
            answer = str(result_obj)
            
            # Extract Internals
            context_text, tool_sequence = extract_context_and_tools(result_obj)
            # Log extracted data for debug
            span.set_attribute("evaluation.tool_sequence", json.dumps(tool_sequence))

            # 2. Run Evaluators & Log Scores
            
            # A. Helpfulness
            help_res = await evaluate_helpfulness(question, answer)
            lf_client.score_trace(trace_id=trace_id_hex, name="Helpfulness", value=help_res["score"], comment=help_res["reason"])
            print(f"   ✅ Helpfulness: {help_res['score']:.2f}")

            # B. Faithfulness
            faith_res = await evaluate_faithfulness(question, answer, context_text)
            lf_client.score_trace(trace_id=trace_id_hex, name="Faithfulness", value=faith_res["score"], comment=faith_res["reason"])
            print(f"   📖 Faithfulness: {faith_res['score']:.2f}")

            # C. Trajectory
            rubric = "1. Retrieve data (summary). 2. Analyze specifics/compare. 3. Check compliance if relevant. 4. Explain."
            traj_res = await evaluate_trajectory(question, tool_sequence, rubric)
            lf_client.score_trace(trace_id=trace_id_hex, name="Trajectory", value=traj_res["score"], comment=traj_res["reason"])
            print(f"   👣 Trajectory:  {traj_res['score']:.2f} ({len(tool_sequence)} tools)")

            # D. Goal Success
            hits = 0
            if expected_key_points:
                for point in expected_key_points:
                    if point.lower() in answer.lower():
                        hits += 1
                    elif any(word in answer.lower() for word in point.split() if len(word) > 4):
                         hits += 0.5 
                success_rate = min(1.0, hits / len(expected_key_points))
            else:
                success_rate = 1.0
            
            lf_client.score_trace(trace_id=trace_id_hex, name="Goal Success", value=success_rate, comment=f"Matched {hits}/{len(expected_key_points)} key points")
            print(f"   🎯 Goal Success: {success_rate:.2f}")

            results.append({
                "case_id": case_id,
                "trace_id": trace_id_hex,
                "helpfulness": help_res["score"],
                "faithfulness": faith_res["score"],
                "trajectory": traj_res["score"],
                "goal_success": success_rate
            })
            
    # Summary
    print("\n📊 Evaluation Summary")
    print(json.dumps(results, indent=2))

if __name__ == "__main__":
    asyncio.run(run_evaluation())