from modaic import PrecompiledAgent, PrecompiledConfig
from agent import SignatureGenerator
import dspy


class PromptToSignatureConfig(PrecompiledConfig):
    lm: str = "gemini/gemini-2.5-pro-preview-03-25"
    max_tokens: int = 4096
    temperature: float = 0.7


class PromptToSignatureAgent(PrecompiledAgent):
    config: PromptToSignatureConfig

    def __init__(self, config: PromptToSignatureConfig, **kwargs):
        super().__init__(config, **kwargs)
        self.signature_generator = SignatureGenerator()

        lm = dspy.LM(
            model=config.lm,
            max_tokens=config.max_tokens,
            temperature=config.temperature,
        )

        self.signature_generator.set_lm(lm)

    def forward(self, prompt: str, as_dict: bool = False) -> dspy.Prediction:
        # returns dspy.Prediction object or dict
        return (
            self.signature_generator.generate_signature(prompt)
            if as_dict
            else self.signature_generator(prompt)
        )


agent = PromptToSignatureAgent(PromptToSignatureConfig())


def main():
    agent.push_to_hub("fadeleke/prompt-to-signature", with_code=True)


if __name__ == "__main__":
    main()