Chat

Looking for example notebooks?

For example notebooks, check out examples/ai/chat on our GitHub.

The chat UI element provides an interactive chatbot interface for conversations. It can be customized with different models, including built-in AI models from popular providers or custom functions.

class marimo.ui.chat(model: Callable[[List[ChatMessage], ChatModelConfig], object], *, prompts: List[str] | None = None, on_message: Callable[[List[ChatMessage]], None] | None = None, show_configuration_controls: bool = False, config: ChatModelConfigDict | None = None, allow_attachments: bool | List[str] = False, max_height: int | None = None)

A chatbot UI element for interactive conversations.

Example: Using a custom model.

Define a chatbot by implementing a function that takes a list of ChatMessages and optionally a config object as input, and returns the chat response. The response can be any object, including text, plots, or marimo UI elements.

def my_rag_model(messages, config):
    # Each message has a `content` attribute, as well as a `role`
    # attribute ("user", "system", "assistant");
    question = messages[-1].content
    docs = find_docs(question)
    prompt = template(question, docs, messages)
    response = query(prompt)
    if is_dataset(response):
        return dataset_to_chart(response)
    return response


chat = mo.ui.chat(my_rag_model)

Async functions and async generators are also supported, meaning these are both valid chat functions:

async def my_rag_model(messages):
    return await my_async_function(messages)

async def my_rag_model(messages):
    for response in my_async_iterator(messages):
        yield response

The last value yielded by the async generator is treated as the model response. ui.chat does not yet support streaming responses to the frontend. Please file a GitHub issue if this is important to you: https://github.com/marimo-team/marimo/issues

Example: Using a built-in model.

Instead of defining a chatbot function, you can use a built-in model from the mo.ai.llm module.

chat = mo.ui.chat(
    mo.ai.llm.openai(
        "gpt-4o",
        system_message="You are a helpful assistant.",
    ),
)

You can also allow the user to include attachments in their messages.

chat = mo.ui.chat(
    mo.ai.llm.openai(
        "gpt-4o",
    ),
    allow_attachments=["image/png", "image/jpeg"],
)

Attributes.

  • value: the current chat history, a list of ChatMessage objects.

Initialization Args.

  • model: (Callable[[List[ChatMessage], ChatModelConfig], object]) a callable that takes in the chat history and returns a response

  • prompts: optional list of initial prompts to present to the user

  • on_message: optional callback function to handle new messages

  • show_configuration_controls: whether to show the configuration controls

  • config: optional ChatModelConfigDict to override the default configuration. Keys include:

    • max_tokens

    • temperature

    • top_p

    • top_k

    • frequency_penalty

    • presence_penalty

  • allow_attachments: (bool | List[str]) allow attachments. True for any attachments types, or pass a list of mime types

  • max_height: optional maximum height for the chat element

Public methods

Inherited from UIElement

form([label, bordered, loading, ...])

Create a submittable form out of this UIElement.

send_message(message, buffers)

Send a message to the element rendered on the frontend from the backend.
Inherited from Html

batch(**elements)

Convert an HTML object with templated text into a UI element.

center()

Center an item.

right()

Right-justify.

left()

Left-justify.

callout([kind])

Create a callout containing this HTML element.

style([style])

Wrap an object in a styled container.

Public Data Attributes:

Inherited from UIElement

value

The element’s current value.
Inherited from Html

text

A string of HTML representing this element.

Basic Usage

Here’s a simple example using a custom echo model:

import marimo as mo

def echo_model(messages, config):
    return f"Echo: {messages[-1].content}"

chat = mo.ui.chat(echo_model, prompts=["Hello", "How are you?"])
chat

Here, messages is a list of ChatMessage objects, which has role ("user", "assistant", or "system") and content (the message string) attributes; config is a ChatModelConfig object with various configuration parameters, which you are free to ignore.

Using a Built-in AI Model

You can use marimo’s built-in AI models, such as OpenAI’s GPT:

import marimo as mo

chat = mo.ui.chat(
    mo.ai.llm.openai(
        "gpt-4",
        system_message="You are a helpful assistant.",
    ),
    show_configuration_controls=True
)
chat

Accessing Chat History

You can access the chat history using the value attribute:

chat.value

This returns a list of ChatMessage objects, each containing role, content, and optional attachments attributes.

Custom Model with Additional Context

Here’s an example of a custom model that uses additional context:

import marimo as mo

def rag_model(messages, config):
    question = messages[-1].content
    docs = find_relevant_docs(question)
    context = "\n".join(docs)
    prompt = f"Context: {context}\n\nQuestion: {question}\n\nAnswer:"
    response = query_llm(prompt, config)
    return response

mo.ui.chat(rag_model)

This example demonstrates how you can implement a Retrieval-Augmented Generation (RAG) model within the chat interface.

Templated Prompts

You can pass sample prompts to mo.ui.chat to allow users to select from a list of predefined prompts. By including a {{var}} in the prompt, you can dynamically insert values into the prompt; a form will be generated to allow users to fill in the variables.

mo.ui.chat(
    mo.ai.llm.openai("gpt-4o"),
    prompts=[
        "What is the capital of France?",
        "What is the capital of Germany?",
        "What is the capital of {{country}}?",
    ],
)

Including Attachments

You can allow users to upload attachments to their messages by passing an allow_attachments parameter to mo.ui.chat.

mo.ui.chat(
    rag_model,
    allow_attachments=["image/png", "image/jpeg"],
    # or True for any attachment type
    # allow_attachments=True,
)

Built-in Models

marimo provides several built-in AI models that you can use with the chat UI element.

OpenAI

import marimo as mo

mo.ui.chat(
    mo.ai.llm.openai(
        "gpt-4o",
        system_message="You are a helpful assistant.",
        api_key="sk-proj-...",
    ),
    show_configuration_controls=True
)
class marimo.ai.llm.openai(model: str, *, system_message: str = 'You are a helpful assistant specializing in data science.', api_key: str | None = None, base_url: str | None = None)

OpenAI ChatModel

Args:

  • model: The model to use. Can be found on the OpenAI models page

  • system_message: The system message to use

  • api_key: The API key to use. If not provided, the API key will be retrieved from the OPENAI_API_KEY environment variable or the user’s config.

  • base_url: The base URL to use

Anthropic

import marimo as mo

mo.ui.chat(
    mo.ai.llm.anthropic(
        "claude-3-5-sonnet-20240620",
        system_message="You are a helpful assistant.",
        api_key="sk-ant-...",
    ),
    show_configuration_controls=True
)
class marimo.ai.llm.anthropic(model: str, *, system_message: str = 'You are a helpful assistant specializing in data science.', api_key: str | None = None, base_url: str | None = None)

Anthropic ChatModel

Args:

  • model: The model to use. Can be found on the Anthropic models page

  • system_message: The system message to use

  • api_key: The API key to use. If not provided, the API key will be retrieved from the ANTHROPIC_API_KEY environment variable or the user’s config.

  • base_url: The base URL to use

Google AI

import marimo as mo

mo.ui.chat(
    mo.ai.llm.google(
        "gemini-1.5-pro-latest",
        system_message="You are a helpful assistant.",
        api_key="AI..",
    ),
    show_configuration_controls=True
)
class marimo.ai.llm.google(model: str, *, system_message: str = 'You are a helpful assistant specializing in data science.', api_key: str | None = None)

Google AI ChatModel

Args:

  • model: The model to use. Can be found on the Gemini models page

  • system_message: The system message to use

  • api_key: The API key to use. If not provided, the API key will be retrieved from the GOOGLE_AI_API_KEY environment variable or the user’s config.

Groq

import marimo as mo

mo.ui.chat(
    mo.ai.llm.groq(
        "llama-3.1-70b-versatile",
        system_message="You are a helpful assistant.",
        api_key="gsk-...",
    ),
    show_configuration_controls=True
)
class marimo.ai.llm.groq(model: str, *, system_message: str = 'You are a helpful assistant specializing in data science.', api_key: str | None = None, base_url: str | None = None)

Groq ChatModel

Args:

  • model: The model to use. Can be found on the Groq models page

  • system_message: The system message to use

  • api_key: The API key to use. If not provided, the API key will be retrieved from the GROQ_API_KEY environment variable or the user’s config.

  • base_url: The base URL to use

Types

Chatbots can be implemented with a function that receives a list of ChatMessage objects and a ChatModelConfig.

class marimo.ai.ChatMessage(role: Literal['user', 'assistant', 'system'], content: object, attachments: List[ChatAttachment] | None = None)

A message in a chat.

class marimo.ai.ChatModelConfig(max_tokens: 'Optional[int]' = None, temperature: 'Optional[float]' = None, top_p: 'Optional[float]' = None, top_k: 'Optional[int]' = None, frequency_penalty: 'Optional[float]' = None, presence_penalty: 'Optional[float]' = None)

mo.ui.chat can be instantiated with an initial configuration with a dictionary conforming to the config.

ChatMessages can also include attachments.

class marimo.ai.ChatAttachment(url: 'str', name: 'str' = 'attachment', content_type: 'Optional[str]' = None)