import openai

from agentdojo.agent_pipeline import (
    AgentPipeline,
    InitQuery,
    OpenAILLM,
    PromptInjectionDetector,
    SystemMessage,
    ToolsExecutionLoop,
    ToolsExecutor,
)
from agentdojo.functions_runtime import FunctionsRuntime

# --8<-- [start:init_pipeline]
llm = OpenAILLM(openai.OpenAI(), "gpt-4o-2024-05-13")  # (1)!
tools_loop = ToolsExecutionLoop(  # (2)!
    [
        ToolsExecutor(),
        PromptInjectionDetector(),  # (3)!
        llm,  # (4)!
    ]
)
tools_pipeline = AgentPipeline(  # (5)!
    [
        SystemMessage("You are a helpful assistant."),
        InitQuery(),
        llm,
        tools_loop,
    ]
)
# --8<-- [end:init_pipeline]

runtime = FunctionsRuntime()


@runtime.register_function
def send_email(address: str, subject: str, body: str) -> str:
    """Sends an email.

    :param address: the recipient.
    :param subject: the subject.
    :param body: the body.
    """
    return "Email sent successfully!"


# --8<-- [start:run_pipeline]
prompt = "Can you please send an email to 'email@example.com?'"
print(tools_pipeline.query(prompt, runtime))

# --8<-- [end:run_pipeline]
