From f8f3b9b81fad5f24c3cca07770ba45ca81e44396 Mon Sep 17 00:00:00 2001 From: Kimi Agent Date: Sat, 14 Mar 2026 19:43:11 -0400 Subject: [PATCH] feat: inject session_id into system prompt for session identity awareness Timmy can now introspect which session he's running in (cli, dashboard, loop). - Add {session_id} placeholder to both lite and full system prompts - get_system_prompt() accepts session_id param (default: 'unknown') - create_timmy() accepts session_id param, forwards to prompt - CLI chat/think/status pass their session_id to create_timmy() - session.py passes _DEFAULT_SESSION_ID to create_timmy() - 7 new tests in test_session_identity.py - Updated 2 existing CLI test mocks Closes #64 --- src/timmy/agent.py | 3 +- src/timmy/cli.py | 6 ++-- src/timmy/prompts.py | 11 +++--- src/timmy/session.py | 2 +- tests/timmy/test_cli.py | 4 +-- tests/timmy/test_session_identity.py | 54 ++++++++++++++++++++++++++++ 6 files changed, 69 insertions(+), 11 deletions(-) create mode 100644 tests/timmy/test_session_identity.py diff --git a/src/timmy/agent.py b/src/timmy/agent.py index 36063bc..eeb02ef 100644 --- a/src/timmy/agent.py +++ b/src/timmy/agent.py @@ -203,6 +203,7 @@ def create_timmy( model_size: str | None = None, *, skip_mcp: bool = False, + session_id: str = "unknown", ) -> TimmyAgent: """Instantiate the agent — Ollama or AirLLM, same public interface. @@ -286,7 +287,7 @@ def create_timmy( logger.debug("MCP tools unavailable: %s", exc) # Select prompt tier based on tool capability - base_prompt = get_system_prompt(tools_enabled=use_tools) + base_prompt = get_system_prompt(tools_enabled=use_tools, session_id=session_id) # Try to load memory context try: diff --git a/src/timmy/cli.py b/src/timmy/cli.py index 3f55c20..cbaa6e5 100644 --- a/src/timmy/cli.py +++ b/src/timmy/cli.py @@ -137,7 +137,7 @@ def think( model_size: str | None = _MODEL_SIZE_OPTION, ): """Ask Timmy to think carefully about a topic.""" - timmy = create_timmy(backend=backend, model_size=model_size) + timmy = create_timmy(backend=backend, model_size=model_size, session_id=_CLI_SESSION_ID) timmy.print_response(f"Think carefully about: {topic}", stream=True, session_id=_CLI_SESSION_ID) @@ -181,7 +181,7 @@ def chat( session_id = str(uuid.uuid4()) else: session_id = _CLI_SESSION_ID - timmy = create_timmy(backend=backend, model_size=model_size) + timmy = create_timmy(backend=backend, model_size=model_size, session_id=session_id) # Use agent.run() so we can intercept paused runs for tool confirmation. run_output = timmy.run(message, stream=False, session_id=session_id) @@ -203,7 +203,7 @@ def status( model_size: str | None = _MODEL_SIZE_OPTION, ): """Print Timmy's operational status.""" - timmy = create_timmy(backend=backend, model_size=model_size) + timmy = create_timmy(backend=backend, model_size=model_size, session_id=_CLI_SESSION_ID) timmy.print_response(STATUS_PROMPT, stream=False, session_id=_CLI_SESSION_ID) diff --git a/src/timmy/prompts.py b/src/timmy/prompts.py index 9b2fae4..6efc6a7 100644 --- a/src/timmy/prompts.py +++ b/src/timmy/prompts.py @@ -31,6 +31,7 @@ Rules: "feel free to ask." - When your values conflict (e.g. honesty vs. helpfulness), lead with honesty. - Sometimes the right answer is nothing. Do not fill silence with noise. +- You are running in session "{session_id}". """ # --------------------------------------------------------------------------- @@ -79,28 +80,30 @@ IDENTITY: - If a request is ambiguous, ask one brief clarifying question. - When you state a fact, commit to it. - Never show raw tool call JSON or function syntax in responses. +- You are running in session "{session_id}". Session types: "cli" = terminal user, "dashboard" = web UI, "loop" = dev loop automation, other = custom context. """ # Default to lite for safety SYSTEM_PROMPT = SYSTEM_PROMPT_LITE -def get_system_prompt(tools_enabled: bool = False) -> str: +def get_system_prompt(tools_enabled: bool = False, session_id: str = "unknown") -> str: """Return the appropriate system prompt based on tool capability. Args: tools_enabled: True if the model supports reliable tool calling. + session_id: The session identifier (cli, dashboard, loop, etc.) Returns: - The system prompt string with model name injected from config. + The system prompt string with model name and session_id injected. """ from config import settings model_name = settings.ollama_model if tools_enabled: - return SYSTEM_PROMPT_FULL.format(model_name=model_name) - return SYSTEM_PROMPT_LITE.format(model_name=model_name) + return SYSTEM_PROMPT_FULL.format(model_name=model_name, session_id=session_id) + return SYSTEM_PROMPT_LITE.format(model_name=model_name, session_id=session_id) STATUS_PROMPT = """Give a one-sentence status report confirming diff --git a/src/timmy/session.py b/src/timmy/session.py index 13b80f3..cfd1dce 100644 --- a/src/timmy/session.py +++ b/src/timmy/session.py @@ -53,7 +53,7 @@ def _get_agent(): from timmy.agent import create_timmy try: - _agent = create_timmy() + _agent = create_timmy(session_id=_DEFAULT_SESSION_ID) logger.info("Session: Timmy agent initialized (singleton)") except Exception as exc: logger.error("Session: Failed to create Timmy agent: %s", exc) diff --git a/tests/timmy/test_cli.py b/tests/timmy/test_cli.py index 83e7018..e6e1679 100644 --- a/tests/timmy/test_cli.py +++ b/tests/timmy/test_cli.py @@ -62,7 +62,7 @@ def test_think_passes_model_size_option(): with patch("timmy.cli.create_timmy", return_value=mock_timmy) as mock_create: runner.invoke(app, ["think", "topic", "--model-size", "70b"]) - mock_create.assert_called_once_with(backend=None, model_size="70b") + mock_create.assert_called_once_with(backend=None, model_size="70b", session_id="cli") # --------------------------------------------------------------------------- @@ -119,7 +119,7 @@ def test_chat_passes_backend_option(): with patch("timmy.cli.create_timmy", return_value=mock_timmy) as mock_create: runner.invoke(app, ["chat", "test", "--backend", "airllm"]) - mock_create.assert_called_once_with(backend="airllm", model_size=None) + mock_create.assert_called_once_with(backend="airllm", model_size=None, session_id="cli") def test_chat_cleans_response(): diff --git a/tests/timmy/test_session_identity.py b/tests/timmy/test_session_identity.py new file mode 100644 index 0000000..c22a30b --- /dev/null +++ b/tests/timmy/test_session_identity.py @@ -0,0 +1,54 @@ +"""Tests for session identity awareness in system prompts. + +Issue #64: Timmy should know what session it is running in. +""" + +from timmy.prompts import get_system_prompt + + +class TestSessionIdentity: + """Test that session_id is properly injected into system prompts.""" + + def test_lite_prompt_includes_session_id(self): + """Lite prompt should include the session_id in the output.""" + prompt = get_system_prompt(tools_enabled=False, session_id="cli") + assert "cli" in prompt + assert 'session "cli"' in prompt + + def test_full_prompt_includes_session_id(self): + """Full prompt should include the session_id in the output.""" + prompt = get_system_prompt(tools_enabled=True, session_id="dashboard") + assert "dashboard" in prompt + assert 'session "dashboard"' in prompt + + def test_default_session_id_is_unknown(self): + """When no session_id is provided, default should be 'unknown'.""" + prompt = get_system_prompt() + assert "unknown" in prompt + assert 'session "unknown"' in prompt + + def test_lite_prompt_session_format(self): + """Lite prompt should have session info after the rules.""" + prompt = get_system_prompt(tools_enabled=False, session_id="loop") + # Should contain the session line + assert '- You are running in session "loop".' in prompt + + def test_full_prompt_session_format(self): + """Full prompt should have session info in IDENTITY section.""" + prompt = get_system_prompt(tools_enabled=True, session_id="custom") + # Should contain session type explanation + assert 'session "custom"' in prompt + assert "Session types:" in prompt + + def test_various_session_ids(self): + """Test that different session IDs are properly handled.""" + test_cases = ["cli", "dashboard", "loop", "custom", "test-session-123"] + for sid in test_cases: + prompt = get_system_prompt(session_id=sid) + assert f'"{sid}"' in prompt, f"Session ID '{sid}' not found in prompt" + + def test_session_id_with_special_chars(self): + """Test session IDs with special characters.""" + sid = "user_123-test.session" + prompt = get_system_prompt(session_id=sid) + assert sid in prompt