dspy-types.py

import dspy
from pydantic import BaseModel
from typing import List

class State(BaseModel):
    name: str
    abbreviation: str
    capital: str

class States(BaseModel):
    states: List[State]

class ExampleSignature(dspy.Signature):
    question = dspy.InputField()
    states = dspy.OutputField(desc=f"The answer written as a single JSON object using the schema {States.schema()}")

def main():
    lm = dspy.OpenAI(model='gpt-3.5-turbo-instruct', max_tokens=400)
    dspy.settings.configure(lm=lm, trace=["Test"])

    program = dspy.Predict(ExampleSignature)
    res = program(question="What are the three largest states in the US?")

    try:
        states = States.parse_raw(res.states)
    except ValueError as e:
        dspy.Assert(False, e)

    print(states.states)
    lm.inspect_history(n=3)


if __name__ == '__main__':
    main()

#[
# State(name='Alaska', abbreviation='AK', capital='Juneau'),
# State(name='Texas', abbreviation='TX', capital='Austin'),
# State(name='California', abbreviation='CA', capital='Sacramento')
#]