example: Tools using ell-ai (#2525)

* example: Tools using ell-ai

* Add README, rename example, update README

* Update README.md

* update readme

* fies

* fixes

---------

Co-authored-by: Akshay Agrawal <akshay@marimo.io>

examples/ai/README.md

@@ -4,6 +4,7 @@ These examples showcase a few simple applications of AI.
 
 - 💬 [`chat/`](chat/): creating chatbots with marimo, using [`mo.ui.chat`](https://docs.marimo.io/api/inputs/chat.html#marimo.ui.chat)
 - 🛢️ [`data/`](data/): making data labeling and model comparison tools
+- 🛠 [`tools/`](tools/): interacting with external functions and services with function calling, returning rich responses
 - 🍿 [`misc/`](misc/): miscellaneous AI examples
 
 > [!TIP]

examples/ai/tools/README.md

@@ -0,0 +1,24 @@
+# AI tool use 🛠
+
+These are examples of using AI that interact with external functions and
+services.
+
+> [!TIP]
+> Submit a
+> [pull request](https://github.com/marimo-team/marimo/pulls) to add an example!
+
+## Running examples
+
+The requirements of each notebook are serialized in them as a top-level
+comment. Here are the steps to open an example notebook:
+
+1. [Install marimo](https://docs.marimo.io/getting_started/index.html#installation)
+2. [Install `uv`](https://github.com/astral-sh/uv/?tab=readme-ov-file#installation)
+3. Open an example with `marimo edit --sandbox <notebook.py>`.
+
+> [!TIP]
+> The [`--sandbox` flag](https://docs.marimo.io/guides/editor_features/package_management.html) opens the notebook in an isolated virtual environment,
+> automatically installing the notebook's dependencies 📦
+
+You can also open notebooks without `uv`, with just `marimo edit <notebook.py>`;
+however, you'll need to install the requirements yourself.

examples/ai/tools/dataset_analysis.py

@@ -0,0 +1,160 @@
+# /// script
+# requires-python = ">=3.12"
+# dependencies = [
+#     "altair==5.4.1",
+#     "beautifulsoup4==4.12.3",
+#     "ell-ai==0.0.13",
+#     "marimo",
+#     "openai==1.51.0",
+#     "polars==1.9.0",
+#     "pyarrow==17.0.0",
+#     "pydantic==2.9.2",
+#     "requests==2.32.3",
+#     "vega-datasets==0.9.0",
+# ]
+# ///
+
+import marimo
+
+__generated_with = "0.9.1"
+app = marimo.App()
+
+
+@app.cell
+def __():
+    import marimo as mo
+    import ell
+    from pydantic import Field
+    import requests
+    import altair as alt
+    import pyarrow
+    import polars as pl
+    from vega_datasets import data
+    return Field, alt, data, ell, mo, pl, pyarrow, requests
+
+
+@app.cell(hide_code=True)
+def __(mo):
+    mo.md(
+        """
+        # Using tools with ell
+
+        This example shows how to use [`ell`](https://docs.ell.so/) with tools to analyze a dataset and return rich responses like charts and tables.
+        """
+    )
+    return
+
+
+@app.cell
+def __(mo):
+    import os
+
+    os_key = os.environ.get("OPENAI_API_KEY")
+    input_key = mo.ui.text(label="OpenAI API key", kind="password")
+    input_key if not os_key else None
+    return input_key, os, os_key
+
+
+@app.cell
+def __(input_key, mo, os_key):
+    openai_key = os_key or input_key.value
+
+    import openai
+
+    client = openai.Client(api_key=openai_key)
+
+    mo.stop(
+        not openai_key,
+        mo.md(
+            "Please set the `OPENAI_API_KEY` environment variable or provide it in the input field"
+        ),
+    )
+    return client, openai, openai_key
+
+
+@app.cell
+def __(data, pl):
+    cars = pl.DataFrame(data.cars())
+    schema = cars.schema
+    return cars, schema
+
+
+@app.cell
+def __(alt, cars, client, ell, schema):
+    @ell.tool()
+    def get_chart(
+        x_encoding: str,
+        y_encoding: str,
+        color: str,
+    ):
+        """Generate an altair chart. For each encoding, please include the type after the colon. For example,"""
+        return (
+            alt.Chart(cars)
+            .mark_circle()
+            .encode(
+                x=x_encoding,
+                y=y_encoding,
+                color=color,
+            )
+            .properties(width="container")
+        )
+
+
+    @ell.tool()
+    def get_filtered_table(sql_query: str):
+        """Filter a pandas dataframe using SQL. Please only use fields from the schema. When referring to the dataframe, call it 'data'."""
+        print(sql_query)
+        filtered = cars.sql(sql_query, table_name="data")
+        return filtered
+
+
+    @ell.complex(
+        model="gpt-4-turbo", tools=[get_chart, get_filtered_table], client=client
+    )
+    def analyze_dataset(prompt: str) -> str:
+        """You are an agent that can analayze the a dataset"""
+        return f"I have a dataset with schema: {schema}. \n{prompt}"
+    return analyze_dataset, get_chart, get_filtered_table
+
+
+@app.cell
+def __(input_key, mo, schema):
+    text = mo.ui.text(
+        full_width=True,
+        disabled=not input_key.value,
+    ).form(bordered=False)
+
+    mo.md(f"""
+    ## **Ask a question!**
+
+    {mo.accordion({
+        "View schema": schema,
+        "View sample questions": mo.md('''
+        * What is the relationship between Cylinders and Horsepower?"
+        * How many cars with MPG great than 30?
+        ''')
+    })}
+
+    {text}
+    """)
+    return (text,)
+
+
+@app.cell
+def __(analyze_dataset, mo, text):
+    mo.stop(not text.value)
+
+    with mo.status.spinner(title=f"Thinking...", subtitle=text.value):
+        response = analyze_dataset(text.value)
+        summary = "Nothing found"
+        if response.tool_calls:
+            try:
+                summary = response.tool_calls[0]()
+                mo.output.replace(summary)
+            except Exception as e:
+                mo.output.replace(mo.callout(str(e)))
+    return response, summary
+
+
+if __name__ == "__main__":
+    app.run()