diff --git a/src/agent_engine/core/spec.py b/src/agent_engine/core/spec.py index f69f6eb3..a7b501d8 100644 --- a/src/agent_engine/core/spec.py +++ b/src/agent_engine/core/spec.py @@ -14,6 +14,7 @@ class ModelConfig: region: str | None = None max_tokens: int | None = None top_p: float | None = None + cache_system_prompt: bool = True @dataclass(frozen=True) diff --git a/src/agent_engine/engine/langgraph/engine.py b/src/agent_engine/engine/langgraph/engine.py index 12cb7202..1bce094b 100644 --- a/src/agent_engine/engine/langgraph/engine.py +++ b/src/agent_engine/engine/langgraph/engine.py @@ -169,6 +169,7 @@ def _model_factory_kwargs(factory: ModelFactory, model: NodeModelConfig) -> dict "region": model.region, "max_tokens": model.max_tokens, "top_p": model.top_p, + "cache_system_prompt": model.cache_system_prompt, } present: dict[str, object] = { key: value for key, value in optional.items() if value is not None diff --git a/src/agent_engine/models/factory.py b/src/agent_engine/models/factory.py index ac898c43..b5c8366f 100644 --- a/src/agent_engine/models/factory.py +++ b/src/agent_engine/models/factory.py @@ -34,6 +34,7 @@ def build_chat_model( region: str | None = None, max_tokens: int | None = None, top_p: float | None = None, + cache_system_prompt: bool = True, ) -> BaseChatModel: """Construct a chat model from flat config fields. @@ -59,6 +60,7 @@ def build_chat_model( temperature=temperature, max_tokens=max_tokens, top_p=top_p, + cache_system_prompt=cache_system_prompt, ) if normalized_provider == "bedrock": return _build_bedrock_model( @@ -94,6 +96,7 @@ def _build_anthropic_model( temperature: float | None, max_tokens: int | None, top_p: float | None, + cache_system_prompt: bool, ) -> BaseChatModel: kwargs = _without_none( { @@ -103,6 +106,8 @@ def _build_anthropic_model( "top_p": top_p, } ) + if cache_system_prompt: + kwargs["cache_control"] = {"type": "ephemeral"} return init_chat_model(name, **kwargs) diff --git a/src/agent_engine/parsers/yaml/parser.py b/src/agent_engine/parsers/yaml/parser.py index 69f96ea3..f541be6e 100644 --- a/src/agent_engine/parsers/yaml/parser.py +++ b/src/agent_engine/parsers/yaml/parser.py @@ -373,6 +373,9 @@ def _validate_model(path: str, raw: Any, errors: list[ValidationError]) -> None: ): errors.append(ValidationError(f"{path}.top_p", "Must be between 0 and 1")) + cache_system_prompt = raw.get("cache_system_prompt") + if cache_system_prompt is not None and not isinstance(cache_system_prompt, bool): + errors.append(ValidationError(f"{path}.cache_system_prompt", "Must be a boolean")) def _validate_models( defaults: Any, @@ -605,6 +608,7 @@ def _build_model(self, raw: dict[str, Any]) -> ModelConfig: region=raw.get("region"), max_tokens=raw.get("max_tokens"), top_p=raw.get("top_p"), + cache_system_prompt=raw.get("cache_system_prompt", True), ) def _build_resolvers( diff --git a/tests/models/test_factory.py b/tests/models/test_factory.py index 3ced573c..cbe91c21 100644 --- a/tests/models/test_factory.py +++ b/tests/models/test_factory.py @@ -73,6 +73,35 @@ def fake_init_chat_model(name: str, **kwargs: Any) -> _FakeAnthropicModel: result = build_chat_model("anthropic", "claude-haiku-4-5", temperature=0.0) + assert result is model + assert calls == [ + ( + "claude-haiku-4-5", + { + "model_provider": "anthropic", + "temperature": 0.0, + "cache_control": {"type": "ephemeral"}, + }, + ) + ] + + +def test_anthropic_provider_disabled_cache_system_prompt( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[tuple[str, dict[str, Any]]] = [] + model = _FakeAnthropicModel() + + def fake_init_chat_model(name: str, **kwargs: Any) -> _FakeAnthropicModel: + calls.append((name, kwargs)) + return model + + monkeypatch.setattr(factory_mod, "init_chat_model", fake_init_chat_model) + + result = build_chat_model( + "anthropic", "claude-haiku-4-5", temperature=0.0, cache_system_prompt=False + ) + assert result is model assert calls == [ ( diff --git a/tests/parsers/test_cache_system_prompt.py b/tests/parsers/test_cache_system_prompt.py new file mode 100644 index 00000000..2eb1b7a0 --- /dev/null +++ b/tests/parsers/test_cache_system_prompt.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pytest + +from agent_engine.parsers.errors import ParseError +from agent_engine.parsers.yaml.parser import YAMLParser + + +def _spec_with_cache_config(cache_val: str | bool | None) -> str: + if cache_val is not None: + cache_line = f" cache_system_prompt: {str(cache_val).lower()}" + else: + cache_line = "" + return f""" +system: + name: t +orchestrators: + router: + description: "Routes requests" + model: + provider: anthropic + name: claude-3-5-sonnet-20240620 +{cache_line} + prompts: + orchestrator: prompts/r.md +agents: + agent_one: + description: "Agent one" + model: + provider: anthropic + name: claude-3-5-sonnet-20240620 +{cache_line} + prompts: + system: prompts/a.md +graph: + router: + agent_one: +""" + + +def _parse(tmp_path: Path, body: str) -> Any: + cfg = tmp_path / "spec.yml" + cfg.write_text(body, encoding="utf-8") + # Write mock prompt files so validation of files passes + (tmp_path / "prompts").mkdir(exist_ok=True) + (tmp_path / "prompts" / "r.md").write_text("router prompt") + (tmp_path / "prompts" / "a.md").write_text("agent prompt") + return YAMLParser().parse(str(cfg)) + + +def test_cache_system_prompt_enabled_passes(tmp_path: Path) -> None: + spec = _parse(tmp_path, _spec_with_cache_config(True)) + assert spec.graph.node.model.cache_system_prompt is True + + +def test_cache_system_prompt_disabled_passes(tmp_path: Path) -> None: + spec = _parse(tmp_path, _spec_with_cache_config(False)) + assert spec.graph.node.model.cache_system_prompt is False + + +def test_cache_system_prompt_defaults_to_true(tmp_path: Path) -> None: + spec = _parse(tmp_path, _spec_with_cache_config(None)) + assert spec.graph.node.model.cache_system_prompt is True + + +def test_cache_system_prompt_invalid_type_fails(tmp_path: Path) -> None: + body = _spec_with_cache_config("not-a-boolean") + with pytest.raises(ParseError, match="Must be a boolean"): + _parse(tmp_path, body)