pip install arize arize-phoenix arize-phoenix-evals pandas openai nest_asyncio

from typing import Dict, Any
import pandas as pd

# A reusable helper that applies both trace-level and span-level filters
from agent_trajectory_utils import filter_spans_by_trace_criteria  # provided in the notebook

trajectory_spans = filter_spans_by_trace_criteria(
    df            = df,
    trace_filters = {"name": {"contains": "searchrouter"}},      # tailor to your app
    span_filters  = {"attributes.openinference.span.kind": {"==": "LLM"}},
)

from agent_trajectory_utils import (
    extract_tool_calls,
    prepare_trace_data_for_evaluation,
)

# Parse `attributes.llm.output_messages` → list of {name, arguments}
trajectory_spans["tool_calls"] = trajectory_spans[
    "attributes.llm.output_messages"
].apply(extract_tool_calls)

# Collapse every trace into a single row that contains its ordered tool calls
trace_df = prepare_trace_data_for_evaluation(
    df = trajectory_spans,
    extract_cols = {
        "tool_calls": "tool_calls",
        "attributes.llm.tools": "attributes.llm.tools",           # reference tool schema
        "attributes.input.value": "attributes.input.value",       # original user input
    },
)

TRAJECTORY_ACCURACY_PROMPT = """
You are a helpful AI bot that checks whether an AI agent's internal trajectory is accurate and effective.

You will be given:
1. The agent's actual trajectory of tool calls
2. The user input that initiated the trajectory
3. The definition of each tool that can be called

An accurate trajectory:
- Progresses logically from step to step
- Uses the right tools for the task
- Is reasonably efficient (no unnecessary detours)

##
Actual Trajectory:
{tool_calls}

User Input:
{attributes.input.value}

Tool Definitions:
{attributes.llm.tools}
##

Respond with **exactly** one word: `correct` or `incorrect`.
- `correct` → trajectory adheres to the rubric and achieves the task.
- `incorrect` → trajectory is confusing, inefficient, or fails the task.
"""

from phoenix.evals import llm_classify, OpenAIModel
import nest_asyncio, os

nest_asyncio.apply()

model = OpenAIModel(
    api_key = os.environ["OPENAI_API_KEY"],
    model   = "gpt-4o-mini",
    temperature = 0.0,
)

rails = ["correct", "incorrect"]
results = llm_classify(
    dataframe           = trace_df,
    template            = TRAJECTORY_ACCURACY_PROMPT,
    model               = model,
    rails               = rails,
    provide_explanation = True,   # add a free-text rationale for debugging
    verbose             = False,
)