Fix Gemini streaming and deriver cacheability

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),
             },
         ],
     )

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>
 {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)

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,
                 )
 

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