from ddtrace.llmobs import LLMObs
from agents import Agent, ModelSettings, function_tool

@function_tool
def add_numbers(a: int, b: int) -> int:
    """
    Adds two numbers together.
    """
    return a + b

@function_tool
def subtract_numbers(a: int, b: int) -> int:
    """
    Subtracts two numbers.
    """
    return a - b


# Lista de herramientas disponibles para el Agent
math_tutor_agent = Agent(
    name="Math Tutor",
    handoff_description="Specialist agent for math questions",
    instructions="You provide help with math problems. Please use the tools to find the answer.",
    model="o3-mini",
    tools=[
        add_numbers, subtract_numbers
    ],
)

history_tutor_agent = Agent(
    name="History Tutor",
    handoff_description="Specialist agent for history questions",
    instructions="You provide help with history problems.",
    model="o3-mini",
)

# El agent de triaje decide a qué agent especializado entregar la tarea, otro tipo de selección de herramientas cubierto por esta evaluación.
triage_agent = Agent(  
    'openai:gpt-4o',
    model_settings=ModelSettings(temperature=0),
    instructions='What is the sum of 1 to 10?',  
    handoffs=[math_tutor_agent, history_tutor_agent],
)

import os

from ddtrace.llmobs import LLMObs
from pydantic_ai import Agent


# Definir herramientas como funciones regulares con sugerencias de tipo
def add_numbers(a: int, b: int) -> int:
    """
    Adds two numbers together.
    """
    return a + b


def subtract_numbers(a: int, b: int) -> int:
    """
    Subtracts two numbers.
    """
    return a - b


def multiply_numbers(a: int, b: int) -> int:
    """
    Multiplies two numbers.
    """
    return a * b


def divide_numbers(a: int, b: int) -> float:
    """
    Divides two numbers.
    """
    return a / b


# Activar LLMObs
LLMObs.enable(
    ml_app="jenn_test",
    api_key=os.environ["DD_API_KEY"],
    site=os.environ["DD_SITE"],
    agentless_enabled=True,
)


# Crear el agent Math Tutor con herramientas
math_tutor_agent = Agent(
    'openai:gpt-5-nano',
    instructions="You provide help with math problems. Please use the tools to find the answer.",
    tools=[add_numbers, subtract_numbers, multiply_numbers, divide_numbers],
)

# Crea el agent History Tutor (nota: gpt-5-nano no existe, utiliza gpt-4o-mini)
history_tutor_agent = Agent(
    'openai:gpt-5-nano',
    instructions="You provide help with history problems.",
)

# Crear el agent de triaje
# Nota: pydantic_ai gestiona los traspasos de forma diferente; normalmente se utiliza result_type 
# o lógica personalizada para enrutar entre agents
triage_agent = Agent(
    'openai:gpt-5-nano',
    instructions=(
        'DO NOT RELY ON YOUR OWN MATHEMATICAL KNOWLEDGE, '
        'MAKE SURE TO CALL AVAILABLE TOOLS TO SOLVE EVERY SUBPROBLEM.'
    ),
    tools=[add_numbers, subtract_numbers, multiply_numbers, divide_numbers],
)


# Ejecutar el agent de forma sincrónica
result = triage_agent.run_sync(
    '''
    Help me solve the following problem:
    What is the sum of the numbers between 1 and 100?
    Make sure you list out all the mathematical operations (addition, subtraction, multiplication, division) in order before you start calling tools in that order.
    '''
)