diff --git a/CHANGELOG.md b/CHANGELOG.md index 20a1bd5..cfe51c4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- `BACKEND_SLOT_PINNING` (default off): pins each pipeline phase to a + llama-server slot (steward=0, orchestrator=1, synthesizer=2) so a + phase's stable prompt prefix stays in that slot's KV cache. No-op on + Claude; ignored by Ollama. + ## [2.5.0] - 2026-09-11 ### Added diff --git a/src/agents/steward/agent.py b/src/agents/steward/agent.py index 4d405be..31e965c 100644 --- a/src/agents/steward/agent.py +++ b/src/agents/steward/agent.py @@ -162,15 +162,20 @@ class StewardAgent: OpenAI-compatible server (Ollama, llama-server) can sit behind OLLAMA_HOST without this method knowing which. """ + payload: dict = { + "model": self.ollama_model, + "messages": [{"role": "user", "content": prompt}], + "temperature": 0.3, # Lower = more consistent + "top_p": 0.9, + } + if config.BACKEND_SLOT_PINNING: + # The steward owns engine slot 0 (see BACKEND_SLOT_PINNING) + payload["id_slot"] = 0 + async with httpx.AsyncClient(timeout=self.timeout) as client: response = await client.post( f"{self.ollama_host}/v1/chat/completions", - json={ - "model": self.ollama_model, - "messages": [{"role": "user", "content": prompt}], - "temperature": 0.3, # Lower = more consistent - "top_p": 0.9, - }, + json=payload, ) response.raise_for_status() diff --git a/src/agents/tatlock.py b/src/agents/tatlock.py index 821d2fa..4b2ce08 100644 --- a/src/agents/tatlock.py +++ b/src/agents/tatlock.py @@ -536,13 +536,13 @@ class TatlockAgent(AgentInterface): # Run with scoped tools and tracker # Force tool_choice to make LLM actually call tools - from src.anthropic.model_selector import get_tool_choice_settings + from src.anthropic.model_selector import get_tool_choice_settings, with_slot_pinning result = await scoped_agent.run( enriched_message, message_history=pydantic_history if pydantic_history else None, deps=tool_tracker, - model_settings=get_tool_choice_settings(), + model_settings=with_slot_pinning(get_tool_choice_settings(), slot=1), ) logger.info( @@ -725,13 +725,13 @@ class TatlockAgent(AgentInterface): pydantic_history.append(ModelResponse(parts=[TextPart(content=content)])) # Run with scoped tools and tracker - from src.anthropic.model_selector import get_tool_choice_settings + from src.anthropic.model_selector import get_tool_choice_settings, with_slot_pinning result = await scoped_agent.run( enriched_message, message_history=pydantic_history if pydantic_history else None, deps=tool_tracker, - model_settings=get_tool_choice_settings(), + model_settings=with_slot_pinning(get_tool_choice_settings(), slot=1), ) # Extract tool calls and results from the agent's messages @@ -888,9 +888,12 @@ class TatlockAgent(AgentInterface): pydantic_history.append(ModelResponse(parts=[TextPart(content=content)])) # Run synthesis + from src.anthropic.model_selector import with_slot_pinning + result = await synthesis_agent.run( synthesis_prompt, message_history=pydantic_history if pydantic_history else None, + model_settings=with_slot_pinning(None, slot=2), ) logger.info( diff --git a/src/anthropic/model_selector.py b/src/anthropic/model_selector.py index a256e0a..8b1360a 100644 --- a/src/anthropic/model_selector.py +++ b/src/anthropic/model_selector.py @@ -259,6 +259,28 @@ def get_tool_choice_settings() -> ModelSettings: return ModelSettings(extra_body={"tool_choice": "required"}) +def with_slot_pinning(settings: ModelSettings | None, slot: int) -> ModelSettings | None: + """ + Merge llama-server slot pinning into model settings when enabled. + + Each pipeline phase owns one engine slot (steward=0, orchestrator=1, + synthesizer=2), so the phase's stable prompt prefix stays in that + slot's KV cache and a turn re-prefills only its new tokens. Off by + default (BACKEND_SLOT_PINNING); a no-op on the Claude backend, and + Ollama ignores the field, so enabling it is safe on any backend. + """ + if not config.BACKEND_SLOT_PINNING or resolve_backend() == "claude": + return settings + + from pydantic_ai.settings import ModelSettings + + merged = dict(settings or {}) + extra_body = dict(merged.get("extra_body") or {}) + extra_body["id_slot"] = slot + merged["extra_body"] = extra_body + return ModelSettings(**merged) + + def get_sampling_settings(temperature: float) -> ModelSettings: """ Get model_settings with a sampling temperature where the backend allows it. diff --git a/src/core/config.py b/src/core/config.py index 3d33730..5b1979f 100644 --- a/src/core/config.py +++ b/src/core/config.py @@ -147,6 +147,10 @@ class Config(BaseSettings): default=None, description="OpenAI-compatible embeddings host; falls back to OLLAMA_HOST so gen and embed can live on different servers", ) + BACKEND_SLOT_PINNING: bool = Field( + default=False, + description="Pin each pipeline phase to a llama-server slot (steward=0, orchestrator=1, synthesizer=2) so a phase's stable prompt prefix stays in that slot's KV cache; no-op on Claude, ignored by Ollama", + ) # Redis Memory Database REDIS_MEMORY_DB: int = Field(default=1, description="Redis database number for memory cache") diff --git a/tests/anthropic/test_model_selector.py b/tests/anthropic/test_model_selector.py index 0a3086f..9fc598e 100644 --- a/tests/anthropic/test_model_selector.py +++ b/tests/anthropic/test_model_selector.py @@ -91,3 +91,32 @@ class TestGetModelInfo: info = model_selector.get_model_info() assert info["backend"] == "claude" assert info["model"] == config.ANTHROPIC_MODEL + + +class TestWithSlotPinning: + def test_disabled_returns_settings_unchanged(self, local_first, monkeypatch): + monkeypatch.setattr(config, "BACKEND_SLOT_PINNING", False) + base = model_selector.get_tool_choice_settings() + assert model_selector.with_slot_pinning(base, slot=1) is base + assert model_selector.with_slot_pinning(None, slot=2) is None + + def test_enabled_merges_slot_with_existing_extra_body(self, local_first, monkeypatch): + monkeypatch.setattr(config, "BACKEND_SLOT_PINNING", True) + settings = model_selector.with_slot_pinning( + model_selector.get_tool_choice_settings(), slot=1 + ) + extra_body = settings["extra_body"] + assert extra_body["id_slot"] == 1 + # tool_choice from the base settings survives the merge + assert extra_body["tool_choice"] == "required" + + def test_enabled_pins_bare_settings(self, local_first, monkeypatch): + monkeypatch.setattr(config, "BACKEND_SLOT_PINNING", True) + settings = model_selector.with_slot_pinning(None, slot=2) + assert settings["extra_body"] == {"id_slot": 2} + + def test_claude_backend_never_pinned(self, local_first, monkeypatch): + monkeypatch.setattr(config, "BACKEND_SLOT_PINNING", True) + monkeypatch.setattr(config, "PREFER_CLOUD_BACKEND", True) + base = model_selector.get_tool_choice_settings() + assert model_selector.with_slot_pinning(base, slot=1) is base