o
    S۷i"                     @   sf  d dl Z d dlZd dlmZmZmZmZmZ d dlm	Z	m
Z
mZ d dlmZ d dlmZ d dlmZ d dlmZmZ d dlmZ d d	lmZmZ d d
lmZ erjd dlmZ d dlmZmZ d dl m!Z! d dl"m#Z# zd dlZW n e$y{   edw d-ddZ%ddde&ddfddZ'd.ddZ(d/d!d"Z)			#	$		d0d%d&Z*d1d(d)Z+			'	*		d2d+d,Z,dS )3    N)GEN_AI_ALLOWED_MESSAGE_ROLESnormalize_message_rolesset_data_normalizednormalize_message_roletruncate_and_annotate_messages)SPANDATA
SPANSTATUSOP)DidNotEnable)should_send_default_pii)set_span_errored)event_from_exceptionsafe_serialize)_transform_system_instructions)_is_system_instruction_get_system_instructions)TYPE_CHECKING)Any)UsageTResponseInputItem)Span)TextPartzOpenAI Agents not installedexcr   returnc                 C   s6   t   t| t jdddd\}}tj||d d S )Nopenai_agentsF)typehandled)client_options	mechanism)hint)r   r   
sentry_sdk
get_clientoptionscapture_event)r   eventr    r%   a/home/ubuntu/vllm_env/lib/python3.10/site-packages/sentry_sdk/integrations/openai_agents/utils.py_capture_exception&   s   
r'   spanr   errorc                 C   s\   t |  | dd t|dr| d|jj t|dr*t|}|r,| d| d S d S d S )Nzspan.statusr)   	__class__z
error.type__str__zerror.message)r   set_datahasattrr*   __name__str)r(   r)   error_messager%   r%   r&   _record_exception_on_span1   s   

r1   sentry_sdk.tracing.Spanagentagents.Agentc                 C   s  |  tjd |  tj|j |jjr|  tj|jj d }|jr/t	|jdr+|jjn|j}nt	|dr7|j
}|r@|  tj| |jjrM|  tj|jj |jjrZ|  tj|jj |jjrg|  tj|jj |jjrt|  tj|jj t|jdkr|  tjtdd |jD  d S d S )Nopenaimodel_sentry_request_modelr   c                 S   s   g | ]}t |qS r%   )vars).0toolr%   r%   r&   
<listcomp>k   s    z#_set_agent_data.<locals>.<listcomp>)r,   r   GEN_AI_SYSTEMGEN_AI_AGENT_NAMEnamemodel_settings
max_tokensGEN_AI_REQUEST_MAX_TOKENSr6   r-   r7   GEN_AI_REQUEST_MODELpresence_penaltyGEN_AI_REQUEST_PRESENCE_PENALTYtemperatureGEN_AI_REQUEST_TEMPERATUREtop_pGEN_AI_REQUEST_TOP_Pfrequency_penalty GEN_AI_REQUEST_FREQUENCY_PENALTYlentoolsGEN_AI_REQUEST_AVAILABLE_TOOLSr   )r(   r3   
model_namer%   r%   r&   _set_agent_data>   sJ   


rO   usager   c                 C   sX   |  tj|j |  tj|jj |  tj|j |  tj	|j
j |  tj|j d S )N)r,   r   GEN_AI_USAGE_INPUT_TOKENSinput_tokens GEN_AI_USAGE_INPUT_TOKENS_CACHEDinput_tokens_detailscached_tokensGEN_AI_USAGE_OUTPUT_TOKENSoutput_tokens$GEN_AI_USAGE_OUTPUT_TOKENS_REASONINGoutput_tokens_detailsreasoning_tokensGEN_AI_USAGE_TOTAL_TOKENStotal_tokens)r(   rP   r%   r%   r&   _set_usage_datao   s   r]   get_response_kwargsdict[str, Any]c                 C   sX  t  sd S g }|dg }g }|d}|d ur |d|d t|}|t|7 }t|dkr:| tjt	
| dd |D }|D ]I}d|v rit|d}	|d	}
||	t|
trcd|
d
gn|
d qC|ddkr{|tj|gd qC|ddkr|tj|gd qCt|}t }t|| |}|d urt| tj|dd d S d S )Ninputsystem_instructionstext)r   contentr   c                 S   s   g | ]}t |s|qS r%   )r   )r9   messager%   r%   r&   r;      s
    z#_set_input_data.<locals>.<listcomp>rolerc   )r   rb   )re   rc   r   function_callfunction_call_outputF)unpack)r   getappendr   r   rK   r,   r   GEN_AI_SYSTEM_INSTRUCTIONSjsondumpsr   
isinstancer/   r   	ASSISTANTTOOLr   r    get_current_scoper   r   GEN_AI_REQUEST_MESSAGES)r(   r^   request_messagesmessagesinstructions_text_partsexplicit_instructionsra   non_system_messagesrd   normalized_rolerc   normalized_messagesscopemessages_datar%   r%   r&   _set_input_data}   sr   


r|   resultc              
   C   s   t  sd S g g d}|jD ]8}|jdkr|d |  q|jdkrE|jD ]}z
|d |j W q& tyD   |d |  Y q&w qt|d dkrY| 	t
jt|d  t|d dkrlt| t
j|d  d S d S )N)responser:   rf   r:   rd   r~   r   )r   outputr   rj   dictrc   rb   AttributeErrorrK   r,   r   GEN_AI_RESPONSE_TOOL_CALLSr   r   GEN_AI_RESPONSE_TEXT)r(   r}   output_messagesr   output_messager%   r%   r&   _set_output_data   s2   



r   agents.Resultc              	   C   s   |j D ]Q}|jjdkrTtjtjd|j | jd3}|	t
jd |	t
j|j t r<|	t
j|j |	t
j|j  |jrE|tj W d    n1 sOw   Y  qd S )NMcpCallzexecute_tool )opdescriptionstart_timestampmcp)r   r*   r.   r    
start_spanr	   GEN_AI_EXECUTE_TOOLr>   r   r,   r   GEN_AI_TOOL_TYPEGEN_AI_TOOL_NAMEr   GEN_AI_TOOL_INPUT	argumentsGEN_AI_TOOL_OUTPUTr)   
set_statusr   INTERNAL_ERROR)r(   r}   r   execute_tool_spanr%   r%   r&   _create_mcp_execute_tool_spans   s.   

r   )r   r   r   N)r(   r2   r3   r4   r   N)r(   r2   rP   r   r   N)r(   r2   r^   r_   r   N)r(   r2   r}   r   r   N)r(   r2   r}   r   r   N)-rl   r    sentry_sdk.ai.utilsr   r   r   r   r   sentry_sdk.constsr   r   r	   sentry_sdk.integrationsr
   sentry_sdk.scoper   sentry_sdk.tracing_utilsr   sentry_sdk.utilsr   r   %sentry_sdk.ai._openai_completions_apir   #sentry_sdk.ai._openai_responses_apir   r   typingr   r   agentsr   r   sentry_sdk.tracingr   sentry_sdk._typesr   ImportErrorr'   	Exceptionr1   rO   r]   r|   r   r   r%   r%   r%   r&   <module>   sP    


1

L