From e08e5d5b1019871f10111e2542683623ad1a254e Mon Sep 17 00:00:00 2001 From: Joaquin Coromina Date: Fri, 19 Sep 2025 19:22:09 +0200 Subject: [PATCH 1/2] Add session_id to prompt for fast access of current notebook --- marimo/_server/ai/prompts.py | 10 ++++++++++ marimo/_server/api/endpoints/ai.py | 2 ++ 2 files changed, 12 insertions(+) diff --git a/marimo/_server/ai/prompts.py b/marimo/_server/ai/prompts.py index 6f30c432a69..772a9cab034 100644 --- a/marimo/_server/ai/prompts.py +++ b/marimo/_server/ai/prompts.py @@ -10,6 +10,7 @@ SchemaTable, VariableContext, ) +from marimo._types.ids import SessionId FIM_PREFIX_TAG = "<|fim_prefix|>" FIM_SUFFIX_TAG = "<|fim_suffix|>" @@ -234,15 +235,24 @@ def _get_mode_intro_message(mode: CopilotMode) -> str: ) +def _get_session_info(session_id: SessionId) -> str: + return ( + f"Current notebook session ID: {session_id}. " + "Use this session_id with tools that require it." + ) + + def get_chat_system_prompt( *, custom_rules: Optional[str], context: Optional[AiCompletionContext], include_other_code: str, mode: CopilotMode, + session_id: SessionId, ) -> str: system_prompt: str = f""" {_get_mode_intro_message(mode)} +{_get_session_info(session_id)} Your goal is to do one of the following two things: diff --git a/marimo/_server/api/endpoints/ai.py b/marimo/_server/api/endpoints/ai.py index 75ba6ee0cfe..e8dc56f2f1c 100644 --- a/marimo/_server/api/endpoints/ai.py +++ b/marimo/_server/api/endpoints/ai.py @@ -209,6 +209,7 @@ async def ai_chat( """ app_state = AppState(request) app_state.require_current_session() + session_id = app_state.require_current_session_id() config = app_state.app_config_manager.get_config(hide_secrets=False) body = await parse_request( request, cls=ChatRequest, allow_unknown_keys=True @@ -223,6 +224,7 @@ async def ai_chat( context=body.context, include_other_code=body.include_other_code, mode=ai_config.get("mode", "manual"), + session_id=session_id, ) max_tokens = get_max_tokens(config) From ebf0af235ccbaadf9a3c9bc71db3fd9c2683c9f0 Mon Sep 17 00:00:00 2001 From: Joaquin Coromina Date: Fri, 19 Sep 2025 19:34:34 +0200 Subject: [PATCH 2/2] Fix prompt tests --- tests/_server/ai/snapshots/chat_system_prompts.txt | 7 +++++++ tests/_server/ai/test_prompts.py | 13 ++++++++++++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/_server/ai/snapshots/chat_system_prompts.txt b/tests/_server/ai/snapshots/chat_system_prompts.txt index 6894ab2bc57..76c36fbd401 100644 --- a/tests/_server/ai/snapshots/chat_system_prompts.txt +++ b/tests/_server/ai/snapshots/chat_system_prompts.txt @@ -12,6 +12,7 @@ Your primary function is to help users create, analyze, and improve data science - You do NOT have access to any external tools, plugins, or APIs. - You may not perform any actions beyond generating text and code suggestions. +Current notebook session ID: s_test. Use this session_id with tools that require it. Your goal is to do one of the following two things: @@ -155,6 +156,7 @@ Your primary function is to help users create, analyze, and improve data science - You do NOT have access to any external tools, plugins, or APIs. - You may not perform any actions beyond generating text and code suggestions. +Current notebook session ID: s_test. Use this session_id with tools that require it. Your goal is to do one of the following two things: @@ -301,6 +303,7 @@ Your primary function is to help users create, analyze, and improve data science - You do NOT have access to any external tools, plugins, or APIs. - You may not perform any actions beyond generating text and code suggestions. +Current notebook session ID: s_test. Use this session_id with tools that require it. Your goal is to do one of the following two things: @@ -447,6 +450,7 @@ Your primary function is to help users create, analyze, and improve data science - You do NOT have access to any external tools, plugins, or APIs. - You may not perform any actions beyond generating text and code suggestions. +Current notebook session ID: s_test. Use this session_id with tools that require it. Your goal is to do one of the following two things: @@ -599,6 +603,7 @@ Your primary function is to help users create, analyze, and improve data science - You do NOT have access to any external tools, plugins, or APIs. - You may not perform any actions beyond generating text and code suggestions. +Current notebook session ID: s_test. Use this session_id with tools that require it. Your goal is to do one of the following two things: @@ -759,6 +764,7 @@ Your primary function is to help users create, analyze, and improve data science - You do NOT have access to any external tools, plugins, or APIs. - You may not perform any actions beyond generating text and code suggestions. +Current notebook session ID: s_test. Use this session_id with tools that require it. Your goal is to do one of the following two things: @@ -907,6 +913,7 @@ Your primary function is to help users create, analyze, and improve data science - You do NOT have access to any external tools, plugins, or APIs. - You may not perform any actions beyond generating text and code suggestions. +Current notebook session ID: s_test. Use this session_id with tools that require it. Your goal is to do one of the following two things: diff --git a/tests/_server/ai/test_prompts.py b/tests/_server/ai/test_prompts.py index 2b5eb135f29..0e52a34dcd3 100644 --- a/tests/_server/ai/test_prompts.py +++ b/tests/_server/ai/test_prompts.py @@ -18,6 +18,7 @@ SchemaTable, VariableContext, ) +from marimo._types.ids import SessionId from tests.mocks import snapshotter snapshot = snapshotter(__file__) @@ -216,7 +217,11 @@ def test_chat_system_prompts(): result: str = "" result += _header("no custom rules") result += get_chat_system_prompt( - custom_rules=None, include_other_code="", context=None, mode="manual" + custom_rules=None, + include_other_code="", + context=None, + mode="manual", + session_id=SessionId("s_test"), # stable fake session id for snapshot ) result += _header("with custom rules") @@ -225,6 +230,7 @@ def test_chat_system_prompts(): include_other_code="", context=None, mode="manual", + session_id=SessionId("s_test"), ) result += _header("with variables") @@ -235,6 +241,7 @@ def test_chat_system_prompts(): variables=["var1", "var2"], ), mode="manual", + session_id=SessionId("s_test"), ) result += _header("with VariableContext objects") @@ -256,6 +263,7 @@ def test_chat_system_prompts(): ] ), mode="manual", + session_id=SessionId("s_test"), ) result += _header("with context") @@ -291,6 +299,7 @@ def test_chat_system_prompts(): ], ), mode="manual", + session_id=SessionId("s_test"), ) result += _header("with other code") @@ -299,6 +308,7 @@ def test_chat_system_prompts(): include_other_code="import pandas as pd\nimport numpy as np\n", context=None, mode="manual", + session_id=SessionId("s_test"), ) result += _header("kitchen sink") @@ -324,6 +334,7 @@ def test_chat_system_prompts(): ], ), mode="manual", + session_id=SessionId("s_test"), ) snapshot("chat_system_prompts.txt", result)