# Answer a math problem with code.
# Adapted from Dust [maths-generate-code](https://dust.tt/spolu/a/d12ac33169)

from minichain import Backend, JinjaPrompt, Prompt, start_chain, SimplePrompt, show_log


# Prompt that asks LLM for code from math.

class ColorPrompt(Prompt[str, bool]):
    def parse(inp: str) -> str:
        return f"Answer 'Yes' if this is a color, {inp}. Answer:"

    def parse(out: str, inp) -> bool:
        # Encode the parsing logic
        return out.strip() == "Yes"  
ColorPrompt().show({"inp": "dog"}, "No")


with start_chain("color") as backend:
    question = 'What is the sum of the powers of 3 (3^i) that are smaller than 100?'
    prompt = MathPrompt(backend.OpenAI()).chain(SimplePrompt(backend.Python()))
    result = prompt({"question": question})
    print(result)

    
show_log("math.log")