From 2c4d88c316c4bad89b32864617abbb252dfe3b84 Mon Sep 17 00:00:00 2001 From: adavyas Date: Fri, 27 Mar 2026 07:45:56 -0700 Subject: [PATCH] Tune dialectic cache layout for Gemini --- src/dialectic/core.py | 41 ++++++--- src/utils/prompt_cache_layouts.py | 85 +++++++++++++++++++ tests/utils/test_prompt_cache_layouts.py | 103 +++++++++++++++++++++++ 3 files changed, 218 insertions(+), 11 deletions(-) create mode 100644 src/utils/prompt_cache_layouts.py create mode 100644 tests/utils/test_prompt_cache_layouts.py diff --git a/src/dialectic/core.py b/src/dialectic/core.py index 893f71a8..91379ca2 100644 --- a/src/dialectic/core.py +++ b/src/dialectic/core.py @@ -37,6 +37,8 @@ from src.utils.clients import ( honcho_llm_call, ) from src.utils.formatting import format_new_turn_with_timestamp +from src.utils.prompt_cache_layouts import build_system_messages +from src.utils.types import SupportedProviders logger = logging.getLogger(__name__) @@ -85,22 +87,41 @@ class DialecticAgent: self.observed_peer_card: list[str] | None = observed_peer_card self.metric_key: str | None = metric_key self.reasoning_level: ReasoningLevel = reasoning_level + self._provider: SupportedProviders = settings.DIALECTIC.LEVELS[ + self.reasoning_level + ].PROVIDER + self._base_system_prompt: str = prompts.agent_system_prompt( + observer, observed, observer_peer_card, observed_peer_card + ) # Initialize conversation history with system prompt - self.messages: list[dict[str, str]] = [ - { - "role": "system", - "content": prompts.agent_system_prompt( - observer, observed, observer_peer_card, observed_peer_card - ), - } - ] + self.messages: list[dict[str, str]] = self._build_system_messages() self._session_history_initialized: bool = False self._prefetched_conclusion_count: int = 0 self._run_id: str = str(uuid.uuid4())[ :8 ] # Always generate for event correlation + def _build_system_messages( + self, + session_history_section: str | None = None, + ) -> list[dict[str, str]]: + """Build provider-aware system messages for the dialectic prompt prefix.""" + return build_system_messages( + self._provider, + self._base_system_prompt, + session_history_section, + wrapper_tag="rolling_history", + ) + + def _set_system_messages(self, session_history_section: str | None = None) -> None: + """Update system messages without disturbing conversation state.""" + conversation_messages = [ + message for message in self.messages if message.get("role") != "system" + ] + self.messages = self._build_system_messages(session_history_section) + self.messages.extend(conversation_messages) + async def _initialize_session_history(self) -> None: """Fetch and inject session history into the system prompt if configured.""" if self._session_history_initialized: @@ -141,9 +162,7 @@ class DialecticAgent: "" ) - # Keep session history in its own system message so the stable base - # instructions can be cached independently of rolling session context. - self.messages.append({"role": "system", "content": session_history_section}) + self._set_system_messages(session_history_section) async def _prefetch_relevant_observations(self, query: str) -> str | None: """ diff --git a/src/utils/prompt_cache_layouts.py b/src/utils/prompt_cache_layouts.py new file mode 100644 index 00000000..e09fe193 --- /dev/null +++ b/src/utils/prompt_cache_layouts.py @@ -0,0 +1,85 @@ +from typing import Literal + +from src.utils.types import SupportedProviders + +PromptCacheLayoutMode = Literal[ + "auto", + "split", + "split_reverse", + "merged_system", + "history_in_user", + "base_in_user", + "all_user", +] + +_MERGED_SYSTEM_PROVIDERS: frozenset[SupportedProviders] = frozenset({"google"}) + + +def provider_default_prompt_cache_layout( + provider: SupportedProviders, +) -> PromptCacheLayoutMode: + """Return the default prompt-cache layout for a provider.""" + if provider in _MERGED_SYSTEM_PROVIDERS: + return "merged_system" + return "split" + + +def resolve_prompt_cache_layout_mode( + provider: SupportedProviders, + layout_mode: PromptCacheLayoutMode, +) -> PromptCacheLayoutMode: + """Resolve ``auto`` to the provider's default layout.""" + if layout_mode == "auto": + return provider_default_prompt_cache_layout(provider) + return layout_mode + + +def merge_system_prompt_with_rolling_context( + base_prompt: str, + rolling_context: str, + *, + wrapper_tag: str = "rolling_history", +) -> str: + """Combine stable and rolling system context into a single system prompt.""" + stable = base_prompt.strip() + rolling = rolling_context.strip() + + if not rolling: + return stable + + wrapped_context = f"<{wrapper_tag}>\n{rolling}\n" + if not stable: + return wrapped_context + return f"{stable}\n\n{wrapped_context}" + + +def build_system_messages( + provider: SupportedProviders, + base_prompt: str, + rolling_context: str | None = None, + *, + wrapper_tag: str = "rolling_history", +) -> list[dict[str, str]]: + """Build provider-aware system messages for cacheable prompt prefixes.""" + base_content = base_prompt.strip() + rolling_content = rolling_context.strip() if rolling_context else "" + + if not rolling_content: + return [{"role": "system", "content": base_content}] + + if provider_default_prompt_cache_layout(provider) == "merged_system": + return [ + { + "role": "system", + "content": merge_system_prompt_with_rolling_context( + base_content, + rolling_content, + wrapper_tag=wrapper_tag, + ), + } + ] + + return [ + {"role": "system", "content": base_content}, + {"role": "system", "content": rolling_content}, + ] diff --git a/tests/utils/test_prompt_cache_layouts.py b/tests/utils/test_prompt_cache_layouts.py new file mode 100644 index 00000000..c76e4cb6 --- /dev/null +++ b/tests/utils/test_prompt_cache_layouts.py @@ -0,0 +1,103 @@ +from unittest.mock import AsyncMock + +from src.dialectic.core import DialecticAgent +from src.utils.prompt_cache_layouts import ( + build_system_messages, + merge_system_prompt_with_rolling_context, + provider_default_prompt_cache_layout, +) + + +def test_provider_default_prompt_cache_layout_is_google_specific() -> None: + assert provider_default_prompt_cache_layout("google") == "merged_system" + assert provider_default_prompt_cache_layout("anthropic") == "split" + assert provider_default_prompt_cache_layout("openai") == "split" + + +def test_build_system_messages_merges_google_history() -> None: + messages = build_system_messages( + "google", + "stable instructions", + "rolling history", + wrapper_tag="rolling_history", + ) + + assert messages == [ + { + "role": "system", + "content": ( + "stable instructions\n\n" + "\nrolling history\n" + ), + } + ] + + +def test_build_system_messages_keeps_anthropic_history_separate() -> None: + messages = build_system_messages( + "anthropic", + "stable instructions", + "rolling history", + wrapper_tag="rolling_history", + ) + + assert messages == [ + {"role": "system", "content": "stable instructions"}, + {"role": "system", "content": "rolling history"}, + ] + + +def test_merge_system_prompt_with_rolling_context_strips_noise() -> None: + merged = merge_system_prompt_with_rolling_context( + "\n stable instructions \n", + "\n rolling history \n", + wrapper_tag="rolling_history", + ) + + assert ( + merged + == "stable instructions\n\n\nrolling history\n" + ) + + +def test_dialectic_agent_rebuilds_google_system_messages() -> None: + agent = DialecticAgent( + db=AsyncMock(), + workspace_name="workspace", + session_name="session", + observer="Mira", + observed="Jon", + reasoning_level="low", + ) + + agent._provider = "google" + agent._set_system_messages("rolling history") + + assert agent.messages == [ + { + "role": "system", + "content": ( + agent._base_system_prompt.strip() + + "\n\n\nrolling history\n" + ), + } + ] + + +def test_dialectic_agent_rebuilds_anthropic_system_messages() -> None: + agent = DialecticAgent( + db=AsyncMock(), + workspace_name="workspace", + session_name="session", + observer="Mira", + observed="Jon", + reasoning_level="medium", + ) + + agent._provider = "anthropic" + agent._set_system_messages("rolling history") + + assert agent.messages == [ + {"role": "system", "content": agent._base_system_prompt.strip()}, + {"role": "system", "content": "rolling history"}, + ]