Skip to content

⚡️ Speed up function _get_mode_intro_message by 31% #395

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

⚡️ Speed up function _get_mode_intro_message by 31% #395

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 25 additions & 23 deletions marimo/_server/ai/prompts.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,29 @@
)
from marimo._types.ids import SessionId

# Pre-constructed intro messages for each mode, avoiding string concatenation on each call
_MANUAL_INTRO: str = (
"You are Marimo Copilot, an AI assistant integrated into the marimo notebook code editor.\n"
"Your primary function is to help users create, analyze, and improve data science notebooks using marimo's reactive programming model.\n"
"## Capabilities\n"
"- Answer questions and provide guidance using only your internal knowledge and the notebook context provided by the user.\n"
"\n"
"## Limitations\n"
"- You do NOT have access to any external tools, plugins, or APIs.\n"
"- You may not perform any actions beyond generating text and code suggestions.\n"
)
_ASK_INTRO: str = (
"You are Marimo Copilot, an AI assistant integrated into the marimo notebook code editor.\n"
"Your primary function is to help users create, analyze, and improve data science notebooks using marimo's reactive programming model.\n"
"## Capabilities\n"
"- You can use a set of read-only tools to gather additional context from the notebook or environment (e.g., searching code, summarizing data, or reading documentation).\n"
"- You may use these tools ONLY to gather information, not to modify code or state.\n"
"\n"
"## Limitations\n"
"- All tool use is strictly read-only. You may not perform write, edit, or execution actions.\n"
"- You must always explain to the user why you are using a tool before invoking it.\n"
)

FIM_PREFIX_TAG = "<|fim_prefix|>"
FIM_SUFFIX_TAG = "<|fim_suffix|>"
FIM_MIDDLE_TAG = "<|fim_middle|>"
Expand Down Expand Up @@ -219,31 +242,10 @@ def get_inline_system_prompt(*, language: Language) -> str:


def _get_mode_intro_message(mode: CopilotMode) -> str:
base_intro = (
"You are Marimo Copilot, an AI assistant integrated into the marimo notebook code editor.\n"
"Your primary function is to help users create, analyze, and improve data science notebooks using marimo's reactive programming model.\n"
)
if mode == "manual":
return (
f"{base_intro}"
"## Capabilities\n"
"- Answer questions and provide guidance using only your internal knowledge and the notebook context provided by the user.\n"
"\n"
"## Limitations\n"
"- You do NOT have access to any external tools, plugins, or APIs.\n"
"- You may not perform any actions beyond generating text and code suggestions.\n"
)
return _MANUAL_INTRO
elif mode == "ask":
return (
f"{base_intro}"
"## Capabilities\n"
"- You can use a set of read-only tools to gather additional context from the notebook or environment (e.g., searching code, summarizing data, or reading documentation).\n"
"- You may use these tools ONLY to gather information, not to modify code or state.\n"
"\n"
"## Limitations\n"
"- All tool use is strictly read-only. You may not perform write, edit, or execution actions.\n"
"- You must always explain to the user why you are using a tool before invoking it.\n"
)
return _ASK_INTRO


def _get_session_info(session_id: SessionId) -> str:
Expand Down