Structured Ouput with Claudette

from claudette.core import *
import claudette_pydantic
from pydantic import BaseModel, Field, create_model
from typing import List, Union, Literal, Type
from textwrap import dedent
import json

class Pet(BaseModel):
    name: str
    price: float = Field(description="Full price of the Pet")

class Dog(Pet):
    type: Literal["dog"] = "dog"
    breed: str

class Fish(Pet):
    type: Literal["fish"] = "fish"
    species: str
    tank_size: float

dogs = [Dog(name="Buddy", breed="Labrador", price=500), Dog(name="Max", breed="German Shepherd", price=600)]
fish = [Fish(name="Nemo", species="Clownfish", price=110, tank_size=20), Fish(name="Bubbles", species="Betta", price=300, tank_size=5)]

class SalesPersonMessage(BaseModel):
    "The salesperson should first talk to the customer to find their budget and negotiate the best deal for a pet"
    type: Literal["message"] = "message"
    content: str

class SalespersonOffer(BaseModel):
    "Send a SellOffer once a pet at a proposed price has informally been agreed upon."
    type: Literal["offer"] = "offer"
    pet: Union[Dog, Fish] = Field(..., discriminator='type')
    price: float

class SalespersonAction(BaseModel):
    action: Union[SalesPersonMessage, SalespersonOffer] = Field(..., discriminator='type')

class SalespersonInfo(BaseModel):
    name: str
    inventory: List[Union[Dog, Fish]] = Field(..., description="List of pets available for sale")


    def create_custom_action(self) -> Type[SalespersonAction]:
        CustomSalespersonAction = create_model(
            "CustomSalespersonAction",
            __base__=SalespersonAction,
        )

        inventory_pretty = '- ' + '\n- '.join(repr(o) for o in self.inventory)

        CustomSalespersonAction.__doc__ = dedent(f"""\
        Represents a salesperson in the pet store simulation.

        The salesperson has access to a list of pets (dogs and fish) available in the store,
        but does not know the customer's budget or preferences. The salesperson's goal is
        to make a sale to maximizes profit.

        ## Discounts

        May offer a 10% discount on all pets
        25% on FISH ONLY

        ## Sales

        You should sell the pet included ONLY. Do not sell accessories, only the items currently in inventory.

        ## Details

        Perform an action based on conversation history and salesperson details.

        Salesperson details:
        Name: {self.name}

        Inventory:
        {inventory_pretty}""")

        return CustomSalespersonAction

salesperson = SalespersonInfo(name="Bob", inventory=dogs + fish)

class CustomerOfferResponse(BaseModel):
    """Determine whether to accept the sale.
    Accept only if the offer is consistent to a previously agreed sale.
    this may **only** be given as an immediate response to a SellOffer"""
    type: Literal["offer"] = "offer"
    accept: bool

class CustomerOfferAction(BaseModel):
    action: CustomerOfferResponse

class CustomerMessage(BaseModel):
    type: Literal["message"] = "message"
    content: str

class CustomerLeave(BaseModel):
    "If all other avenues are exhausted, nothing in the store is acceptable, you may leave the store"
    type: Literal["leave"] = "leave"

class CustomerAction(BaseModel):
    action: Union[CustomerMessage, CustomerLeave] = Field(..., discriminator='type')


class CustomerInfo(BaseModel):
    name: str
    pet_preference: str = Field(..., description="The type of pet the customer prefers ('dog' or 'fish')")
    budget: float = Field(..., description="The maximum amount the customer is willing to spend")

    def create_custom_action(self, offer=False) -> Type[CustomerAction]:
        CustomCustomerAction = create_model(
            "CustomCustomerAction",
            __base__=CustomerOfferAction if offer else CustomerAction,
        )

        CustomCustomerAction.__doc__ = dedent(f"""\
        Represents a customer in the pet store simulation.

        The customer has a pet preference and a budget, but does not know what pets
        the store has in stock or their prices. The customer's goal is to negotiate
        the best deal for their preferred type of pet within their budget.

        Perform an action based on conversation history and customer details.

        Customer Information:
        - Name: {self.name}
        - Pet Preference: {self.pet_preference}
        - Budget: ${self.budget}
        """)

        return CustomCustomerAction


pet_preference = f"""\
Alice really wants a dog, but has a hard limit of 100 dollars.
She may also be open to buying a fish"""

customer = CustomerInfo(name="Alice", pet_preference=pet_preference, budget=100)

sp = """\
You are participating in a role-playing simulation of a pet store negotiation.
If you are the customer, your goal is to purchase your preferred pet within your budget.
If you are the salesperson, your goal is to make a profitable sale while ensuring customer satisfaction.
Be friendly and professional, but also try to get the best deal for yourself.

The negotiation continues until a sale is agreed upon or one party decides to end the conversation.
You should be polite, but do not discuss the specifics of the pets, work out a sale based on price and pet preferences. Each response should be 3 sentences max.
"""


def chat_loop(model, steps): # similar to Chat.toolloop
    chat = Chat(model, sp=sp)
    responses = []
    res = None
    for i in range(steps):
        if i == 0: # prompt required to switch to assistant
            pr = "A customer walks into a store"
        else:
            pr = None

        if i % 2 == 0:
            role = "salesperson"
            s = salesperson.create_custom_action()
        else:
            role = "customer"
            if res is not None and res.action.type == 'offer':
                offer = True
            else:
                offer = False
            s = customer.create_custom_action(offer)

        res = chat.struct(s, treat_as_output=False, pr=pr)
        responses.append(res)
        if res.action.type == 'message':
            print(f"{role}: {res.action.content}\n")
        if i % 2 == 1:
            if res.action.type == 'offer' and res.action.accept:
                return responses[-2].action
            elif res.action.type == 'leave':
                return responses[-1].action
    return None

def sim(model, steps):
    print(f"Model: {model} | Steps: {steps}\n=====\n")
    res = chat_loop(model, steps)
    if res is None:
        print("Timeout")
    elif res.type == 'offer':
        print("Sold:", repr(res))
    elif res.type == 'leave':
        print("Customer left.")
    else:
        raise ValueError("invalid return", res)

sim(models[-1], 20)