# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "ell-ai==0.0.13",
#     "marimo",
#     "openai==1.51.0",
# ]
# ///

import marimo

__generated_with = "0.9.20"
app = marimo.App(width="medium")


@app.cell(hide_code=True)
def __():
    import os
    import textwrap

    import ell
    import marimo as mo
    import openai

    return ell, mo, openai, os, textwrap


@app.cell(hide_code=True)
def __(mo):
    mo.md(
        """
        # Creating a code interpreter

        This example shows how to create a code-interpreter in a few lines of code.
        """
    )
    return


@app.cell(hide_code=True)
def __(mo, os):
    api_key = mo.ui.text(
        label="OpenAI API key",
        kind="password",
        value=os.environ.get("OPENAI_API_KEY", ""),
    )
    api_key
    return (api_key,)


@app.cell(hide_code=True)
def __(api_key, openai):
    client = openai.Client(api_key=api_key.value)
    model = "gpt-4-turbo"
    return client, model


@app.cell(hide_code=True)
def __():
    # https://stackoverflow.com/questions/33908794/get-value-of-last-expression-in-exec-call
    def exec_with_result(script, globals=None, locals=None):
        """Execute a script and return the value of the last expression"""
        import ast

        stmts = list(ast.iter_child_nodes(ast.parse(script)))
        if not stmts:
            return None
        if isinstance(stmts[-1], ast.Expr):
            # the last one is an expression and we will try to return the results
            # so we first execute the previous statements
            if len(stmts) > 1:
                exec(
                    compile(ast.Module(body=stmts[:-1]), filename="<ast>", mode="exec"),
                    globals,
                    locals,
                )
            # then we eval the last one
            return eval(
                compile(
                    ast.Expression(body=stmts[-1].value),
                    filename="<ast>",
                    mode="eval",
                ),
                globals,
                locals,
            )
        else:
            # otherwise we just execute the entire code
            return exec(script, globals, locals)

    return (exec_with_result,)


@app.cell
def __(ell, exec_with_result, mo):
    def code_fence(code):
        return f"```python\n\n{code}\n\n```"

    @ell.tool()
    def execute_code(code: str):
        """
        Execute python. The last line should be the result, don't use print().
        Please make sure it is safe before executing.
        """
        with mo.capture_stdout() as out:
            result = exec_with_result(code)
            output = out.getvalue()
            results = [
                "**Code**",
                code_fence(code),
                "**Result**",
                code_fence(result if result is not None else output),
            ]
            return mo.md("\n\n".join(results))

    return code_fence, execute_code


@app.cell
def __(client, ell, execute_code, mo, model):
    @ell.complex(model=model, tools=[execute_code], client=client)
    def custom_chatbot(messages, config) -> str:
        """You are data scientist with access to writing python code."""
        return [
            ell.user(message.content)
            if message.role == "user"
            else ell.assistant(message.content)
            for message in messages
        ]

    def my_model(messages, config):
        response = custom_chatbot(messages, config)
        if response.tool_calls:
            return response.tool_calls[0]()
        return mo.md(response.text)

    return custom_chatbot, my_model


@app.cell
def __(mo, my_model):
    numbers = [x for x in range(1, 10)]

    mo.ui.chat(
        my_model,
        prompts=[
            "What is the square root of {{number}}?",
            f"Can you sum this list using python: {numbers}",
        ],
    )
    return (numbers,)


if __name__ == "__main__":
    app.run()