From 6a746f1de4edfe1efcd47ef69097b026e1530b56 Mon Sep 17 00:00:00 2001 From: adavyas Date: Fri, 13 Mar 2026 23:20:00 -0700 Subject: [PATCH] Implement deriver custom instructions --- src/deriver/deriver.py | 11 ++- src/deriver/prompts.py | 46 ++++++++++++- src/schemas/configuration.py | 1 + tests/deriver/test_deriver_processing.py | 86 ++++++++++++++++++++++++ tests/deriver/test_prompts.py | 30 +++++++++ tests/utils/test_config_helpers.py | 36 ++++++++++ 6 files changed, 205 insertions(+), 5 deletions(-) create mode 100644 tests/deriver/test_prompts.py create mode 100644 tests/utils/test_config_helpers.py diff --git a/src/deriver/deriver.py b/src/deriver/deriver.py index c236cdff..fc37ff41 100644 --- a/src/deriver/deriver.py +++ b/src/deriver/deriver.py @@ -23,7 +23,7 @@ from src.utils.representation import PromptRepresentation, Representation from src.utils.tokens import track_deriver_input_tokens from .prompts import ( - estimate_minimal_deriver_prompt_tokens, + estimate_deriver_prompt_tokens, minimal_deriver_system_prompt, minimal_deriver_user_prompt, ) @@ -98,7 +98,8 @@ async def process_representation_tasks_batch( ) # Track token usage - count only tokens from messages being processed - prompt_tokens = estimate_minimal_deriver_prompt_tokens() + custom_instructions = message_level_configuration.reasoning.custom_instructions + prompt_tokens = estimate_deriver_prompt_tokens(custom_instructions) queue_item_message_ids_set = set(queue_item_message_ids) messages_tokens = sum( msg.token_count for msg in messages if msg.id in queue_item_message_ids_set @@ -143,7 +144,11 @@ async def process_representation_tasks_batch( {"role": "system", "content": minimal_deriver_system_prompt()}, { "role": "user", - "content": minimal_deriver_user_prompt(observed, formatted_messages), + "content": minimal_deriver_user_prompt( + observed, + formatted_messages, + custom_instructions=custom_instructions, + ), }, ], ) diff --git a/src/deriver/prompts.py b/src/deriver/prompts.py index 92e0eb75..a11071d9 100644 --- a/src/deriver/prompts.py +++ b/src/deriver/prompts.py @@ -11,6 +11,19 @@ from inspect import cleandoc as c from src.utils.tokens import estimate_tokens +def _custom_instructions_section(custom_instructions: str | None) -> str: + """Render the optional custom instructions block for the deriver prompt.""" + if not custom_instructions or not custom_instructions.strip(): + return "" + + return c( + f""" + CUSTOM INSTRUCTIONS: + {custom_instructions.strip()} + """ + ) + + def minimal_deriver_system_prompt() -> str: """Generate the cacheable instructions for observation extraction.""" return c( @@ -38,12 +51,20 @@ EXAMPLES: ) -def minimal_deriver_user_prompt(peer_id: str, messages: str) -> str: +def minimal_deriver_user_prompt( + peer_id: str, + messages: str, + *, + custom_instructions: str | None = None, +) -> str: """Generate the per-request message payload for observation extraction.""" + instructions_section = _custom_instructions_section(custom_instructions) return c( f""" Peer identifier: {peer_id} +{instructions_section} + Messages to analyze: {messages} @@ -55,6 +76,8 @@ Messages to analyze: def minimal_deriver_prompt( peer_id: str, messages: str, + *, + custom_instructions: str | None = None, ) -> str: """ Generate the combined prompt for fast observation extraction. @@ -66,7 +89,7 @@ def minimal_deriver_prompt( f""" {minimal_deriver_system_prompt()} -{minimal_deriver_user_prompt(peer_id, messages)} +{minimal_deriver_user_prompt(peer_id, messages, custom_instructions=custom_instructions)} """ ) @@ -84,3 +107,22 @@ def estimate_minimal_deriver_prompt_tokens() -> int: return estimate_tokens(prompt) except ValueError: return 300 + + +def estimate_deriver_prompt_tokens(custom_instructions: str | None = None) -> int: + """Estimate deriver prompt tokens, including optional custom instructions.""" + if not custom_instructions or not custom_instructions.strip(): + return estimate_minimal_deriver_prompt_tokens() + + return estimate_tokens( + "\n\n".join( + [ + minimal_deriver_system_prompt(), + minimal_deriver_user_prompt( + peer_id="", + messages="", + custom_instructions=custom_instructions, + ), + ] + ) + ) diff --git a/src/schemas/configuration.py b/src/schemas/configuration.py index 29180247..5cdffd52 100644 --- a/src/schemas/configuration.py +++ b/src/schemas/configuration.py @@ -127,6 +127,7 @@ class MessageConfiguration(BaseModel): class ResolvedReasoningConfiguration(BaseModel): enabled: bool + custom_instructions: str | None = None class ResolvedPeerCardConfiguration(BaseModel): diff --git a/tests/deriver/test_deriver_processing.py b/tests/deriver/test_deriver_processing.py index 0cde8a68..fb44989d 100644 --- a/tests/deriver/test_deriver_processing.py +++ b/tests/deriver/test_deriver_processing.py @@ -1,9 +1,22 @@ import signal +from datetime import datetime, timezone +from types import SimpleNamespace from typing import Any +from unittest.mock import AsyncMock import pytest from src import models +from src.deriver.deriver import process_representation_tasks_batch +from src.schemas import ( + ResolvedConfiguration, + ResolvedDreamConfiguration, + ResolvedPeerCardConfiguration, + ResolvedReasoningConfiguration, + ResolvedSummaryConfiguration, +) +from src.utils.clients import HonchoLLMCallResponse +from src.utils.representation import ExplicitObservationBase, PromptRepresentation from src.utils.representation import Representation from src.utils.work_unit import construct_work_unit_key, parse_work_unit_key @@ -183,6 +196,79 @@ class TestBackwardsCompatibility: assert observers == [] + +@pytest.mark.asyncio +class TestCustomInstructions: + async def test_deriver_passes_custom_instructions_into_prompt( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + captured: dict[str, Any] = {} + + def fake_prompt( + peer_id: str, + messages: str, + *, + custom_instructions: str | None = None, + ) -> str: + captured["peer_id"] = peer_id + captured["messages"] = messages + captured["custom_instructions"] = custom_instructions + return "prompt" + + mock_response = HonchoLLMCallResponse( + content=PromptRepresentation( + explicit=[ExplicitObservationBase(content="Alice likes tea")] + ), + output_tokens=5, + finish_reasons=["stop"], + ) + + monkeypatch.setattr("src.deriver.deriver.minimal_deriver_prompt", fake_prompt) + monkeypatch.setattr( + "src.deriver.deriver.honcho_llm_call", + AsyncMock(return_value=mock_response), + ) + monkeypatch.setattr( + "src.crud.representation.RepresentationManager.save_representation", + AsyncMock(), + ) + + message = SimpleNamespace( + id=1, + content="I like tea.", + created_at=datetime.now(timezone.utc), + peer_name="alice", + token_count=4, + session_name="session-1", + workspace_name="workspace-1", + ) + configuration = ResolvedConfiguration( + reasoning=ResolvedReasoningConfiguration( + enabled=True, + custom_instructions="Focus on durable preferences only.", + ), + peer_card=ResolvedPeerCardConfiguration(use=True, create=True), + summary=ResolvedSummaryConfiguration( + enabled=True, + messages_per_short_summary=10, + messages_per_long_summary=20, + ), + dream=ResolvedDreamConfiguration(enabled=True), + ) + + await process_representation_tasks_batch( + messages=[message], + message_level_configuration=configuration, + observers=["alice"], + observed="alice", + queue_item_message_ids=[1], + ) + + assert captured["peer_id"] == "alice" + assert captured["custom_instructions"] == "Focus on durable preferences only." + assert "I like tea." in captured["messages"] + # async def test_representation_batch_uses_earliest_cutoff( # self, # db_session: AsyncSession, diff --git a/tests/deriver/test_prompts.py b/tests/deriver/test_prompts.py new file mode 100644 index 00000000..122d2a39 --- /dev/null +++ b/tests/deriver/test_prompts.py @@ -0,0 +1,30 @@ +from src.deriver.prompts import ( + estimate_deriver_prompt_tokens, + estimate_minimal_deriver_prompt_tokens, + minimal_deriver_prompt, +) + + +class TestMinimalDeriverPrompt: + def test_includes_custom_instructions_section_when_present(self) -> None: + prompt = minimal_deriver_prompt( + peer_id="alice", + messages="alice: hello", + custom_instructions="Focus on durable preferences only.", + ) + + assert "CUSTOM INSTRUCTIONS:" in prompt + assert "Focus on durable preferences only." in prompt + + def test_omits_custom_instructions_section_when_absent(self) -> None: + prompt = minimal_deriver_prompt(peer_id="alice", messages="alice: hello") + + assert "CUSTOM INSTRUCTIONS:" not in prompt + + def test_custom_instructions_increase_prompt_token_estimate(self) -> None: + base_tokens = estimate_minimal_deriver_prompt_tokens() + custom_tokens = estimate_deriver_prompt_tokens( + "Focus on durable preferences only." + ) + + assert custom_tokens > base_tokens diff --git a/tests/utils/test_config_helpers.py b/tests/utils/test_config_helpers.py new file mode 100644 index 00000000..26d6ed13 --- /dev/null +++ b/tests/utils/test_config_helpers.py @@ -0,0 +1,36 @@ +from types import SimpleNamespace + +from src.schemas import MessageConfiguration, ReasoningConfiguration +from src.utils.config_helpers import get_configuration + + +class TestGetConfiguration: + def test_preserves_workspace_custom_instructions(self) -> None: + workspace = SimpleNamespace( + configuration={ + "reasoning": { + "enabled": True, + "custom_instructions": "Focus on durable preferences.", + } + } + ) + + config = get_configuration(None, None, workspace) + + assert config.reasoning.enabled is True + assert config.reasoning.custom_instructions == "Focus on durable preferences." + + def test_message_custom_instructions_override_session_and_workspace(self) -> None: + workspace = SimpleNamespace( + configuration={"reasoning": {"custom_instructions": "workspace scope"}} + ) + session = SimpleNamespace( + configuration={"reasoning": {"custom_instructions": "session scope"}} + ) + message = MessageConfiguration( + reasoning=ReasoningConfiguration(custom_instructions="message scope") + ) + + config = get_configuration(message, session, workspace) + + assert config.reasoning.custom_instructions == "message scope"