Getting started with agent simulation

from typing import Dict

from snowglobe.client import CompletionRequest, CompletionFunctionOutputs
from snowglobe.tools import snowglobe_tool
import litellm
import json
import time

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "get_customer_info",
            "parameters": {
                "type": "object",
                "description": "Get information about a customer by name or email or phone",
                "properties": {
                    "email_address": {
                        "type": "string",
                        "description": "The email address of the customer to look up with @ and fully qualified domain such as bob@example.com"
                    }
                },
                "required": []
            },
            "returns": {
                "type": "object",
                "description": "Information about the customer",
                "properties": {
                    "customer_id": {
                        "type": "string",
                        "format": "uuid",
                        "description": "The unique identifier for the customer"
                    },
                    "name": {
                        "type": "string",
                        "description": "The name of the customer"
                    },
                    "email": {
                        "type": "string",
                        "description": "The email address of the customer"
                    }
                }
            },
            "examples": [{
                "input": {
                    "email_address": "bsmith@example.com"
                },
                "output": {
                    "customer_id": "3d03ebc1-eb76-498e-9109-c911840e2ac1",
                    "name": "Bob Smith",
                    "email": "bsmith@example.com"
                }
            }]
        }
    },
    {
        "type": "function",
        "function": {
            "name": "get_order_status",
            "parameters": {
                "type": "object",
                "description": "Get the status of a pizza order by order ID",
                "properties": {
                    "order_id": {
                        "type": "string",
                        "description": "The unique identifier for the pizza order"
                    }
                },
                "required": ["order_id"]
            },
            "returns": {
                "type": "object",
                "description": "The status of the pizza order",
                "properties": {
                    "order_id": {
                        "type": "string",
                        "format": "uuid",
                        "description": "The unique identifier for the pizza order"
                    },
                    "status": {
                        "type": "string",
                        "description": "The current status of the order (e.g., 'preparing', 'baking', 'out for delivery', 'delivered')"
                    },
                    "estimated_delivery_time": {
                        "type": "string",
                        "description": "The estimated delivery time for the order"
                    }
                }
            },
            "examples": [{
                "input": {"order_id": "3d03ebc1-eb76-498e-9109-c911840e2ac3"},
                "output": {
                    "order_id": "3d03ebc1-eb76-498e-9109-c911840e2ac3",
                    "status": "out for delivery",
                    "estimated_delivery_time": "2024-06-01T18:30:00Z"
                }
            }]
        }
    }
]

# Snowglobe-Connect introspects this function at startup to learn your tools
def tool_defs():
    return TOOLS


# --- Tool implementations ---

CUSTOMERS = {
    "c001": {"customer_id": "c001", "name": "John Doe", "email": "johndoe@example.com"},
    "c002": {"customer_id": "c002", "name": "Jane Smith", "email": "janesmith@example.com"},
    "c003": {"customer_id": "c003", "name": "Bob Johnson", "email": "bobj@example.com"},
}

ORDERS = {
    "ord-1001": {"order_id": "ord-1001", "status": "delivered", "total": 29.98},
    "ord-1002": {"order_id": "ord-1002", "status": "out for delivery", "estimated_delivery_time": "2026-03-26T13:00:00Z", "total": 17.98},
    "ord-1003": {"order_id": "ord-1003", "status": "baking", "estimated_delivery_time": "2026-03-26T18:15:00Z", "total": 14.49},
}


@snowglobe_tool
def get_order_status(*, order_id: str) -> dict:
    order = ORDERS.get(order_id)
    if not order:
        return {"error": f"No order found with id {order_id}"}
    return {
        "order_id": order["order_id"],
        "status": order["status"],
        "estimated_delivery_time": order.get("estimated_delivery_time", "N/A"),
        "total": order["total"],
    }


@snowglobe_tool
def get_customer_info(*, email_address: str = "") -> dict:
    for c in CUSTOMERS.values():
        if email_address and email_address.lower() == c["email"].lower():
            return json.dumps({"customer_id": c["customer_id"], "name": c["name"], "email": c["email"]})
    return json.dumps({"error": "Customer not found"})


TOOLS_MAP = {
    "get_order_status": get_order_status,
    "get_customer_info": get_customer_info,
}


# --- Entry point for Snowglobe ---

def completion(request: CompletionRequest) -> CompletionFunctionOutputs:
    messages = request.to_openai_messages()
    response = main(messages)
    return CompletionFunctionOutputs(response=response)


def main(messages: list[Dict]) -> str:
    system_prompt = """You are a customer support agent for Amattos Pizza Restaurant.
ONLY pass valid UUID4 strings as order_id, customer_id, product_id, article_id, and refund_id when they are an argument.
If you do not know the required id field make another tool call to get it or ask the customer for it.
If you do not know the value for a tool call do not make up a value or pass in unknown or a variation of it as an argument.
Always try to get the customers name or email or phone number and use that to look up their customer_id with the get_customer_info tool before calling any other.
Use tools available to assist customers."""

    messages = [{"role": "system", "content": system_prompt}] + messages

    while True:
        response = litellm.completion(
            model="gpt-5-nano",
            messages=messages,
            tools=TOOLS,
        )
        message = response.choices[0].message

        if not message.tool_calls:
            return message.content

        messages.append(message)
        for tool_call in message.tool_calls:
            if tool_call.function.name not in TOOLS_MAP:
                result = f"Error: Tool {tool_call.function.name} not found"
            else:
                result = TOOLS_MAP[tool_call.function.name](
                    **json.loads(tool_call.function.arguments)
                )

            messages.append({
                "role": "tool",
                "tool_call_id": tool_call.id,
                "content": str(result),
            })