Prompt Caching

import os
import json
from cerebras.cloud.sdk import Cerebras

client = Cerebras(api_key=os.environ.get("CEREBRAS_API_KEY"))

# Mock order database
ORDERS = {
    "ORD-123456": {"status": "processing", "eta_days": 5},
}

def get_order_status(order_id: str):
    """Look up an order's status"""
    order = ORDERS.get(order_id)
    if not order:
        return {"error": "Order not found"}
    return {"order_id": order_id, "status": order["status"], "eta_days": order.get("eta_days")}

def cancel_order(order_id: str):
    """Cancel an order"""
    order = ORDERS.get(order_id)
    if not order:
        return {"error": "Order not found"}
    if order["status"] in ["shipped", "delivered"]:
        return {"error": f"Cannot cancel - order already {order['status']}"}
    
    order["status"] = "cancelled"
    order.pop("eta_days", None)
    return {"order_id": order_id, "status": "cancelled"}

tools = [
    {
        "type": "function",
        "function": {
            "name": "get_order_status",
            "description": "Look up an order status by order ID",
            "parameters": {
                "type": "object",
                "properties": {
                    "order_id": {"type": "string", "description": "Order ID (e.g., ORD-123456)"}
                },
                "required": ["order_id"]
            },
            "strict": True
        }
    },
    {
        "type": "function",
        "function": {
            "name": "cancel_order",
            "description": "Cancel an order by order ID",
            "parameters": {
                "type": "object",
                "properties": {
                    "order_id": {"type": "string", "description": "Order ID (e.g., ORD-123456)"}
                },
                "required": ["order_id"]
            },
            "strict": True
        }
    }
]

available_functions = {
    "get_order_status": get_order_status,
    "cancel_order": cancel_order
}

messages = [
    {"role": "system", "content": "You are a shopping assistant. Help users check order status and cancel orders."},
    {"role": "user", "content": "Where is my order ORD-123456?"}
]

# Turn 1 - creates cache
response = client.chat.completions.create(model="gpt-oss-120b", messages=messages, tools=tools)
print("Turn 1 usage:", response.usage)

msg = response.choices[0].message
messages.append(msg.model_dump())

if msg.tool_calls:
    for call in msg.tool_calls:
        func = available_functions[call.function.name]
        result = func(**json.loads(call.function.arguments))
        messages.append({"role": "tool", "tool_call_id": call.id, "content": json.dumps(result)})
    
    response = client.chat.completions.create(model="gpt-oss-120b", messages=messages, tools=tools)
    messages.append(response.choices[0].message.model_dump())
    print("Turn 1 response:", response.choices[0].message.content)

# Turn 2 - uses cache
messages.append({"role": "user", "content": "Please cancel it, I ordered by mistake."})
response = client.chat.completions.create(model="gpt-oss-120b", messages=messages, tools=tools)
print("\nTurn 2 usage:", response.usage)

msg = response.choices[0].message
messages.append(msg.model_dump())

if msg.tool_calls:
    for call in msg.tool_calls:
        func = available_functions[call.function.name]
        result = func(**json.loads(call.function.arguments))
        messages.append({"role": "tool", "tool_call_id": call.id, "content": json.dumps(result)})
    
    response = client.chat.completions.create(model="gpt-oss-120b", messages=messages, tools=tools)
    print("Turn 2 response:", response.choices[0].message.content)