chore: import upstream snapshot with attribution

tests/unittests/models/test_llm_response.py

@@ -0,0 +1,459 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for LlmResponse, including log probabilities feature."""
+
+from google.adk.models.llm_response import LlmResponse
+from google.genai import types
+
+
+def test_llm_response_create_with_logprobs():
+  """Test LlmResponse.create() extracts logprobs from candidate."""
+  avg_logprobs = -0.75
+  logprobs_result = types.LogprobsResult(
+      chosen_candidates=[], top_candidates=[]
+  )
+
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(parts=[types.Part(text='Response text')]),
+              finish_reason=types.FinishReason.STOP,
+              avg_logprobs=avg_logprobs,
+              logprobs_result=logprobs_result,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.avg_logprobs == avg_logprobs
+  assert response.logprobs_result == logprobs_result
+  assert response.content.parts[0].text == 'Response text'
+  assert response.finish_reason == types.FinishReason.STOP
+
+
+def test_llm_response_create_without_logprobs():
+  """Test LlmResponse.create() handles missing logprobs gracefully."""
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(parts=[types.Part(text='Response text')]),
+              finish_reason=types.FinishReason.STOP,
+              avg_logprobs=None,
+              logprobs_result=None,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.avg_logprobs is None
+  assert response.logprobs_result is None
+  assert response.content.parts[0].text == 'Response text'
+
+
+def test_llm_response_create_error_case_with_logprobs():
+  """Test LlmResponse.create() includes logprobs in error cases."""
+  avg_logprobs = -2.1
+
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=None,  # No content - error case
+              finish_reason=types.FinishReason.SAFETY,
+              finish_message='Safety filter triggered',
+              avg_logprobs=avg_logprobs,
+              logprobs_result=None,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.avg_logprobs == avg_logprobs
+  assert response.logprobs_result is None
+  assert response.error_code == types.FinishReason.SAFETY
+  assert response.error_message == 'Safety filter triggered'
+
+
+def test_llm_response_create_no_candidates():
+  """Test LlmResponse.create() with no candidates."""
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[],
+      prompt_feedback=types.GenerateContentResponsePromptFeedback(
+          block_reason=types.BlockedReason.SAFETY,
+          block_reason_message='Prompt blocked for safety',
+      ),
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  # No candidates means no logprobs
+  assert response.avg_logprobs is None
+  assert response.logprobs_result is None
+  assert response.error_code == types.BlockedReason.SAFETY
+  assert response.error_message == 'Prompt blocked for safety'
+
+
+def test_llm_response_create_no_candidates_without_prompt_feedback():
+  """Test LlmResponse.create() for empty successful model responses."""
+  usage_metadata = types.GenerateContentResponseUsageMetadata(
+      prompt_token_count=10,
+      candidates_token_count=0,
+      total_token_count=10,
+  )
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[],
+      usage_metadata=usage_metadata,
+      model_version='gemini-2.5-flash',
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.error_code is None
+  assert response.error_message is None
+  assert response.finish_reason is None
+  assert response.content is not None
+  assert response.content.role == 'model'
+  assert not response.content.parts
+  assert response.usage_metadata == usage_metadata
+  assert response.model_version == 'gemini-2.5-flash'
+
+
+def test_llm_response_create_with_concrete_logprobs_result():
+  """Test LlmResponse.create() with detailed logprobs_result containing actual token data."""
+  # Create realistic logprobs data
+  chosen_candidates = [
+      types.LogprobsResultCandidate(
+          token='The', log_probability=-0.1, token_id=123
+      ),
+      types.LogprobsResultCandidate(
+          token=' capital', log_probability=-0.5, token_id=456
+      ),
+      types.LogprobsResultCandidate(
+          token=' of', log_probability=-0.2, token_id=789
+      ),
+  ]
+
+  top_candidates = [
+      types.LogprobsResultTopCandidates(
+          candidates=[
+              types.LogprobsResultCandidate(
+                  token='The', log_probability=-0.1, token_id=123
+              ),
+              types.LogprobsResultCandidate(
+                  token='A', log_probability=-2.3, token_id=124
+              ),
+              types.LogprobsResultCandidate(
+                  token='This', log_probability=-3.1, token_id=125
+              ),
+          ]
+      ),
+      types.LogprobsResultTopCandidates(
+          candidates=[
+              types.LogprobsResultCandidate(
+                  token=' capital', log_probability=-0.5, token_id=456
+              ),
+              types.LogprobsResultCandidate(
+                  token=' city', log_probability=-1.2, token_id=457
+              ),
+              types.LogprobsResultCandidate(
+                  token=' main', log_probability=-2.8, token_id=458
+              ),
+          ]
+      ),
+  ]
+
+  avg_logprobs = -0.27  # Average of -0.1, -0.5, -0.2
+  logprobs_result = types.LogprobsResult(
+      chosen_candidates=chosen_candidates, top_candidates=top_candidates
+  )
+
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  parts=[types.Part(text='The capital of France is Paris.')]
+              ),
+              finish_reason=types.FinishReason.STOP,
+              avg_logprobs=avg_logprobs,
+              logprobs_result=logprobs_result,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.avg_logprobs == avg_logprobs
+  assert response.logprobs_result is not None
+
+  # Test chosen candidates
+  assert len(response.logprobs_result.chosen_candidates) == 3
+  assert response.logprobs_result.chosen_candidates[0].token == 'The'
+  assert response.logprobs_result.chosen_candidates[0].log_probability == -0.1
+  assert response.logprobs_result.chosen_candidates[0].token_id == 123
+  assert response.logprobs_result.chosen_candidates[1].token == ' capital'
+  assert response.logprobs_result.chosen_candidates[1].log_probability == -0.5
+  assert response.logprobs_result.chosen_candidates[1].token_id == 456
+
+  # Test top candidates
+  assert len(response.logprobs_result.top_candidates) == 2
+  assert (
+      len(response.logprobs_result.top_candidates[0].candidates) == 3
+  )  # 3 alternatives for first token
+  assert response.logprobs_result.top_candidates[0].candidates[0].token == 'The'
+  assert (
+      response.logprobs_result.top_candidates[0].candidates[0].token_id == 123
+  )
+  assert response.logprobs_result.top_candidates[0].candidates[1].token == 'A'
+  assert (
+      response.logprobs_result.top_candidates[0].candidates[1].token_id == 124
+  )
+  assert (
+      response.logprobs_result.top_candidates[0].candidates[2].token == 'This'
+  )
+  assert (
+      response.logprobs_result.top_candidates[0].candidates[2].token_id == 125
+  )
+
+
+def test_llm_response_create_with_partial_logprobs_result():
+  """Test LlmResponse.create() with logprobs_result having only chosen_candidates."""
+  chosen_candidates = [
+      types.LogprobsResultCandidate(
+          token='Hello', log_probability=-0.05, token_id=111
+      ),
+      types.LogprobsResultCandidate(
+          token=' world', log_probability=-0.8, token_id=222
+      ),
+  ]
+
+  logprobs_result = types.LogprobsResult(
+      chosen_candidates=chosen_candidates,
+      top_candidates=[],  # Empty top candidates
+  )
+
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(parts=[types.Part(text='Hello world')]),
+              finish_reason=types.FinishReason.STOP,
+              avg_logprobs=-0.425,  # Average of -0.05 and -0.8
+              logprobs_result=logprobs_result,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.avg_logprobs == -0.425
+  assert response.logprobs_result is not None
+  assert len(response.logprobs_result.chosen_candidates) == 2
+  assert len(response.logprobs_result.top_candidates) == 0
+  assert response.logprobs_result.chosen_candidates[0].token == 'Hello'
+  assert response.logprobs_result.chosen_candidates[1].token == ' world'
+
+
+def test_llm_response_create_with_citation_metadata():
+  """Test LlmResponse.create() extracts citation_metadata from candidate."""
+  citation_metadata = types.CitationMetadata(
+      citations=[
+          types.Citation(
+              start_index=0,
+              end_index=10,
+              uri='https://example.com',
+          )
+      ]
+  )
+
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(parts=[types.Part(text='Response text')]),
+              finish_reason=types.FinishReason.STOP,
+              citation_metadata=citation_metadata,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.citation_metadata == citation_metadata
+  assert response.content.parts[0].text == 'Response text'
+
+
+def test_llm_response_create_without_citation_metadata():
+  """Test LlmResponse.create() handles missing citation_metadata gracefully."""
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(parts=[types.Part(text='Response text')]),
+              finish_reason=types.FinishReason.STOP,
+              citation_metadata=None,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.citation_metadata is None
+  assert response.content.parts[0].text == 'Response text'
+
+
+def test_llm_response_create_error_case_with_citation_metadata():
+  """Test LlmResponse.create() includes citation_metadata in error cases."""
+  citation_metadata = types.CitationMetadata(
+      citations=[
+          types.Citation(
+              start_index=0,
+              end_index=10,
+              uri='https://example.com',
+          )
+      ]
+  )
+
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=None,  # No content - blocked case
+              finish_reason=types.FinishReason.RECITATION,
+              finish_message='Response blocked due to recitation triggered',
+              citation_metadata=citation_metadata,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.citation_metadata == citation_metadata
+  assert response.error_code == types.FinishReason.RECITATION
+  assert (
+      response.error_message == 'Response blocked due to recitation triggered'
+  )
+
+
+def test_llm_response_create_empty_content_with_stop_reason():
+  """Empty content + STOP stays a successful response at the model layer.
+
+  Surfacing the empty turn as an error is the flow's job (non-streaming only);
+  the model/streaming layer must not classify a terminal finish-only chunk as
+  an error or it breaks streaming consumers that batch parts across chunks.
+  """
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(parts=[]),
+              finish_reason=types.FinishReason.STOP,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.error_code is None
+  assert response.content is not None
+  assert response.finish_reason == types.FinishReason.STOP
+
+
+def test_llm_response_create_non_empty_parts_with_stop_is_success():
+  """Regression guard: real text + STOP must remain a successful response."""
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  role='model', parts=[types.Part(text='ok')]
+              ),
+              finish_reason=types.FinishReason.STOP,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.error_code is None
+  assert response.content is not None
+
+
+def test_llm_response_create_includes_model_version():
+  """Test LlmResponse.create() includes model version."""
+  generate_content_response = types.GenerateContentResponse(
+      model_version='gemini-2.5-flash',
+      candidates=[
+          types.Candidate(
+              content=types.Content(parts=[types.Part(text='Response text')]),
+              finish_reason=types.FinishReason.STOP,
+          )
+      ],
+  )
+  response = LlmResponse.create(generate_content_response)
+  assert response.model_version == 'gemini-2.5-flash'
+
+
+def test_get_function_calls_returns_calls_in_order():
+  fc1 = types.FunctionCall(name='a', args={})
+  fc2 = types.FunctionCall(name='b', args={'x': 1})
+  response = LlmResponse(
+      content=types.Content(
+          parts=[
+              types.Part(function_call=fc1),
+              types.Part(text='ignored'),
+              types.Part(function_call=fc2),
+          ]
+      )
+  )
+  assert response.get_function_calls() == [fc1, fc2]
+
+
+def test_get_function_calls_empty_when_no_content():
+  assert LlmResponse().get_function_calls() == []
+
+
+def test_get_function_calls_empty_when_no_parts():
+  response = LlmResponse(content=types.Content(parts=None))
+  assert response.get_function_calls() == []
+
+
+def test_get_function_responses_returns_responses_in_order():
+  fr1 = types.FunctionResponse(name='a', response={'r': 1})
+  fr2 = types.FunctionResponse(name='b', response={'r': 2})
+  response = LlmResponse(
+      content=types.Content(
+          parts=[
+              types.Part(function_response=fr1),
+              types.Part(text='ignored'),
+              types.Part(function_response=fr2),
+          ]
+      )
+  )
+  assert response.get_function_responses() == [fr1, fr2]
+
+
+def test_get_function_responses_empty_when_no_content():
+  assert LlmResponse().get_function_responses() == []
+
+
+def test_get_function_responses_empty_when_no_parts():
+  response = LlmResponse(content=types.Content(parts=None))
+  assert response.get_function_responses() == []
+
+
+def test_environment_id_defaults_to_none_and_roundtrips():
+  resp = LlmResponse()
+  assert resp.environment_id is None
+
+  resp.environment_id = 'env_abc'
+  dumped = resp.model_dump(exclude_none=True)
+  assert dumped['environment_id'] == 'env_abc'
+  assert LlmResponse.model_validate(dumped).environment_id == 'env_abc'