feat: Expose log probs of candidates in LlmResponse

fixes #2764

PiperOrigin-RevId: 807516910

Author: seanzhougoogle

src/google/adk/models/llm_response.py

@@ -42,6 +42,8 @@ class LlmResponse(BaseModel):
     custom_metadata: The custom metadata of the LlmResponse.
     input_transcription: Audio transcription of user input.
     output_transcription: Audio transcription of model output.
+    avg_logprobs: Average log probability of the generated tokens.
+    logprobs_result: Detailed log probabilities for chosen and top candidate tokens.
   """
 
   model_config = ConfigDict(
@@ -109,6 +111,12 @@ class LlmResponse(BaseModel):
   output_transcription: Optional[types.Transcription] = None
   """Audio transcription of model output."""
 
+  avg_logprobs: Optional[float] = None
+  """Average log probability of the generated tokens."""
+
+  logprobs_result: Optional[types.LogprobsResult] = None
+  """Detailed log probabilities for chosen and top candidate tokens."""
+
   @staticmethod
   def create(
       generate_content_response: types.GenerateContentResponse,
@@ -131,13 +139,17 @@ def create(
             grounding_metadata=candidate.grounding_metadata,
             usage_metadata=usage_metadata,
             finish_reason=candidate.finish_reason,
+            avg_logprobs=candidate.avg_logprobs,
+            logprobs_result=candidate.logprobs_result,
         )
       else:
         return LlmResponse(
             error_code=candidate.finish_reason,
             error_message=candidate.finish_message,
             usage_metadata=usage_metadata,
             finish_reason=candidate.finish_reason,
+            avg_logprobs=candidate.avg_logprobs,
+            logprobs_result=candidate.logprobs_result,
         )
     else:
       if generate_content_response.prompt_feedback:

tests/unittests/models/test_llm_response.py

@@ -0,0 +1,241 @@
+# Copyright 2025 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_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'