class JokeIdea(BaseModel):
    setup: str
    contradiction: str
    punchline: str

class QueryToIdea(dspy.Signature):
    """Generate a joke idea with setup, contradiction, and punchline."""
    query = dspy.InputField()
    joke_idea: JokeIdea = dspy.OutputField()

class IdeaToJoke(dspy.Signature):
    """Convert a joke idea into a full comedian delivery."""
    joke_idea: JokeIdea = dspy.InputField()
    joke = dspy.OutputField()

class JokeGenerator(dspy.Module):
    def __init__(self):
        self.query_to_idea = dspy.Predict(QueryToIdea)
        self.idea_to_joke = dspy.Predict(IdeaToJoke)
        
        self.query_to_idea.set_lm(lm=dspy.LM("openai/gpt-4.1-mini"))
        self.idea_to_joke.set_lm(lm=dspy.LM("openai/gpt-4.1"))

    
    def forward(self, query):
        idea = self.query_to_idea(query=query)
        joke = self.idea_to_joke(joke_idea=idea.joke_idea)
        return joke