import dspy
from pydantic import BaseModel, Field

class NutritionalInfo(BaseModel):
    name: str
    calories: float
    protein: float
    carbs: float
    fats: float
    micros: dict | None = None

class ExtractNutrition(dspy.Signature):
    """Extract nutritional information from a food description."""
    description: str = dspy.InputField(desc="Description of the food or meal")
    nutritional_info: NutritionalInfo = dspy.OutputField(desc="Nutritional information as a structured object")

class NutritionModule(dspy.Module):
    def __init__(self):
        super().__init__()
        self.extract = dspy.ChainOfThought(ExtractNutrition)

    def forward(self, description: str):
        return self.extract(description=description)

nutrition_module = NutritionModule()