Skip to content

Commit 70af8fa

Browse files
committed
feat: use reasoning field in StreamingChunk for Google GenAI
Populate StreamingChunk.reasoning with ReasoningContent instead of storing reasoning deltas as dicts in meta. Update aggregation to read from chunk.reasoning instead of chunk.meta["reasoning_deltas"].
1 parent 24935e4 commit 70af8fa

2 files changed

Lines changed: 47 additions & 22 deletions

File tree

integrations/google_genai/src/haystack_integrations/components/generators/google_genai/chat/utils.py

Lines changed: 15 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -555,7 +555,7 @@ def _convert_google_chunk_to_streaming_chunk(
555555
content = ""
556556
tool_calls: list[ToolCallDelta] = []
557557
finish_reason = None
558-
reasoning_deltas: list[dict[str, str]] = []
558+
reasoning_deltas: list[str] = []
559559
thought_signature_deltas: list[dict[str, Any]] = [] # Track thought signatures in streaming
560560

561561
if chunk.candidates:
@@ -606,39 +606,39 @@ def _convert_google_chunk_to_streaming_chunk(
606606

607607
# Handle thought parts for Gemini 2.5 series
608608
elif hasattr(part, "thought") and part.thought:
609-
thought_delta = {
610-
"type": "reasoning",
611-
"content": part.text if part.text else "",
612-
}
613-
reasoning_deltas.append(thought_delta)
609+
reasoning_deltas.append(part.text if part.text else "")
610+
611+
# Combine reasoning deltas into a single ReasoningContent
612+
reasoning = ReasoningContent(reasoning_text="".join(reasoning_deltas)) if reasoning_deltas else None
614613

615614
# start is only used by print_streaming_chunk. We try to make a reasonable assumption here but it should not be
616615
# a problem if we change it in the future.
617616
start = index == 0 or len(tool_calls) > 0
618617

619-
# Create meta with reasoning deltas and thought signatures if available
618+
# Create meta with thought signatures if available
620619
meta: dict[str, Any] = {
621620
"received_at": datetime.now(timezone.utc).isoformat(),
622621
"model": model,
623622
"usage": usage,
624623
}
625624

626-
# Add reasoning deltas to meta if available
627-
if reasoning_deltas:
628-
meta["reasoning_deltas"] = reasoning_deltas
629-
630625
# Add thought signature deltas to meta if available (for multi-turn context)
631626
if thought_signature_deltas:
632627
meta["thought_signature_deltas"] = thought_signature_deltas
633628

629+
# StreamingChunk allows only one of content/tool_calls/reasoning to be set.
630+
# Determine the effective content: tool_calls and reasoning take priority.
631+
effective_content = "" if tool_calls or reasoning else content
632+
634633
return StreamingChunk(
635-
content="" if tool_calls else content, # prioritize tool calls over content when both are present
634+
content=effective_content,
636635
tool_calls=tool_calls,
637636
component_info=component_info,
638637
index=index,
639638
start=start,
640639
finish_reason=FINISH_REASON_MAPPING.get(finish_reason or ""),
641640
meta=meta,
641+
reasoning=reasoning,
642642
)
643643

644644

@@ -662,13 +662,9 @@ def _aggregate_streaming_chunks_with_reasoning(chunks: list[StreamingChunk]) ->
662662
thoughts_token_count = None
663663

664664
for chunk in chunks:
665-
# Extract reasoning deltas
666-
if chunk.meta and "reasoning_deltas" in chunk.meta:
667-
reasoning_deltas = chunk.meta["reasoning_deltas"]
668-
if isinstance(reasoning_deltas, list):
669-
for delta in reasoning_deltas:
670-
if delta.get("type") == "reasoning":
671-
reasoning_text_parts.append(delta.get("content", ""))
665+
# Extract reasoning from the StreamingChunk.reasoning field
666+
if chunk.reasoning and chunk.reasoning.reasoning_text:
667+
reasoning_text_parts.append(chunk.reasoning.reasoning_text)
672668

673669
# Extract thought signature deltas (for multi-turn context preservation)
674670
if chunk.meta and "thought_signature_deltas" in chunk.meta:

integrations/google_genai/tests/test_chat_generator_utils.py

Lines changed: 32 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -489,6 +489,38 @@ def test_convert_google_chunk_to_streaming_chunk_real_example(self, monkeypatch)
489489
assert streaming_chunk.tool_calls[5].id is None
490490
assert streaming_chunk.tool_calls[5].index == 5
491491

492+
def test_convert_google_chunk_to_streaming_chunk_with_thought(self, monkeypatch):
493+
"""Test that thought parts populate StreamingChunk.reasoning instead of meta."""
494+
monkeypatch.setenv("GOOGLE_API_KEY", "test-api-key")
495+
component = GoogleGenAIChatGenerator()
496+
component_info = ComponentInfo.from_component(component)
497+
498+
# Simulate a chunk with a thought part (reasoning-only chunk)
499+
mock_thought_part = Mock()
500+
mock_thought_part.text = "Let me think about this..."
501+
mock_thought_part.thought = True
502+
mock_thought_part.function_call = None
503+
504+
mock_content = Mock()
505+
mock_content.parts = [mock_thought_part]
506+
mock_candidate = Mock()
507+
mock_candidate.content = mock_content
508+
mock_candidate.finish_reason = None
509+
510+
mock_chunk = Mock()
511+
mock_chunk.candidates = [mock_candidate]
512+
mock_chunk.usage_metadata = None
513+
514+
streaming_chunk = _convert_google_chunk_to_streaming_chunk(
515+
chunk=mock_chunk, index=0, component_info=component_info, model="gemini-2.5-flash"
516+
)
517+
518+
# Reasoning should be in the reasoning field, not in meta
519+
assert streaming_chunk.reasoning is not None
520+
assert streaming_chunk.reasoning.reasoning_text == "Let me think about this..."
521+
assert "reasoning_deltas" not in streaming_chunk.meta
522+
assert streaming_chunk.content == ""
523+
492524
def test_aggregate_streaming_chunks_with_reasoning(self):
493525
"""Test the _aggregate_streaming_chunks_with_reasoning function for reasoning content aggregation."""
494526

@@ -515,9 +547,6 @@ def test_aggregate_streaming_chunks_with_reasoning(self):
515547
}
516548
final_chunk.reasoning = ReasoningContent(reasoning_text="I should greet the user politely")
517549

518-
# Add reasoning deltas to the final chunk meta (this is how the real method works)
519-
final_chunk.meta["reasoning_deltas"] = [{"type": "reasoning", "content": "I should greet the user politely"}]
520-
521550
# Test aggregation
522551
result = _aggregate_streaming_chunks_with_reasoning([chunk1, chunk2, final_chunk])
523552

0 commit comments

Comments
 (0)