diff --git a/src/talemate/agents/world_state/__init__.py b/src/talemate/agents/world_state/__init__.py index 19d55c39..9e2acce0 100644 --- a/src/talemate/agents/world_state/__init__.py +++ b/src/talemate/agents/world_state/__init__.py @@ -133,6 +133,15 @@ class WorldStateAgent(MemoryRAGMixin, CharacterProgressionMixin, AvatarMixin, Ag description="When on, Look at and Investigate messages also count as part of the current moment and can show inline highlights.", value=False, ), + "examine_length": AgentActionConfig( + type="number", + label="Investigate length", + description="Token length for the Investigate result when taking a closer look at a highlighted entity. Shorter lengths also tell the model to keep the description briefer.", + value=256, + min=32, + max=1024, + step=32, + ), "inject_as_scene_memory": AgentActionConfig( type="bool", label="Pin to context", @@ -236,6 +245,10 @@ class WorldStateAgent(MemoryRAGMixin, CharacterProgressionMixin, AvatarMixin, Ag "update_world_state", "include_context_investigation" ) + @property + def update_world_state_examine_length(self) -> int: + return self.resolve_config("update_world_state", "examine_length") + @property def update_world_state_durable_snapshot(self) -> bool: return self.resolve_config("update_world_state", "durable_snapshot") @@ -498,7 +511,8 @@ class WorldStateAgent(MemoryRAGMixin, CharacterProgressionMixin, AvatarMixin, Ag entity. Takes the 2-3 sentence snapshot text and expands it into grounded prose describing what the player observes about the entity at this moment in the scene. Output length is governed by the - `create` kind's token budget via response-length.jinja2. + `examine_length` config, which both caps the tokens and scales the + auto-appended response-length instruction. Returns the synthesized text. Caller is responsible for surfacing it in the UI — typically as a ContextInvestigationMessage anchored to @@ -509,10 +523,12 @@ class WorldStateAgent(MemoryRAGMixin, CharacterProgressionMixin, AvatarMixin, Ag if not snapshot_text or not snapshot_text.strip(): raise ValueError("examine_entity requires non-empty snapshot_text") + response_length = self.update_world_state_examine_length + _, extracted = await Prompt.request( "world_state.examine-entity", self.client, - "create", + f"create_{response_length}", vars={ "scene": self.scene, "max_tokens": self.client.max_token_length, diff --git a/src/talemate/client/presets.py b/src/talemate/client/presets.py index ec2b612a..3b4db2a9 100644 --- a/src/talemate/client/presets.py +++ b/src/talemate/client/presets.py @@ -13,6 +13,7 @@ __all__ = [ "set_max_tokens", "set_preset", "preset_for_kind", + "preset_name_for_kind", "make_kind", "max_tokens_for_kind", ] @@ -84,6 +85,11 @@ PRESET_SUBSTRING_MAPPINGS = { "visualize": "creative_instruction", "visual": "creative_instruction", "world_state": "analytical", + # Must stay after "create"/"creative": preset_for_kind takes the LAST + # matching substring, so this lets a parametric kind like + # "creative_instruction_256" resolve to the creative_instruction preset + # instead of being captured by the broader "create"/"creative" entries. + "creative_instruction": "creative_instruction", } PRESET_MAPPING = { @@ -105,11 +111,14 @@ PRESET_MAPPING = { } -def preset_for_kind(kind: str, client: "ClientBase") -> dict: - # Check the substrings first(based on order of the original elifs) - - preset_name = None +def preset_name_for_kind(kind: str) -> str | None: + """ + Resolve a kind string to an inference preset name. + Exact matches in PRESET_MAPPING win first; otherwise the LAST matching + substring in PRESET_SUBSTRING_MAPPINGS is used (so more specific entries + placed later override broader ones). Returns None when nothing matches. + """ preset_name = PRESET_MAPPING.get(kind) if not preset_name: @@ -117,6 +126,12 @@ def preset_for_kind(kind: str, client: "ClientBase") -> dict: if substring in kind: preset_name = value + return preset_name + + +def preset_for_kind(kind: str, client: "ClientBase") -> dict: + preset_name = preset_name_for_kind(kind) + if not preset_name: log.warning( f"No preset found for kind {kind}, defaulting to 'scene_direction'", diff --git a/tests/prompts/test_world_state_templates.py b/tests/prompts/test_world_state_templates.py index 53a5d8a1..cec1c76f 100644 --- a/tests/prompts/test_world_state_templates.py +++ b/tests/prompts/test_world_state_templates.py @@ -733,6 +733,20 @@ class TestWorldStateAgentExamineMethods: entity_name="X", entity_kind="item", snapshot_text=" " ) + @pytest.mark.asyncio + async def test_examine_entity_uses_configured_length_in_kind(self, active_context): + agent = active_context + agent.actions["update_world_state"].config["examine_length"].value = 128 + agent.client.send_prompt = AsyncMock(return_value="A short look.") + + await agent.examine_entity( + entity_name="The Silver Dagger", + entity_kind="item", + snapshot_text="A worn silver dagger with an etched pommel sigil.", + ) + + assert agent.client.send_prompt.call_args.kwargs["kind"] == "create_128" + class TestWorldStateAgentReinforcementMethods: """Tests for world_state agent reinforcement methods.""" diff --git a/tests/test_client_presets.py b/tests/test_client_presets.py new file mode 100644 index 00000000..61f62c07 --- /dev/null +++ b/tests/test_client_presets.py @@ -0,0 +1,48 @@ +"""Tests for kind -> preset/token resolution in talemate.client.presets.""" + +import pytest + +from talemate.client.presets import ( + max_tokens_for_kind, + preset_name_for_kind, +) + + +class TestPresetNameForKind: + def test_exact_match_wins(self): + assert preset_name_for_kind("create") == "creative_instruction" + assert preset_name_for_kind("conversation") == "conversation" + + def test_substring_fallback(self): + assert preset_name_for_kind("narrate_512") == "creative" + assert preset_name_for_kind("create_256") == "creative" + + def test_unknown_kind_returns_none(self): + assert preset_name_for_kind("totally_unknown") is None + + @pytest.mark.parametrize("length", [128, 256, 512]) + def test_creative_instruction_parametric_kind_keeps_preset(self, length): + # A parametric creative_instruction kind must resolve to the + # creative_instruction preset, not be captured by the broader + # "create"/"creative" substring entries. + assert ( + preset_name_for_kind(f"creative_instruction_{length}") + == "creative_instruction" + ) + + def test_create_parametric_kind_unaffected(self): + # The narrower creative_instruction entry must not change how plain + # create_/creative_ parametric kinds resolve. + assert preset_name_for_kind("create_256") == "creative" + assert preset_name_for_kind("creative_256") == "creative" + + +class TestMaxTokensForKind: + def test_trailing_digit_is_used(self): + assert max_tokens_for_kind("creative_instruction_256", 8192) == 256 + assert max_tokens_for_kind("create_128", 8192) == 128 + + def test_callable_budget_mapping(self): + # "create" caps at min(1024, budget * 0.35) + assert max_tokens_for_kind("create", 1000) == 350 + assert max_tokens_for_kind("create", 100000) == 1024