num_samples = 5

class JokeJudge(dspy.Signature):
    """Given a list of joke ideas, you must pick the best joke"""
    joke_ideas: list[JokeIdeas] = dspy.InputField()
    best_idx: int = dspy.OutputField(
        le=num_samples,
        ge=1,
        description="The index of the funniest joke")

class ConditionalJokeGenerator(dspy.Module):
    def __init__(self):
        self.query_to_idea = dspy.ChainOfThought(QueryToIdea)
        self.judge = dspy.ChainOfThought(JokeJudge)
        self.idea_to_joke = dspy.ChainOfThought(IdeaToJoke)
    
    async def forward(self, query):
        # Generate multiple ideas in parallel
        ideas = await asyncio.gather(*[
            self.query_to_idea.acall(query=query) 
            for _ in range(num_samples)
        ])
        
        # Judge and rank ideas
        best_idx = (await self.judge.acall(joke_ideas=ideas)).best_idx
        
        # Select best idea and generate final joke
        best_idea = ideas[best_idx]

        # Convert from idea to joke
        return await self.idea_to_joke.acall(joke_idea=best_idea)