from getpass import getpass
import phoenix as px
from openinference.instrumentation.haystack import HaystackInstrumentor
from phoenix.otel import register

px.launch_app().view()

tracer_provider = register()

# Use Phoenix's autoinstrumentor to automatically track traces from Haystack
HaystackInstrumentor().instrument(tracer_provider=tracer_provider, skip_dep_check=True)

import os

os.environ["OPENAI_API_KEY"] = getpass("Enter your OpenAI API key: ")

from haystack.dataclasses import ChatMessage
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.components.builders import ChatPromptBuilder
from haystack import Pipeline

messages = [
    ChatMessage.from_system(
        "You are a toy maker elf. Your job is to make toys for the nice kids on the nice list. If the child is on the naughty list, give them a 'Rabbit R1'. {{name}} is on the {{list}} list"
    ),
    ChatMessage.from_user(
        "Create a toy for {{name}} that they will like. {{name}} is {{age}} years old and likes {{likes}} and dislikes {{dislikes}}."
    ),
]

builder = ChatPromptBuilder(messages)
chat_generator = OpenAIChatGenerator(model="gpt-4o")

pipeline = Pipeline()
pipeline.add_component("builder", builder)
pipeline.add_component("chat_generator", chat_generator)

pipeline.connect("builder", "chat_generator")


def make_toy(child):
    return pipeline.run({"builder": {**child}})["chat_generator"]["replies"]

pipeline.show()

for child in children:
    make_toy(child)

llm_judge_prompt = """
Evaluate the toy for this child, based on their likes and dislikes

All children on the naughty list get a 'Rabbit R1'. Any other toy given to a naughty child is incorrect.

Respond with a single word: 'correct' or 'incorrect'. Also include a short explanation for your answer.

Description of the child: {description}
Toy: {toy}

*****
Example output:
label: 'correct'
explanation: 'The toy is a Lego set, which is one of the child's likes.'
*****
"""

import phoenix as px

# Download the traces from Phoenix
# HINT: https://docs.arize.com/phoenix/evaluation/how-to-evals/evaluating-phoenix-traces#download-trace-dataset-from-phoenix
spans_df = px.Client().get_spans_dataframe()
spans_df

spans_df.span_kind.value_counts()

spans_df[spans_df.span_kind == "LLM"]

df = spans_df[spans_df.span_kind == "LLM"][
    ["attributes.input.value", "attributes.output.value"]
]

df.head()

import json

json.loads(df["attributes.input.value"].values[0])["messages"]

json.loads(df["attributes.output.value"].values[0])["replies"]

df["description"] = df["attributes.input.value"].apply(
    lambda x: json.loads(x)["messages"]
)
df["toy"] = df["attributes.output.value"].apply(lambda x: json.loads(x)["replies"])

df.drop(["attributes.input.value", "attributes.output.value"], axis=1, inplace=True)

df.head()

import nest_asyncio

nest_asyncio.apply()

from phoenix.evals import (
    llm_classify,
    OpenAIModel,  # can swap for another model supported by Phoenix or run open-source models through LiteLLM and Ollama: https://docs.arize.com/phoenix/evaluation/evaluation-models
)
from phoenix.evals.templates import ClassificationTemplate

# Evaluate the traces with the LLM Judge
# HINT: https://docs.arize.com/phoenix/evaluation/how-to-evals/bring-your-own-evaluator#categorical-llm_classify
# HINT: For evaluation, try using a different language model than the one you used for toy matching
# referred: https://github.com/deepset-ai/haystack/discussions/8579#discussioncomment-11649855
eval_results = llm_classify(
    dataframe=df,
    template=ClassificationTemplate(
        template=llm_judge_prompt,
        rails=["correct", "incorrect"],
        explanation_template="Explanation: ",
    ),
    model=OpenAIModel(model="gpt-4o-mini"),
    rails=["correct", "incorrect"],
    provide_explanation=True,
)
eval_results["score"] = eval_results["label"].apply(
    lambda x: 1 if x == "correct" else 0
)
eval_results.head()

eval_results.score.value_counts()

(17 / 20) * 100

from phoenix.trace import SpanEvaluations

# Upload results into Phoenix
# HINT: https://docs.arize.com/phoenix/evaluation/how-to-evals/evaluating-phoenix-traces#download-trace-dataset-from-phoenix
px.Client().log_evaluations(
    SpanEvaluations(eval_name="evaluate_toy", dataframe=eval_results)
)