1- # Copyright (c) Microsoft. All rights reserved.
1+ # Copyright (c) Microsoft Corporation.
2+ # Licensed under the MIT License.
3+
4+ from typing import List
25
36from .agent_details import AgentDetails
47from .constants import (
58 GEN_AI_INPUT_MESSAGES_KEY ,
9+ GEN_AI_OPERATION_NAME_KEY ,
10+ GEN_AI_OUTPUT_MESSAGES_KEY ,
611 GEN_AI_PROVIDER_NAME_KEY ,
712 GEN_AI_REQUEST_MODEL_KEY ,
813 GEN_AI_RESPONSE_FINISH_REASONS_KEY ,
914 GEN_AI_RESPONSE_ID_KEY ,
15+ GEN_AI_THOUGHT_PROCESS_KEY ,
1016 GEN_AI_USAGE_INPUT_TOKENS_KEY ,
1117 GEN_AI_USAGE_OUTPUT_TOKENS_KEY ,
1218)
1723
1824
1925class InferenceScope (OpenTelemetryScope ):
20- """Provides OpenTelemetry tracing scope for inference call ."""
26+ """Provides OpenTelemetry tracing scope for generative AI inference operations ."""
2127
2228 @staticmethod
2329 def start (
@@ -26,7 +32,7 @@ def start(
2632 tenant_details : TenantDetails ,
2733 request : Request | None = None ,
2834 ) -> "InferenceScope" :
29- """Create and start a new scope for inference call .
35+ """Creates and starts a new scope for inference tracing .
3036
3137 Args:
3238 details: The details of the inference call
@@ -41,40 +47,102 @@ def start(
4147
4248 def __init__ (
4349 self ,
44- inference_call_details : InferenceCallDetails ,
50+ details : InferenceCallDetails ,
4551 agent_details : AgentDetails ,
4652 tenant_details : TenantDetails ,
4753 request : Request | None = None ,
4854 ):
49- """Initialize the agent invocation scope.
55+ """Initialize the inference scope.
5056
5157 Args:
52- inference_call_details : The details of the inference call
58+ details : The details of the inference call
5359 agent_details: The details of the agent making the call
5460 tenant_details: The details of the tenant
5561 request: Optional request details for additional context
5662 """
5763
5864 super ().__init__ (
5965 kind = "Client" ,
60- operation_name = inference_call_details .operationName .value ,
61- activity_name = f"{ inference_call_details .operationName .value } { inference_call_details .model } ,
66+ operation_name = details .operationName .value ,
67+ activity_name = f"{ details .operationName .value } { details .model } ,
6268 agent_details = agent_details ,
6369 tenant_details = tenant_details ,
6470 )
6571
66- # Set request content if provided
6772 if request :
6873 self .set_tag_maybe (GEN_AI_INPUT_MESSAGES_KEY , request .content )
6974
70- self .set_tag_maybe (GEN_AI_REQUEST_MODEL_KEY , inference_call_details .model )
71- self .set_tag_maybe (GEN_AI_PROVIDER_NAME_KEY , inference_call_details .providerName )
72- self .set_tag_maybe (GEN_AI_USAGE_INPUT_TOKENS_KEY , inference_call_details .inputTokens )
73- self .set_tag_maybe (GEN_AI_USAGE_OUTPUT_TOKENS_KEY , inference_call_details .outputTokens )
75+ self .set_tag_maybe (GEN_AI_OPERATION_NAME_KEY , details .operationName .value )
76+ self .set_tag_maybe (GEN_AI_REQUEST_MODEL_KEY , details .model )
77+ self .set_tag_maybe (GEN_AI_PROVIDER_NAME_KEY , details .providerName )
78+ self .set_tag_maybe (
79+ GEN_AI_USAGE_INPUT_TOKENS_KEY ,
80+ str (details .inputTokens ) if details .inputTokens is not None else None ,
81+ )
82+ self .set_tag_maybe (
83+ GEN_AI_USAGE_OUTPUT_TOKENS_KEY ,
84+ str (details .outputTokens ) if details .outputTokens is not None else None ,
85+ )
7486 self .set_tag_maybe (
7587 GEN_AI_RESPONSE_FINISH_REASONS_KEY ,
76- "," .join (inference_call_details .finishReasons )
77- if inference_call_details .finishReasons
78- else None ,
88+ "," .join (details .finishReasons ) if details .finishReasons else None ,
7989 )
80- self .set_tag_maybe (GEN_AI_RESPONSE_ID_KEY , inference_call_details .responseId )
90+ self .set_tag_maybe (GEN_AI_RESPONSE_ID_KEY , details .responseId )
91+
92+ def record_input_messages (self , messages : List [str ]) -> None :
93+ """Records the input messages for telemetry tracking.
94+
95+ Args:
96+ messages: List of input messages
97+ """
98+ self .set_tag_maybe (GEN_AI_INPUT_MESSAGES_KEY , "," .join (messages ))
99+
100+ def record_output_messages (self , messages : List [str ]) -> None :
101+ """Records the output messages for telemetry tracking.
102+
103+ Args:
104+ messages: List of output messages
105+ """
106+ self .set_tag_maybe (GEN_AI_OUTPUT_MESSAGES_KEY , "," .join (messages ))
107+
108+ def record_input_tokens (self , input_tokens : int ) -> None :
109+ """Records the number of input tokens for telemetry tracking.
110+
111+ Args:
112+ input_tokens: Number of input tokens
113+ """
114+ self .set_tag_maybe (GEN_AI_USAGE_INPUT_TOKENS_KEY , str (input_tokens ))
115+
116+ def record_output_tokens (self , output_tokens : int ) -> None :
117+ """Records the number of output tokens for telemetry tracking.
118+
119+ Args:
120+ output_tokens: Number of output tokens
121+ """
122+ self .set_tag_maybe (GEN_AI_USAGE_OUTPUT_TOKENS_KEY , str (output_tokens ))
123+
124+ def record_response_id (self , response_id : str ) -> None :
125+ """Records the response id for telemetry tracking.
126+
127+ Args:
128+ response_id: Response ID to record
129+ """
130+ if response_id :
131+ self .set_tag_maybe (GEN_AI_RESPONSE_ID_KEY , response_id )
132+
133+ def record_finish_reasons (self , finish_reasons : List [str ]) -> None :
134+ """Records the finish reasons for telemetry tracking.
135+
136+ Args:
137+ finish_reasons: List of finish reasons
138+ """
139+ if finish_reasons :
140+ self .set_tag_maybe (GEN_AI_RESPONSE_FINISH_REASONS_KEY , "," .join (finish_reasons ))
141+
142+ def record_thought_process (self , thought_process : str ) -> None :
143+ """Records the thought process.
144+
145+ Args:
146+ thought_process: The thought process to record
147+ """
148+ self .set_tag_maybe (GEN_AI_THOUGHT_PROCESS_KEY , thought_process )
0 commit comments