from modaic import PrecompiledAgent, PrecompiledConfig
import dspy

class ReceiptClassifierConfig(
    PrecompiledConfig
):  # Configurable for each agent instance
    lm: str = "openai/gpt-4o-mini"  # Set OPENAI_API_KEY
    max_tokens: int = 500

class ReceiptClassifier(PrecompiledAgent):  # Extends dspy.module()
    config: ReceiptClassifierConfig

    def __init__(self, config: ReceiptClassifierConfig, **kwargs):
        super().__init__(config, **kwargs)

        self.extract = dspy.Predict("image -> total_cost")
        lm = dspy.LM(config.lm, max_tokens=config.max_tokens)
        self.extract.set_lm(lm)

    def forward(self, image_url: str) -> str:
        receipt_image = dspy.Image(url=image_url)
        return self.extract(image=receipt_image)

agent = ReceiptClassifier(ReceiptClassifierConfig())


if __name__ == "__main__":
    results = agent("https://img.freepik.com/free-vector/receipt-template-collection-with-realistic-design_23-2147922106.jpg?semt=ais_hybrid&w=740&q=80")
    print(results)
    agent.push_to_hub("farouk1/receipt-classifier", with_code=True) # Set MODAIC_TOKEN