diff --git a/src/agentscope/model/_dashscope_model.py b/src/agentscope/model/_dashscope_model.py index c17a08e703..99e2e66512 100644 --- a/src/agentscope/model/_dashscope_model.py +++ b/src/agentscope/model/_dashscope_model.py @@ -368,6 +368,10 @@ async def _parse_dashscope_stream_response( input_tokens=chunk.usage.input_tokens, output_tokens=chunk.usage.output_tokens, time=(datetime.now() - start_datetime).total_seconds(), + prompt_tokens_details=chunk.usage.get( + "prompt_tokens_details", + None, + ), ) parsed_chunk = ChatResponse( @@ -469,6 +473,10 @@ async def _parse_dashscope_generation_response( input_tokens=response.usage.input_tokens, output_tokens=response.usage.output_tokens, time=(datetime.now() - start_datetime).total_seconds(), + prompt_tokens_details=response.usage.get( + "prompt_tokens_details", + None, + ), ) parsed_response = ChatResponse( diff --git a/src/agentscope/model/_model_usage.py b/src/agentscope/model/_model_usage.py index fcc76120ff..5e314be6d3 100644 --- a/src/agentscope/model/_model_usage.py +++ b/src/agentscope/model/_model_usage.py @@ -19,5 +19,8 @@ class ChatUsage(DictMixin): time: float """The time used in seconds.""" + prompt_tokens_details: dict = field(default_factory=dict) + """The details of cached tokens. """ + type: Literal["chat"] = field(default_factory=lambda: "chat") """The type of the usage, must be `chat`.""" diff --git a/src/agentscope/tracing/_attributes.py b/src/agentscope/tracing/_attributes.py index 991481b5af..5a92a5411b 100644 --- a/src/agentscope/tracing/_attributes.py +++ b/src/agentscope/tracing/_attributes.py @@ -68,6 +68,9 @@ class SpanAttributes: GEN_AI_USAGE_OUTPUT_TOKENS = GenAIAttributes.GEN_AI_USAGE_OUTPUT_TOKENS """The gen ai usage output tokens.""" + GEN_AI_USAGE_CACHED_TOKENS = "gen_ai.dash.usage.cached_tokens" + """The gen ai usage cached tokens.""" + # GenAI Message Attributes GEN_AI_INPUT_MESSAGES = GenAIAttributes.GEN_AI_INPUT_MESSAGES """The gen ai input messages.""" diff --git a/src/agentscope/tracing/_extractor.py b/src/agentscope/tracing/_extractor.py index 6ca71ea87e..5af90a562f 100644 --- a/src/agentscope/tracing/_extractor.py +++ b/src/agentscope/tracing/_extractor.py @@ -378,6 +378,13 @@ def _get_llm_response_attributes( attributes[ SpanAttributes.GEN_AI_USAGE_OUTPUT_TOKENS ] = chat_response.usage.output_tokens + if ( + hasattr(chat_response.usage, "prompt_tokens_details") + and chat_response.usage.prompt_tokens_details + ): + attributes[ + SpanAttributes.GEN_AI_USAGE_CACHED_TOKENS + ] = _serialize_to_str(chat_response.usage.prompt_tokens_details) output_messages = _get_llm_output_messages(chat_response) if output_messages: