o
    .i                     @   sJ  d dl Z d dlZd dlZd dlmZ d dlmZ d dlZd dlmZ d dl	m
Z
 d dlmZmZmZmZ d dlmZmZmZ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  d dl!m"Z" d dl#m$Z$m%Z%m&Z&m'Z' d dl(m)Z) e)rd dl(m*Z*m+Z+m,Z,m-Z-m.Z.m/Z/m0Z0mZ d dl1m2Z2 d dl3m4Z4 d dl5m6Z6m7Z7 d dl8m9Z9 zDzd dl8m:Z: W n e;y   dZ:Y nw zd dl8m9Z9 W n e;y   dZ9Y nw d dl<m=Z=m>Z> d dl?m@Z@mAZA e)rd dlBmCZCmDZD W n e;y   edw dZEzd dlFmGZGmHZH d dlImJZJ W n e;y   dZEY nw G dd deZKdgdd d!eLd"dfd#d$ZMd%d d&d'd"eNfd(d)ZO	*	+	,	 	-	.	/	0	1	2	"	dhd3d4ZP	-	.	5	6	"	did7d8ZQ	-	.	5	6	9		"	djd:d;ZR	-	.	5	6	9		"	djd<d=ZS	-	.	5	6	9		"	djd>d?ZT	dgd-d.d,d d@d d9ddAeLd"dfdBdCZUdkdFdGZV	dgd-d.d,d d5d6d9ddAeLd"dfdHdIZW		dld-d.d,d d5d6d9ddJdKdAeLd"dfdLdMZX	dgd-d.d,d d5d6d9ddAeLd"dfdNdOZY		dld-d.d,d d5d6d9ddJdKdAeLd"dfdPdQZZ	dgd-d.d,d d5d6d9ddAeLd"dfdRdSZ[dmdTdUZ\dmdVdWZ]dkdXdYZ^dndZd[Z_dnd\d]Z`dkd^d_Zadnd`daZbdndbdcZcddd d"eLfdedfZddS )o    Nwraps)Iterable)consts)record_token_usage)set_data_normalizednormalize_message_rolestruncate_and_annotate_messages&truncate_and_annotate_embedding_inputs)_is_system_instruction_get_system_instructions_transform_system_instructions_get_text_items)r   r   )SPANDATA)DidNotEnableIntegration)should_send_default_pii)set_span_errored)capture_internal_exceptionsevent_from_exceptionsafe_serializereraise)TYPE_CHECKING)AnyListOptionalCallableAsyncIteratorIteratorUnionr   )Span)TextPart)ResponseInputParamSequenceNotStr)Omit)NotGiven)CompletionsAsyncCompletions)
EmbeddingsAsyncEmbeddings)ChatCompletionMessageParamChatCompletionChunkzOpenAI not installedT)	ResponsesAsyncResponses)ResponseCompletedEventFc                	   @   s^   e Zd ZdZde Z		ddd deddd	dfd
dZedddZdd de	d	e
fddZdS )OpenAIIntegrationopenaizauto.ai.TNselfinclude_promptstiktoken_encoding_namezOptional[str]returnc                 C   s0   || _ d | _|d urdd l}||| _d S d S Nr   )r2   tiktoken_encodingtiktokenget_encoding)r1   r2   r3   r7    r9   k/lsinfo/ai/hellotax_ai/llm_service/venv_vllm/lib/python3.10/site-packages/sentry_sdk/integrations/openai.py__init__[   s   zOpenAIIntegration.__init__c                   C   sT   t tjt_ttjt_ttjt_ttjt_t	r(t
tjt_ttjt_d S d S N)_wrap_chat_completion_creater&   create"_wrap_async_chat_completion_creater'   _wrap_embeddings_creater(   _wrap_async_embeddings_creater)   RESPONSES_API_ENABLED_wrap_responses_creater,   _wrap_async_responses_creater-   r9   r9   r9   r:   
setup_onceh   s   zOpenAIIntegration.setup_oncesc                 C   s6   | j d u rdS z	t| j |W S  ty   Y dS w r5   )r6   lenencode_ordinary	Exception)r1   rF   r9   r9   r:   count_tokensv   s   
zOpenAIIntegration.count_tokens)TN)r4   N)__name__
__module____qualname__
identifieroriginboolr;   staticmethodrE   strintrJ   r9   r9   r9   r:   r/   W   s"    

r/   excr   manual_span_cleanupr4   c                 C   sZ   t  }t| |r|d ur|d d d  t| t  jdddd\}}t j||d d S )Nr0   F)typehandled)client_options	mechanism)hint)
sentry_sdkget_current_spanr   __exit__r   
get_clientoptionscapture_event)rT   rU   current_spaneventrZ   r9   r9   r:   _capture_exception   s   
rc   usagenamesz	List[str]c                 C   s6   |D ]}t | |rtt| |trt| |  S qdS r5   )hasattr
isinstancegetattrrS   )rd   re   namer9   r9   r:   
_get_usage   s
   rj   messages.Optional[Iterable[ChatCompletionMessageParam]]responsespanr    streaming_message_responsesOptional[List[str]]rJ   Callable[..., Any]c                    s  d}d}d}d}d}	t |drBt|jddg}t |jdr%t|jjdg}t|jddg}t |jd	r;t|jjd
g}t|jdg}	|dkrz| pIg D ]/}
t|
trX| |
7 }qJt|
try|
d}|d u rgqJt	|}|t
 fdd|D 7 }qJqJ|dkr|d ur|D ]}
| |
7 }qnt |dr|jD ]}t |drt |jdr| |jj7 }q|pd }|pd }|pd }|pd }|	pd }	t||||||	d d S )Nr   rd   input_tokensprompt_tokensinput_tokens_detailscached_tokensoutput_tokenscompletion_tokensoutput_tokens_detailsreasoning_tokenstotal_tokenscontentc                 3   s    | ]} |V  qd S r<   r9   ).0textrJ   r9   r:   	<genexpr>   s    z)_calculate_token_usage.<locals>.<genexpr>choicesmessage)rr   input_tokens_cachedrv   output_tokens_reasoningrz   )rf   rj   rd   rt   rx   rg   rR   dictgetr   sumr   r   r{   r   )rk   rm   rn   ro   rJ   rr   r   rv   r   rz   r   message_content
text_itemschoicer9   r~   r:   _calculate_token_usage   sn   






	


r   kwargsdict[str, Any]c                 C   s   t | tjd tjtjtjtjtjtjtj	d}|
 D ]\}}||}|d ur3t|r3t | || q|d}|d urRt|rTt|dkrVt | tjt| d S d S d S d S )Nr0   )modelstream
max_tokenspresence_penaltyfrequency_penaltytemperaturetop_ptoolsr   )r   r   GEN_AI_SYSTEMGEN_AI_REQUEST_MODELGEN_AI_RESPONSE_STREAMINGGEN_AI_REQUEST_MAX_TOKENSGEN_AI_REQUEST_PRESENCE_PENALTY GEN_AI_REQUEST_FREQUENCY_PENALTYGEN_AI_REQUEST_TEMPERATUREGEN_AI_REQUEST_TOP_Pitemsr   	_is_givenrG   GEN_AI_REQUEST_AVAILABLE_TOOLSr   )rn   r   kwargs_keys_to_attributeskey	attributevaluer   r9   r9   r:   _commmon_set_input_data   s(   	

r   integrationc                 C   s  | d}| d}t r|jst| tjd t| | d S |d u rF|d urFt|rF| tj	t
d|dg t| tjd t| | d S |d u rXt| tjd t| | d S g }|d urjt|rj|d|d t|}|t|7 }t|dkr| tj	t
| t|trt|g}t }t|| |}	|	d urt| tj|	dd t| tjd t| | d S d	d
 |D }
t|
dkrt|
}t }t|| |}	|	d urt| tj|	dd t| tjd t| | d S )Ninstructionsinput	responsesr}   )rV   r{   r   Funpackc                 S      g | ]}t |s|qS r9   ) _is_system_instruction_responsesr|   r   r9   r9   r:   
<listcomp>E  s
    z1_set_responses_api_input_data.<locals>.<listcomp>)r   r   r2   r   r   GEN_AI_OPERATION_NAMEr   r   set_dataGEN_AI_SYSTEM_INSTRUCTIONSjsondumpsappend"_get_system_instructions_responsesr   rG   rg   rR   r   r[   get_current_scoper	   GEN_AI_REQUEST_MESSAGES)rn   r   r   explicit_instructionsrk   instructions_text_partssystem_instructionsnormalized_messagesscopemessages_datanon_system_messagesr9   r9   r:   _set_responses_api_input_data   s   









r   c           	      C   sp  | d}t r|jr|d u rt| tjd t| | d S t|trLt	|g}t
 }t|| |}|d ur>t| tj|dd t| tjd t| | d S t|trVt|trdt| tjd t| | d S t|}||d< t|}t|dkr| tjtt| dd |D }t|dkrt	|}t
 }t|| |}|d urt| tj|dd t| tjd t| | d S )Nrk   chatFr   r   c                 S   r   r9   )"_is_system_instruction_completionsr   r9   r9   r:   r     s    z3_set_completions_api_input_data.<locals>.<listcomp>)r   r   r2   r   r   r   r   rg   rR   r   r[   r   r	   r   r   r   list$_get_system_instructions_completionsrG   r   r   r   r   r   )	rn   r   r   rk   r   r   r   r   r   r9   r9   r:   _set_completions_api_input_dataU  s^   






r   c                 C   s6  | d}t r|jr|d u rt| tjd t| | d S t|trLt| tjd t| | t	|g}t
 }t|| |}|d urJt| tj|dd d S t|trVt|trdt| tjd t| | d S t|}||d< t|dkrt	|}t
 }t|| |}|d urt| tj|dd t| tjd t| | d S )Nr   
embeddingsFr   r   )r   r   r2   r   r   r   r   rg   rR   r   r[   r   r
   GEN_AI_EMBEDDINGS_INPUTr   r   r   rG   )rn   r   r   rk   r   r   r   r9   r9   r:   _set_embeddings_input_data  sT   






r   r   finish_spanc           	   
   C   s  t |drt| tj|j t |drCt r-|jr-dd |jD }t|dkr-t| tj	| t
||| d |j |rA| d d d  d S d S t |drt r|jrg g d}|jD ]8}|jdkrg|d	 |  qV|jd
kr|jD ]}z
|d |j W qo ty   |d |  Y qow qVt|d	 dkrt| tj|d	 dd t|d dkrt| tj	|d  t
||| d |j |r| d d d  d S d S t
||| d |j |r| d d d  d S d S )Nr   r   c                 S   s    g | ]}|j d ur|j  qS r<   )r   
model_dump)r|   r   r9   r9   r:   r     s
    
z+_set_common_output_data.<locals>.<listcomp>r   output)rm   toolfunction_callr   r   rm   Fr   )rf   r   r   GEN_AI_RESPONSE_MODELr   r   r2   r   rG   GEN_AI_RESPONSE_TEXTr   rJ   r]   r   rV   r   r   r{   r}   AttributeErrorGEN_AI_RESPONSE_TOOL_CALLS)	rn   rm   r   r   r   response_textoutput_messagesr   output_messager9   r9   r:   _set_common_output_data  s`   






r   fargsc           	      o   s    t  t}|d u r| |i |S d|vr| |i |S zt|d  W n ty6   | |i | Y S w |d}t jtj	j
d| tjd}|  t||| t }| ||fV }|dd}|rrt|||||dd |S t||||dd |S )	Nrk   r   zchat opri   rO   r   FTr   )r[   r^   get_integrationr/   iter	TypeErrorr   
start_spanr   OPGEN_AI_CHATrO   	__enter__r   timeperf_counter*_set_streaming_completions_api_output_data _set_completions_api_output_data	r   r   r   r   r   rn   
start_timerm   is_streaming_responser9   r9   r:   _new_chat_completion_common  s>   

r   c                 C   6   | d}|d urt|tr|g}t| |||| d S Nrk   r   rg   rR   r   )rn   rm   r   r   r   rk   r9   r9   r:   r   <     
r   r   zOptional[float]c              
      s   | dd urttrgd g  jd
 f	dd}d f	dd}ttjd	krE| _d S | _d S )Nrk   r4   Iterator[ChatCompletionChunk]c               	   3   sh   D ]b}  tj| j t J t| drSd}| jD ]7}t|drNt|jdrNd ur6d u r6t	  |jj
}t |krE g   | |pLd |d7 }qW d    n1 s]w   Y  | V  qt 6 d uruttj t dkrdd  D }t rjrttj| t|j W d    n1 sw   Y  rd d d  d S d S )	Nr   r   deltar{       c                 S      g | ]}d  |qS r   joinr|   chunkr9   r9   r:   r         zT_set_streaming_completions_api_output_data.<locals>.new_iterator.<locals>.<listcomp>r   r   r   r   r   rf   r   r   r   r   r{   rG   r   r   #GEN_AI_RESPONSE_TIME_TO_FIRST_TOKENr   r2   r   r   rJ   r]   xchoice_indexr   r{   all_responses	data_bufr   r   rk   old_iteratorrm   rn   r   ttftr9   r:   new_iteratorc  sV   



z@_set_streaming_completions_api_output_data.<locals>.new_iterator"AsyncIterator[ChatCompletionChunk]c               	     sr  2 zf3 d H W }  tj| j t J t| drWd}| jD ]7}t|drRt|jdrRd ur:d u r:t	  |jj
}t |krI g   | |pPd |d7 }qW d    n1 saw   Y  | V  q6 t 6 d urzttj t dkrdd  D }t rjrttj| t|j W d    n1 sw   Y  rd d d  d S d S )	Nr   r   r   r{   r   r   c                 S   r   r   r   r   r9   r9   r:   r     r   zZ_set_streaming_completions_api_output_data.<locals>.new_iterator_async.<locals>.<listcomp>r   r   r   r9   r:   new_iterator_async  sX   



zF_set_streaming_completions_api_output_data.<locals>.new_iterator_async<class 'async_generator'>r4   r   r4   r  r   rg   rR   	_iteratorrV   rn   rm   r   r   r   r   r  r  r9   r   r:   r   Q  s   
,,r   c                 C   r   Nr   r   rn   rm   r   r   r   r   r9   r9   r:   _set_responses_api_output_data  r   r  c              
      s   | dd urttrgd g  jd
 f	dd}d f	dd}ttjd	krE| _d S | _d S )Nr   r4   r   c               	   3   s`   d} D ]Z}t  J t|dr3d urd u rt  t dkr) g   d |jp1d t|trM	t
j|jj t|jd j d} W d    n1 sWw   Y  |V  qt  8 d urott
j t dkrdd  D }t rjrtt
j| | rt|j W d    n1 sw   Y  rd d d  d S d S )NTr   r   r   Fc                 S   r   r   r   r   r9   r9   r:   r     r   zR_set_streaming_responses_api_output_data.<locals>.new_iterator.<locals>.<listcomp>r   rf   r   r   rG   r   r   rg   r.   r   r   r   rm   r   r   rJ   r   r   r   r2   r   r]   count_tokens_manuallyr   r   	r   r   r   r   r   rm   rn   r   r  r9   r:   r    s^   


z>_set_streaming_responses_api_output_data.<locals>.new_iteratorr  c               	     sj  d} 2 z^3 d H W }t  J t|dr7d ur"d u r"t  t dkr- g   d |jp5d t|trQ	t
j|jj t|jd j d} W d    n1 s[w   Y  |V  q6 t  8 d urttt
j t dkrdd  D }t rjrtt
j| | rt|j W d    n1 sw   Y  rd d d  d S d S )NTr   r   r   Fc                 S   r   r   r   r   r9   r9   r:   r   9  r   zX_set_streaming_responses_api_output_data.<locals>.new_iterator_async.<locals>.<listcomp>r  r  r  r9   r:   r    s`   


zD_set_streaming_responses_api_output_data.<locals>.new_iterator_asyncr  r  r  r  r
  r9   r  r:   (_set_streaming_responses_api_output_data  s   
10r  c                 C   r   r  r   r  r9   r9   r:   _set_embeddings_output_dataO  r   r  c                    &   d
dd t d fdd	}|S )Nr   r   r   r   r4   c                 _     t | g|R i |}z	t|\} }}W n ty) } z	|jW  Y d }~S d }~ww zAz	| |i |}W n1 tye } z%t }t  t| W d    n1 sRw   Y  t	|  W Y d }~nd }~ww |
|W S  ty } z	|jW  Y d }~S d }~ww r<   r   nextStopIterationr   rI   sysexc_infor   rc   r   sendr   r   r   generesultr  r9   r9   r:   _execute_synce  .   
z3_wrap_chat_completion_create.<locals>._execute_syncc                     sB   t  t}|d u sd|vr| i |S  g| R i |S r   r[   r^   r   r/   r   r   r   r   r   r9   r:   _sentry_patched_create_syncz  s   zA_wrap_chat_completion_create.<locals>._sentry_patched_create_syncr   r   r   r   r   r   r4   r   r   r   r   r   r4   r   r   r   r%  r9   r$  r:   r=   d     
r=   c                    r  )Nr   r   r   r   r4   c                      t | g|R i |}z	t|\} }}W n ty- } z|jI d H W  Y d }~S d }~ww zDz| |i |I d H }W n1 tyl } z%t }t  t| W d    n1 sYw   Y  t	|  W Y d }~nd }~ww |
|W S  ty } z	|jW  Y d }~S d }~ww r<   r  r  r9   r9   r:   _execute_async  0   
z:_wrap_async_chat_completion_create.<locals>._execute_asyncc                     sP   t  t}|d u sd|vr| i |I d H S  g| R i |I d H S r   r"  r#  r+  r   r9   r:   _sentry_patched_create_async  s
   zH_wrap_async_chat_completion_create.<locals>._sentry_patched_create_asyncr&  r'  r   r   r.  r9   r-  r:   r?     r)  r?   c                 o   s    t  t}|d u r| |i |S |d}t jtjjd| tj	d}t
||| | ||fV }t||||dd |W  d    S 1 sHw   Y  d S )Nr   zembeddings r   Fr   )r[   r^   r   r/   r   r   r   r   GEN_AI_EMBEDDINGSrO   r   r  )r   r   r   r   r   rn   rm   r9   r9   r:   _new_embeddings_create_common  s$   

$r1  c                    r  )Nr   r   r   r   r4   c                 _   s  t | g|R i |}z	t|\} }}W n ty) } z	|jW  Y d }~S d }~ww zCz	| |i |}W n3 tyg } z't }t  t|dd W d    n1 sTw   Y  t	|  W Y d }~nd }~ww |
|W S  ty } z	|jW  Y d }~S d }~ww NF)rU   r1  r  r  r   rI   r  r  r   rc   r   r  r  r9   r9   r:   r     s.   z._wrap_embeddings_create.<locals>._execute_syncc                     :   t  t}|d u r| i |S  g| R i |S r<   r"  r#  r$  r9   r:   r%       z<_wrap_embeddings_create.<locals>._sentry_patched_create_syncr&  r'  r   r(  r9   r$  r:   r@        
r@   c                    r  )Nr   r   r   r   r4   c                    s  t | g|R i |}z	t|\} }}W n ty- } z|jI d H W  Y d }~S d }~ww zFz| |i |I d H }W n3 tyn } z't }t  t|dd W d    n1 s[w   Y  t	|  W Y d }~nd }~ww |
|W S  ty } z	|jW  Y d }~S d }~ww r2  r3  r  r9   r9   r:   r+    s0   z5_wrap_async_embeddings_create.<locals>._execute_asyncc                     H   t  t}|d u r| i |I d H S  g| R i |I d H S r<   r"  r#  r-  r9   r:   r.    
   zC_wrap_async_embeddings_create.<locals>._sentry_patched_create_asyncr&  r'  r   r/  r9   r-  r:   rA     r6  rA   c           	      o   s    t  t}|d u r| |i |S |d}t jtjjd| tj	d}|
  t||| t }| ||fV }|dd}|rNt|||||dd |S t||||dd |S )Nr   z
responses r   r   FTr   )r[   r^   r   r/   r   r   r   r   GEN_AI_RESPONSESrO   r   r   r   r   r  r  r   r9   r9   r:   _new_responses_create_common  s0   

r:  c                    r  )Nr   r   r   r   r4   c                 _   r  r<   r:  r  r  r   rI   r  r  r   rc   r   r  r  r9   r9   r:   r   "  r!  z-_wrap_responses_create.<locals>._execute_syncc                     r4  r<   r"  r#  r$  r9   r:   r%  7  r5  z;_wrap_responses_create.<locals>._sentry_patched_create_syncr&  r'  r   r(  r9   r$  r:   rC   !  r6  rC   c                    r  )Nr   r   r   r   r4   c                    r*  r<   r;  r  r9   r9   r:   r+  C  r,  z4_wrap_async_responses_create.<locals>._execute_asyncc                     r7  r<   r"  r#  r-  r9   r:   _sentry_patched_responses_asyncX  r8  zE_wrap_async_responses_create.<locals>._sentry_patched_responses_asyncr&  r'  r   )r   r<  r9   r-  r:   rD   B  r6  rD   objc                 C   s0   t durt| t rdS tdurt| trdS dS )zF
    Check for givenness safely across different openai versions.
    NFT)r%   rg   r$   )r=  r9   r9   r:   r   c  s
   r   )T)rk   rl   rm   r   rn   r    ro   rp   rJ   rq   r4   N)rn   r    r   r   r4   N)rn   r    r   r   r   r/   r4   Nr&  )NT)r   rq   r4   rq   )r   r   r4   r   )er  r   r   	functoolsr   collections.abcr   r[   r   sentry_sdk.ai.monitoringr   sentry_sdk.ai.utilsr   r   r	   r
   %sentry_sdk.ai._openai_completions_apir   r   r   r   r   r   #sentry_sdk.ai._openai_responses_apir   r   sentry_sdk.constsr   sentry_sdk.integrationsr   r   sentry_sdk.scoper   sentry_sdk.tracing_utilsr   sentry_sdk.utilsr   r   r   r   typingr   r   r   r   r   r   r   r   sentry_sdk.tracingr    sentry_sdk._typesr!   openai.types.responsesr"   r#   r0   r$   r%   ImportError!openai.resources.chat.completionsr&   r'   openai.resourcesr(   r)   openai.types.chatr*   r+   rB   openai.resources.responsesr,   r-   /openai.types.responses.response_completed_eventr.   r/   rP   rc   rS   rj   r   r   r   r   r   r   r   r   r   r  r  r  r=   r?   r1  r@   rA   r:  rC   rD   r   r9   r9   r9   r:   <module>   s   (
(
G

X
@
@

B/

u

~


"
"

!
!
 
!!