diff --git a/src/deriver/deriver.py b/src/deriver/deriver.py index b623b0b5..c236cdff 100644 --- a/src/deriver/deriver.py +++ b/src/deriver/deriver.py @@ -140,10 +140,10 @@ async def process_representation_tasks_batch( retry_attempts=3, trace_name="minimal_deriver", messages=[ - {"role": "system", "content": minimal_deriver_system_prompt(observed)}, + {"role": "system", "content": minimal_deriver_system_prompt()}, { "role": "user", - "content": minimal_deriver_user_prompt(formatted_messages), + "content": minimal_deriver_user_prompt(observed, formatted_messages), }, ], ) diff --git a/src/deriver/prompts.py b/src/deriver/prompts.py index 4d453066..a8318885 100644 --- a/src/deriver/prompts.py +++ b/src/deriver/prompts.py @@ -11,36 +11,38 @@ from inspect import cleandoc as c from src.utils.tokens import estimate_tokens -def minimal_deriver_system_prompt(peer_id: str) -> str: +def minimal_deriver_system_prompt() -> str: """Generate the cacheable instructions for observation extraction.""" return c( - f""" -Analyze messages from {peer_id} to extract **explicit atomic facts** about them. + """ +Analyze messages to extract **explicit atomic facts** about the target person. -[EXPLICIT] DEFINITION: Facts about {peer_id} that can be derived directly from their messages. +[EXPLICIT] DEFINITION: Facts about the target person that can be derived directly from their messages. - Transform statements into one or multiple conclusions - Each conclusion must be self-contained with enough context - - Do not infer unstated background facts or implications - Use absolute dates/times when possible (e.g. "June 26, 2025" not "yesterday") RULES: -- Properly attribute observations to the correct subject: if it is about {peer_id}, say so. If {peer_id} is referencing someone or something else, make that clear. -- Observations should make sense on their own. Each observation will be used in the future to better understand {peer_id}. -- Extract ALL observations from {peer_id} messages, using others as context. +- Use the target person identifier provided in the user message when attributing observations about them. +- Properly attribute observations to the correct subject: if it is about the target person, say so. If the target person is referencing someone or something else, make that clear. +- Observations should make sense on their own. Each observation will be used in the future to better understand the target person. +- Extract ALL observations from the target person's messages, using others as context. - Contextualize each observation sufficiently (e.g. "Ann is nervous about the job interview at the pharmacy" not just "Ann is nervous") EXAMPLES: -- EXPLICIT: "I just had my 25th birthday last Saturday" → "{peer_id} is 25 years old", "{peer_id} celebrated their birthday last Saturday" -- EXPLICIT: "I took my dog for a walk in NYC" → "{peer_id} has a dog", "{peer_id} said they took their dog for a walk in NYC" -- EXPLICIT: "{peer_id} attended college" → "{peer_id} said they attended college" +- EXPLICIT: "I just had my 25th birthday last Saturday" → "The target person is 25 years old", "The target person's birthday is June 21st" +- EXPLICIT: "I took my dog for a walk in NYC" → "The target person has a dog", "The target person lives in NYC" +- EXPLICIT: "The target person attended college" + general knowledge → "The target person completed high school or equivalent" """ ) -def minimal_deriver_user_prompt(messages: str) -> str: +def minimal_deriver_user_prompt(peer_id: str, messages: str) -> str: """Generate the per-request message payload for observation extraction.""" return c( f""" +Target person identifier: {peer_id} + Messages to analyze: {messages} @@ -61,9 +63,9 @@ def minimal_deriver_prompt( """ return c( f""" -{minimal_deriver_system_prompt(peer_id)} +{minimal_deriver_system_prompt()} -{minimal_deriver_user_prompt(messages)} +{minimal_deriver_user_prompt(peer_id, messages)} """ ) @@ -74,8 +76,8 @@ def estimate_minimal_deriver_prompt_tokens() -> int: try: prompt = "\n\n".join( [ - minimal_deriver_system_prompt(peer_id=""), - minimal_deriver_user_prompt(messages=""), + minimal_deriver_system_prompt(), + minimal_deriver_user_prompt(peer_id="", messages=""), ] ) return estimate_tokens(prompt) diff --git a/src/utils/clients.py b/src/utils/clients.py index b2d1054f..c7c56eba 100644 --- a/src/utils/clients.py +++ b/src/utils/clients.py @@ -2601,17 +2601,25 @@ async def handle_streaming_response( ) case genai.Client(): - prompt_text = params["messages"][0]["content"] if params["messages"] else "" stream_config: GenerateContentConfigDict = { "max_output_tokens": cast(int, params["max_tokens"]), } + raw_messages = params.get("messages") + if isinstance(raw_messages, list): + contents, system_instruction = _build_gemini_contents_from_messages( + cast(list[dict[str, Any]], raw_messages) + ) + if system_instruction: + stream_config["system_instruction"] = system_instruction + else: + contents = cast(str, params.get("prompt", "")) if response_model is not None: stream_config["response_mime_type"] = "application/json" stream_config["response_schema"] = response_model response_stream = await client.aio.models.generate_content_stream( model=params["model"], - contents=prompt_text, + contents=contents, config=stream_config, ) else: @@ -2619,7 +2627,7 @@ async def handle_streaming_response( stream_config["response_mime_type"] = "application/json" response_stream = await client.aio.models.generate_content_stream( model=params["model"], - contents=prompt_text, + contents=contents, config=stream_config, ) diff --git a/tests/utils/test_clients.py b/tests/utils/test_clients.py index 3eb57311..cf81e3a6 100644 --- a/tests/utils/test_clients.py +++ b/tests/utils/test_clients.py @@ -1051,6 +1051,66 @@ class TestGoogleClient: call_args = mock_aio.models.generate_content_stream.call_args assert call_args.kwargs["config"]["max_output_tokens"] == 100 + async def test_google_streaming_uses_system_and_user_messages(self): + """Gemini streaming should preserve split system and user messages.""" + from google import genai + + mock_client = Mock(spec=genai.Client) + + mock_finish_reason = Mock() + mock_finish_reason.name = "STOP" + mock_usage_metadata = Mock(candidates_token_count=12) + mock_chunks = [ + Mock(text="ok"), + Mock( + text="", + candidates=[Mock(finish_reason=mock_finish_reason)], + usage_metadata=mock_usage_metadata, + ), + ] + + async def async_chunk_iterator(): + for chunk in mock_chunks: + yield chunk + + mock_aio = Mock() + mock_aio.models.generate_content_stream = AsyncMock( + return_value=async_chunk_iterator() + ) + mock_client.aio = mock_aio + + with patch.dict(CLIENTS, {"google": mock_client}): + chunks: list[HonchoLLMCallStreamChunk] = [] + async for chunk in handle_streaming_response( + client=mock_client, + params={ + "model": "gemini-1.5-pro", + "max_tokens": 100, + "messages": [ + {"role": "system", "content": "stable instructions"}, + {"role": "system", "content": "session history"}, + {"role": "user", "content": "Hello"}, + ], + }, + json_mode=False, + thinking_budget_tokens=None, + ): + chunks.append(chunk) + + assert len(chunks) == 2 + assert chunks[0].content == "ok" + assert chunks[1].is_done is True + + call_args = mock_aio.models.generate_content_stream.call_args + assert call_args is not None + assert call_args.kwargs["contents"] == [ + {"role": "user", "parts": [{"text": "Hello"}]} + ] + assert call_args.kwargs["config"]["system_instruction"] == ( + "stable instructions\n\nsession history" + ) + assert call_args.kwargs["config"]["max_output_tokens"] == 100 + async def test_google_no_candidates_fallback(self): """Test Google/Gemini fallback when no candidates""" from google import genai