feat(convo): add new fields to conversation api (dapr#902)

* feat(convo): add new fields to conversation api

Signed-off-by: Samantha Coyle <sam@diagrid.io>

* fix: update proto/grpc code generator and add more tests

Signed-off-by: Samantha Coyle <sam@diagrid.io>

* style: appease linter

Signed-off-by: Samantha Coyle <sam@diagrid.io>

* style: tox -e type fixes

Signed-off-by: Samantha Coyle <sam@diagrid.io>

---------

Signed-off-by: Samantha Coyle <sam@diagrid.io>

Author: sicoyle

dapr/clients/grpc/client.py

@@ -25,8 +25,10 @@
 
 import grpc  # type: ignore
 from google.protobuf.any_pb2 import Any as GrpcAny
+from google.protobuf.duration_pb2 import Duration as GrpcDuration
 from google.protobuf.empty_pb2 import Empty as GrpcEmpty
 from google.protobuf.message import Message as GrpcMessage
+from google.protobuf.struct_pb2 import Struct as GrpcStruct
 from grpc import (  # type: ignore
     RpcError,
     StatusCode,
@@ -1880,6 +1882,8 @@ def converse_alpha2(
         temperature: Optional[float] = None,
         tools: Optional[List[conversation.ConversationTools]] = None,
         tool_choice: Optional[str] = None,
+        response_format: Optional[GrpcStruct] = None,
+        prompt_cache_retention: Optional[GrpcDuration] = None,
     ) -> conversation.ConversationResponseAlpha2:
         """Invoke an LLM using the conversation API (Alpha2) with tool calling support.
 
@@ -1893,6 +1897,8 @@ def converse_alpha2(
             temperature: Optional temperature setting for the LLM to optimize for creativity or predictability
             tools: Optional list of tools available for the LLM to call
             tool_choice: Optional control over which tools can be called ('none', 'auto', 'required', or specific tool name)
+            response_format: Optional response format (google.protobuf.struct_pb2.Struct, ex: json_schema for structured output)
+            prompt_cache_retention: Optional retention for prompt cache (google.protobuf.duration_pb2.Duration)
 
         Returns:
             ConversationResponseAlpha2 containing the conversation results with choices and tool calls
@@ -1949,6 +1955,10 @@ def converse_alpha2(
             request.temperature = temperature
         if tool_choice is not None:
             request.tool_choice = tool_choice
+        if response_format is not None and hasattr(request, 'response_format'):
+            request.response_format.CopyFrom(response_format)
+        if prompt_cache_retention is not None and hasattr(request, 'prompt_cache_retention'):
+            request.prompt_cache_retention.CopyFrom(prompt_cache_retention)
 
         try:
             response, call = self.retry_policy.run_rpc(self._stub.ConverseAlpha2.with_call, request)

dapr/clients/grpc/conversation.py

@@ -338,11 +338,46 @@ class ConversationResultAlpha2Choices:
     message: ConversationResultAlpha2Message
 
 
+@dataclass
+class ConversationResultAlpha2CompletionUsageCompletionTokensDetails:
+    """Breakdown of tokens used in the completion."""
+
+    accepted_prediction_tokens: int = 0
+    audio_tokens: int = 0
+    reasoning_tokens: int = 0
+    rejected_prediction_tokens: int = 0
+
+
+@dataclass
+class ConversationResultAlpha2CompletionUsagePromptTokensDetails:
+    """Breakdown of tokens used in the prompt."""
+
+    audio_tokens: int = 0
+    cached_tokens: int = 0
+
+
+@dataclass
+class ConversationResultAlpha2CompletionUsage:
+    """Token usage for one Alpha2 conversation result."""
+
+    completion_tokens: int = 0
+    prompt_tokens: int = 0
+    total_tokens: int = 0
+    completion_tokens_details: Optional[
+        ConversationResultAlpha2CompletionUsageCompletionTokensDetails
+    ] = None
+    prompt_tokens_details: Optional[ConversationResultAlpha2CompletionUsagePromptTokensDetails] = (
+        None
+    )
+
+
 @dataclass
 class ConversationResultAlpha2:
     """One of the outputs in Alpha2 response from conversation input."""
 
     choices: List[ConversationResultAlpha2Choices] = field(default_factory=list)
+    model: Optional[str] = None
+    usage: Optional[ConversationResultAlpha2CompletionUsage] = None
 
 
 @dataclass
@@ -657,5 +692,38 @@ def _get_outputs_from_grpc_response(
                 )
             )
 
-        outputs.append(ConversationResultAlpha2(choices=choices))
+        model: Optional[str] = None
+        usage: Optional[ConversationResultAlpha2CompletionUsage] = None
+        if hasattr(output, 'model') and getattr(output, 'model', None):
+            model = output.model
+        if hasattr(output, 'usage') and output.usage:
+            u = output.usage
+            completion_details: Optional[
+                ConversationResultAlpha2CompletionUsageCompletionTokensDetails
+            ] = None
+            prompt_details: Optional[ConversationResultAlpha2CompletionUsagePromptTokensDetails] = (
+                None
+            )
+            if hasattr(u, 'completion_tokens_details') and u.completion_tokens_details:
+                cd = u.completion_tokens_details
+                completion_details = ConversationResultAlpha2CompletionUsageCompletionTokensDetails(
+                    accepted_prediction_tokens=getattr(cd, 'accepted_prediction_tokens', 0) or 0,
+                    audio_tokens=getattr(cd, 'audio_tokens', 0) or 0,
+                    reasoning_tokens=getattr(cd, 'reasoning_tokens', 0) or 0,
+                    rejected_prediction_tokens=getattr(cd, 'rejected_prediction_tokens', 0) or 0,
+                )
+            if hasattr(u, 'prompt_tokens_details') and u.prompt_tokens_details:
+                pd = u.prompt_tokens_details
+                prompt_details = ConversationResultAlpha2CompletionUsagePromptTokensDetails(
+                    audio_tokens=getattr(pd, 'audio_tokens', 0) or 0,
+                    cached_tokens=getattr(pd, 'cached_tokens', 0) or 0,
+                )
+            usage = ConversationResultAlpha2CompletionUsage(
+                completion_tokens=getattr(u, 'completion_tokens', 0) or 0,
+                prompt_tokens=getattr(u, 'prompt_tokens', 0) or 0,
+                total_tokens=getattr(u, 'total_tokens', 0) or 0,
+                completion_tokens_details=completion_details,
+                prompt_tokens_details=prompt_details,
+            )
+        outputs.append(ConversationResultAlpha2(choices=choices, model=model, usage=usage))
     return outputs

tests/clients/fake_dapr_server.py

@@ -636,6 +636,20 @@ def ConverseAlpha2(self, request, context):
 
             # Create result for this input
             result = api_v1.ConversationResultAlpha2(choices=choices)
+            if hasattr(result, 'model'):
+                result.model = 'test-llm'
+            if hasattr(result, 'usage'):
+                try:
+                    usage_cls = getattr(api_v1, 'ConversationResultAlpha2CompletionUsage', None)
+                    if usage_cls is not None:
+                        u = usage_cls(
+                            completion_tokens=10,
+                            prompt_tokens=5,
+                            total_tokens=15,
+                        )
+                        result.usage.CopyFrom(u)
+                except Exception:
+                    pass
             outputs.append(result)
 
         return api_v1.ConversationResponseAlpha2(

tests/clients/test_conversation.py

@@ -17,7 +17,9 @@
 import json
 import unittest
 import uuid
+from unittest.mock import Mock, patch
 
+from google.protobuf.struct_pb2 import Struct
 from google.rpc import code_pb2, status_pb2
 
 from dapr.aio.clients import DaprClient as AsyncDaprClient
@@ -37,12 +39,16 @@
     ConversationResponseAlpha2,
     ConversationResultAlpha2,
     ConversationResultAlpha2Choices,
+    ConversationResultAlpha2CompletionUsage,
+    ConversationResultAlpha2CompletionUsageCompletionTokensDetails,
+    ConversationResultAlpha2CompletionUsagePromptTokensDetails,
     ConversationResultAlpha2Message,
     ConversationToolCalls,
     ConversationToolCallsOfFunction,
     ConversationTools,
     ConversationToolsFunction,
     FunctionBackend,
+    _get_outputs_from_grpc_response,
     create_assistant_message,
     create_system_message,
     create_tool_message,
@@ -248,6 +254,14 @@ def test_basic_conversation_alpha2(self):
         self.assertEqual(choice.finish_reason, 'stop')
         self.assertIn('Hello Alpha2!', choice.message.content)
 
+        out = response.outputs[0]
+        if out.model is not None:
+            self.assertEqual(out.model, 'test-llm')
+        if out.usage is not None:
+            self.assertGreaterEqual(out.usage.total_tokens, 15)
+            self.assertGreaterEqual(out.usage.prompt_tokens, 5)
+            self.assertGreaterEqual(out.usage.completion_tokens, 10)
+
     def test_conversation_alpha2_with_system_message(self):
         """Test Alpha2 conversation with system message."""
         system_message = create_system_message('You are a helpful assistant.')
@@ -1107,6 +1121,186 @@ def test_empty_and_none_outputs(self):
         self.assertEqual(response_none.to_assistant_messages(), [])
 
 
+class TestConversationResultAlpha2ModelAndUsage(unittest.TestCase):
+    """Tests for model and usage fields on ConversationResultAlpha2 and related types."""
+
+    def test_result_alpha2_has_model_and_usage_attributes(self):
+        """ConversationResultAlpha2 accepts and exposes model and usage."""
+        msg = ConversationResultAlpha2Message(content='Hi', tool_calls=[])
+        choice = ConversationResultAlpha2Choices(finish_reason='stop', index=0, message=msg)
+        usage = ConversationResultAlpha2CompletionUsage(
+            completion_tokens=10,
+            prompt_tokens=5,
+            total_tokens=15,
+        )
+        result = ConversationResultAlpha2(
+            choices=[choice],
+            model='test-model-1',
+            usage=usage,
+        )
+        self.assertEqual(result.model, 'test-model-1')
+        self.assertIsNotNone(result.usage)
+        self.assertEqual(result.usage.completion_tokens, 10)
+        self.assertEqual(result.usage.prompt_tokens, 5)
+        self.assertEqual(result.usage.total_tokens, 15)
+
+    def test_result_alpha2_model_and_usage_default_none(self):
+        """ConversationResultAlpha2 optional fields default to None when not provided.
+
+        When the API returns a response, model and usage are set from the conversation
+        component. This test only checks that the dataclass defaults are None when
+        constructing with choices only.
+        """
+        msg = ConversationResultAlpha2Message(content='Hi', tool_calls=[])
+        choice = ConversationResultAlpha2Choices(finish_reason='stop', index=0, message=msg)
+        result = ConversationResultAlpha2(choices=[choice])
+        self.assertIsNone(result.model)
+        self.assertIsNone(result.usage)
+
+    def test_usage_completion_and_prompt_details(self):
+        """ConversationResultAlpha2CompletionUsage supports details."""
+        completion_details = ConversationResultAlpha2CompletionUsageCompletionTokensDetails(
+            accepted_prediction_tokens=1,
+            audio_tokens=2,
+            reasoning_tokens=3,
+            rejected_prediction_tokens=0,
+        )
+        prompt_details = ConversationResultAlpha2CompletionUsagePromptTokensDetails(
+            audio_tokens=0,
+            cached_tokens=4,
+        )
+        usage = ConversationResultAlpha2CompletionUsage(
+            completion_tokens=10,
+            prompt_tokens=5,
+            total_tokens=15,
+            completion_tokens_details=completion_details,
+            prompt_tokens_details=prompt_details,
+        )
+        self.assertEqual(usage.completion_tokens_details.accepted_prediction_tokens, 1)
+        self.assertEqual(usage.completion_tokens_details.audio_tokens, 2)
+        self.assertEqual(usage.completion_tokens_details.reasoning_tokens, 3)
+        self.assertEqual(usage.completion_tokens_details.rejected_prediction_tokens, 0)
+        self.assertEqual(usage.prompt_tokens_details.audio_tokens, 0)
+        self.assertEqual(usage.prompt_tokens_details.cached_tokens, 4)
+        self.assertEqual(usage.total_tokens, 15)
+        self.assertEqual(usage.completion_tokens, 10)
+        self.assertEqual(usage.prompt_tokens, 5)
+
+    def test_get_outputs_from_grpc_response_populates_model_and_usage(self):
+        """_get_outputs_from_grpc_response sets model and usage when present on proto."""
+        from unittest import mock
+
+        # Build a mock proto response with one output that has model and usage
+        mock_usage = mock.Mock()
+        mock_usage.completion_tokens = 20
+        mock_usage.prompt_tokens = 8
+        mock_usage.total_tokens = 28
+        mock_usage.completion_tokens_details = None
+        mock_usage.prompt_tokens_details = None
+
+        mock_choice_msg = mock.Mock()
+        mock_choice_msg.content = 'Hello'
+        mock_choice_msg.tool_calls = []
+
+        mock_choice = mock.Mock()
+        mock_choice.finish_reason = 'stop'
+        mock_choice.index = 0
+        mock_choice.message = mock_choice_msg
+
+        mock_output = mock.Mock()
+        mock_output.model = 'gpt-4o-mini'
+        mock_output.usage = mock_usage
+        mock_output.choices = [mock_choice]
+
+        mock_response = mock.Mock()
+        mock_response.outputs = [mock_output]
+
+        outputs = _get_outputs_from_grpc_response(mock_response)
+        self.assertEqual(len(outputs), 1)
+        out = outputs[0]
+        self.assertEqual(out.model, 'gpt-4o-mini')
+        self.assertIsNotNone(out.usage)
+        self.assertEqual(out.usage.completion_tokens, 20)
+        self.assertEqual(out.usage.prompt_tokens, 8)
+        self.assertEqual(out.usage.total_tokens, 28)
+        self.assertEqual(len(out.choices), 1)
+        self.assertEqual(out.choices[0].message.content, 'Hello')
+
+    def test_get_outputs_from_grpc_response_without_model_usage(self):
+        """_get_outputs_from_grpc_response leaves model and usage None when absent."""
+        from unittest import mock
+
+        mock_choice_msg = mock.Mock()
+        mock_choice_msg.content = 'Echo'
+        mock_choice_msg.tool_calls = []
+
+        mock_choice = mock.Mock()
+        mock_choice.finish_reason = 'stop'
+        mock_choice.index = 0
+        mock_choice.message = mock_choice_msg
+
+        mock_output = mock.Mock(spec=['choices'])
+        mock_output.choices = [mock_choice]
+        # No model or usage attributes
+
+        mock_response = mock.Mock()
+        mock_response.outputs = [mock_output]
+
+        outputs = _get_outputs_from_grpc_response(mock_response)
+        self.assertEqual(len(outputs), 1)
+        out = outputs[0]
+        self.assertIsNone(out.model)
+        self.assertIsNone(out.usage)
+        self.assertEqual(out.choices[0].message.content, 'Echo')
+
+
+class ConverseAlpha2ResponseFormatTests(unittest.TestCase):
+    """Unit tests for converse_alpha2 response_format parameter."""
+
+    def test_converse_alpha2_passes_response_format_on_request(self):
+        """converse_alpha2 sets response_format on the gRPC request when provided."""
+        user_message = create_user_message('Structured output please')
+        input_alpha2 = ConversationInputAlpha2(messages=[user_message])
+        response_format = Struct()
+        response_format.update(
+            {'type': 'json_schema', 'json_schema': {'name': 'test', 'schema': {}}}
+        )
+
+        captured_requests = []
+        mock_choice_msg = Mock()
+        mock_choice_msg.content = 'ok'
+        mock_choice_msg.tool_calls = []
+        mock_choice = Mock()
+        mock_choice.finish_reason = 'stop'
+        mock_choice.index = 0
+        mock_choice.message = mock_choice_msg
+        mock_output = Mock()
+        mock_output.choices = [mock_choice]
+        mock_response = Mock()
+        mock_response.outputs = [mock_output]
+        mock_response.context_id = ''
+        mock_call = Mock()
+
+        def capture_run_rpc(rpc, request, *args, **kwargs):
+            captured_requests.append(request)
+            return (mock_response, mock_call)
+
+        with patch('dapr.clients.health.DaprHealth.wait_for_sidecar'):
+            client = DaprClient('localhost:50011')
+        with patch.object(client.retry_policy, 'run_rpc', side_effect=capture_run_rpc):
+            client.converse_alpha2(
+                name='test-llm',
+                inputs=[input_alpha2],
+                response_format=response_format,
+            )
+
+        self.assertEqual(len(captured_requests), 1)
+        req = captured_requests[0]
+        self.assertTrue(hasattr(req, 'response_format'))
+        self.assertEqual(req.response_format['type'], 'json_schema')
+        self.assertEqual(req.response_format['json_schema']['name'], 'test')
+
+
 class ExecuteRegisteredToolSyncTests(unittest.TestCase):
     def tearDown(self):
         # Cleanup all tools we may have registered by name prefix