o
    hNs                  	   @   s  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	m
Z
 d dl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mZ d dlmZ erpd dlmZmZmZmZmZm Z m!Z!m"Z" d dl#m$Z$ d dl%m&Z& z d dl'm(Z( d dl)m*Z*m+Z+m,Z,m-Z- d dl.m/Z/ d dl0m1Z1 W n e2y   edw zd dl3m4Z4 W n e2y   dZ4Y nw ej5ej6ej7ej8ej9ej6ej:ej;dZ<G dd deZ=G dd dZ>G dd de*Z?dd Z@dd ZAd d! ZBd"d# ZCd$d% ZDd&d' ZEd(d) ZFd*d+ ZGd,d- ZHd.d/ ZIdS )0    N)OrderedDictwraps)set_ai_pipeline_name)set_data_normalizedget_start_span_function)OPSPANDATA)DidNotEnableIntegration)should_send_default_pii)
_get_valueset_span_errored)loggercapture_internal_exceptions)TYPE_CHECKING)AnyAsyncIteratorCallableDictIteratorListOptionalUnion)UUID)Span)AgentFinish)BaseCallbackHandlerBaseCallbackManager	Callbacksmanager)BaseMessage)	LLMResultzlangchain not installed)AgentExecutor)frequency_penaltyfunction_call
max_tokenspresence_penaltytemperature
tool_callstop_ktop_pc                   @   s4   e Zd ZdZde ZdZd
ddZedd Zd	S )LangchainIntegration	langchainzauto.ai.   Tc                 C   s   || _ || _d S N)include_prompts	max_spans)selfr0   r1    r3   c/var/www/html/pro-man-master/venv/lib/python3.10/site-packages/sentry_sdk/integrations/langchain.py__init__F   s   
zLangchainIntegration.__init__c                   C   s4   t tjt_td urttjt_ttjt_d S d S r/   )_wrap_configurer    
_configurer#   _wrap_agent_executor_invokeinvoke_wrap_agent_executor_streamstreamr3   r3   r3   r4   
setup_onceK   s
   zLangchainIntegration.setup_onceN)Tr.   )	__name__
__module____qualname__
identifieroriginr1   r5   staticmethodr<   r3   r3   r3   r4   r,   ?   s    

r,   c                   @   s    e Zd ZdZg ZdZdd ZdS )WatchedSpanNFc                 C   s
   || _ d S r/   )span)r2   rD   r3   r3   r4   r5   Z   s   
zWatchedSpan.__init__)r=   r>   r?   rD   childrenis_pipeliner5   r3   r3   r3   r4   rC   U   s
    rC   c                   @   s   e Zd ZdZdd Zdd Zdd Zdd	 Zd
d Zdd Z	ddddddZ
dd Zdd Zdd Zdd Zdd Zdd Zdd Zd d! Zd"d# ZdS )$SentryLangchainCallbackz+Callback handler that creates Sentry spans.c                 C   s   t  | _|| _|| _d S r/   )r   span_mapmax_span_map_sizer0   )r2   rI   r0   r3   r3   r4   r5   b   s   
z SentryLangchainCallback.__init__c                 C   sF   t | j| jkr!| jjdd\}}| || t | j| jksd S d S )NF)last)lenrH   rI   popitem
_exit_span)r2   run_idwatched_spanr3   r3   r4   gc_span_maph   s   z#SentryLangchainCallback.gc_span_mapc                 C   s   t  7 |r|| jvr	 W d    d S | j| }|j}t| t||j |d d d  | j|= W d    d S 1 s=w   Y  d S r/   )r   rH   rD   r   
sentry_sdkcapture_exceptionscope__exit__)r2   rN   error	span_datarD   r3   r3   r4   _handle_erroro   s   

"z%SentryLangchainCallback._handle_errorc                 C   s   |j |jd}||j |S )N)rolecontent)typerY   updateadditional_kwargs)r2   messageparsedr3   r3   r4   _normalize_langchain_message~   s   z4SentryLangchainCallback._normalize_langchain_messagec                 K   sv   d }|r| j |}|rt|jjdi |}|j| |d u r+ttjdi |}|j	  || j |< | 
  |S )Nr3   )rH   getrC   rD   start_childrE   appendrQ   
start_span	__enter__rP   )r2   rN   	parent_idkwargsrO   parent_spanr3   r3   r4   _create_span   s   

z$SentryLangchainCallback._create_spanc                 C   s*   |j rtd  |jd d d  | j|= d S r/   )rF   r   rD   rT   rH   )r2   rV   rN   r3   r3   r4   rM      s   z"SentryLangchainCallback._exit_spanN)tagsparent_run_idmetadatac                K   s~  t   |s	 W d   dS |di }||di  |dp.|dp.|dp.d}	| j||tj|dp;d	tjd
}
|
j}|	rL|	t
j|	 |dd}d|v r^|	t
jd nd|v ri|	t
jd t D ]\}}||v r|| durt|||| dd qmt||d t r| jrt|t
j|dd W d   dS W d   dS W d   dS 1 sw   Y  dS )zRun when LLM starts running.Ninvocation_paramsrf   model
model_namemodel_id namezLangchain LLM calloprq   rA   _type	anthropicopenaiFunpacktools)r   r`   r[   rh   r   GEN_AI_PIPELINEr,   rA   rD   set_datar	   GEN_AI_REQUEST_MODELGEN_AI_SYSTEMDATA_FIELDSitemsr   _set_tools_on_spanr   r0   GEN_AI_REQUEST_MESSAGES)r2   
serializedpromptsrN   ri   rj   rk   rf   
all_paramsrm   rO   rD   ai_typekey	attributer3   r3   r4   on_llm_start   sZ   

))"z$SentryLangchainCallback.on_llm_startc                K   s  t   |s	 W d   dS |di }||di  |dp.|dp.|dp.d}| j||dtjd	|  tjd
}|j	}|
tjd |rV|
tj| |dd}	d|	v rh|
tjd nd|	v rs|
tjd t D ]\}
}|
|v r||
 durt||||
 dd qwt||d t r| jrg }|D ]}|D ]
}|| | qqt|tj|dd W d   dS W d   dS W d   dS 1 sw   Y  dS )z#Run when Chat Model starts running.Nrl   rf   rm   rn   ro   rp   rj   zchat rr   chatrt   ru   rv   Frw   ry   )r   r`   r[   rh   r   GEN_AI_CHATstripr,   rA   rD   r{   r	   GEN_AI_OPERATION_NAMEr|   r}   r~   r   r   r   r   r0   rb   r_   r   )r2   r   messagesrN   rf   r   rm   rO   rD   r   r   r   normalized_messageslist_r]   r3   r3   r4   on_chat_model_start   sj   
''"z+SentryLangchainCallback.on_chat_model_startc                K   s   t  ? |r|| jvr	 W d   dS | j| }|j}t r/| jr/t|tjdd |jD  t	|| | 
|| W d   dS 1 sEw   Y  dS )z!Run when Chat Model ends running.Nc                 S      g | ]	}d d |D qS )c                 S      g | ]}|j qS r3   text.0xr3   r3   r4   
<listcomp>      zHSentryLangchainCallback.on_chat_model_end.<locals>.<listcomp>.<listcomp>r3   r   r   r3   r3   r4   r         z=SentryLangchainCallback.on_chat_model_end.<locals>.<listcomp>)r   rH   rD   r   r0   r   r	   GEN_AI_RESPONSE_TEXTgenerations_record_token_usagerM   )r2   responserN   rf   rV   rD   r3   r3   r4   on_chat_model_end  s   

"z)SentryLangchainCallback.on_chat_model_endc          
   	   K   s  t   |r|| jvr	 W d   dS | j| }|j}z	|jd d }W n ty0   d}Y nw |durz|jd}|durG|tj	| W n	 t
yQ   Y nw z|jd}|durd|tj| W n	 t
yn   Y nw z t r| jrt|jdd}	|	dur|	g krt|tj|	dd W n	 t
y   Y nw t r| jrt|tjdd	 |jD  t|| | || W d   dS 1 sw   Y  dS )
zRun when LLM ends running.Nr   rn   finish_reasonr)   Frw   c                 S   r   )c                 S   r   r3   r   r   r3   r3   r4   r   U  r   zASentryLangchainCallback.on_llm_end.<locals>.<listcomp>.<listcomp>r3   r   r3   r3   r4   r   U  r   z6SentryLangchainCallback.on_llm_end.<locals>.<listcomp>)r   rH   rD   r   
IndexErrorgeneration_infor`   r{   r	   GEN_AI_RESPONSE_MODELAttributeErrorGEN_AI_RESPONSE_FINISH_REASONSr   r0   getattrr]   r   GEN_AI_RESPONSE_TOOL_CALLSr   r   rM   )
r2   r   rN   rf   rV   rD   
generationresponse_modelr   r)   r3   r3   r4   
on_llm_end$  sj   

"z"SentryLangchainCallback.on_llm_endc                K      |  || dS )zRun when LLM errors.NrW   r2   rU   rN   rf   r3   r3   r4   on_llm_error[     z$SentryLangchainCallback.on_llm_errorc                K   r   )zRun when Chat Model errors.Nr   r   r3   r3   r4   on_chat_model_error`  r   z+SentryLangchainCallback.on_chat_model_errorc                K   s   t  7 |r|| jvr	 W d    d S | j| }|j}t r,| jr,t|tj|j	  | 
|| W d    d S 1 s=w   Y  d S r/   )r   rH   rD   r   r0   r   r	   r   return_valuesr   rM   )r2   finishrN   rf   rV   rD   r3   r3   r4   on_agent_finishe  s   
"z'SentryLangchainCallback.on_agent_finishc          	      K   s  t  w |s	 W d   dS |dp|dpd}| j||dtjd|  tjd}|j}|	t
jd |	t
j| |d}|durP|	t
j| t rj| jrrt|t
j|d	|g W d   dS W d   dS W d   dS 1 s}w   Y  dS )
zRun when tool starts running.Nrq   rp   rj   zexecute_tool rr   execute_tooldescriptioninputs)r   r`   rh   r   GEN_AI_EXECUTE_TOOLr   r,   rA   rD   r{   r	   r   GEN_AI_TOOL_NAMEGEN_AI_TOOL_DESCRIPTIONr   r0   r   GEN_AI_TOOL_INPUT)	r2   r   	input_strrN   rf   	tool_namerO   rD   tool_descriptionr3   r3   r4   on_tool_startu  s:   
"z%SentryLangchainCallback.on_tool_startc                K   s   t  4 |r|| jvr	 W d   dS | j| }|j}t r)| jr)t|tj| | || W d   dS 1 s:w   Y  dS )zRun when tool ends running.N)	r   rH   rD   r   r0   r   r	   GEN_AI_TOOL_OUTPUTrM   )r2   outputrN   rf   rV   rD   r3   r3   r4   on_tool_end  s   
"z#SentryLangchainCallback.on_tool_endc                O   r   )zRun when tool errors.Nr   )r2   rU   rN   argsrf   r3   r3   r4   on_tool_error  r   z%SentryLangchainCallback.on_tool_error)r=   r>   r?   __doc__r5   rP   rW   r_   rh   rM   r   r   r   r   r   r   r   r   r   r   r3   r3   r3   r4   rG   _   s*    ;87 rG   c                 C   sD   | sdS t | dpt | d}t | dpt | d}t | d}|||fS )NNNNprompt_tokensinput_tokenscompletion_tokensoutput_tokenstotal_tokensr   )token_usager   r   r   r3   r3   r4   _extract_tokens  s   

r   c           
      C   s   | sdS d}d}d}| D ]0}|D ]+}t |}t|\}}}	||dur$|nd7 }||dur.|nd7 }||	dur8|	nd7 }qq|dkrC|nd|dkrJ|nd|dkrR|fS dfS )z8Extract token usage from response.generations structure.r   r   N)_get_token_usager   )
r   total_inputtotal_outputtotal_totalgen_listgenr   r   r   r   r3   r3   r4    _extract_tokens_from_generations  s&   
r   c                 C   s   d}t | d}|dur|D ]}t ||}|dur|  S qt | d}|dur8|D ]}t ||}|dur7|  S q(|D ]}t | |}|durI|  S q:dS )zM
    Check multiple paths to extract token usage from different objects.
    )usager   usage_metadatar]   N
llm_outputr   )objpossible_namesr]   rq   r   r   r3   r3   r4   r     s*   




r   c                 C   sv   t |}|rt|\}}}nt|j\}}}|d ur!| tj| |d ur,| tj| |d ur9| tj| d S d S r/   )	r   r   r   r   r{   r	   GEN_AI_USAGE_INPUT_TOKENSGEN_AI_USAGE_OUTPUT_TOKENSGEN_AI_USAGE_TOTAL_TOKENS)rD   r   r   r   r   r   r3   r3   r4   r     s   
r   c                 C   s   t | dd}t |dd}t |di }t | ddp't |ddp'|dp'|d}|r2t|dkr2|nd}z"d}t|dkrD|d d	}|du rR|d	}W ||fS W ||fS  tyb   Y ||fS w )
z?
    Get the agent name and available tools for the agent.
    agentNrunnableconfigry   available_toolsr      run_name)r   r`   rK   	Exception)r   r   rf   r   r   runnable_configry   
agent_namer3   r3   r4   _get_request_data  s0   
r   c              	   C   s  | sdS t | ttfsdS g }| D ]}zt |trtd|v r=t |d tr=|d }|d|dd}|d r<|| n~d|v rR|d|dd}|| ni|dp`|dp`|d}|rs|||dpo|dd nGt|drt|ddt|ddpt|ddd}|d r|| n%t|d	r||jt|d
dd nt	|}|r|dkr||dd W q t
y   Y qw |r|S dS )z/Parse and simplify tools into a cleaner format.Nfunctionrq   r   )rq   r   r   function_namedescr=   r   rp   )
isinstancelisttupledictr`   rb   hasattrr   r=   strr   )ry   simplified_toolstoolfuncsimplified_toolrq   tool_strr3   r3   r4   _simplify_langchain_tools   sp   









r   c                 C   s2   |durt |}|rt| tj|dd dS dS dS )z9Set available tools data on a span if tools are provided.NFrw   )r   r   r	   GEN_AI_REQUEST_AVAILABLE_TOOLS)rD   ry   r   r3   r3   r4   r   c  s   
r   c                    s   t  	 	 d fdd	}|S )Nc           	         s>  t  t}|d u r | ||g|R i |S |pg }t|tr%|j}n$t|tr.|g}nt|tr6|}nt	
d|  | ||g|R i |S t|trR|j}n
t|trZ|}ng }tdd t||D st|j|j}t|tr| }g |j||_nt|tr||g}ng ||} | ||g|R i |S )NzUnknown callback type: %sc                 s   s    | ]}t |tV  qd S r/   )r   rG   )r   cbr3   r3   r4   	<genexpr>  s
    
z9_wrap_configure.<locals>.new_configure.<locals>.<genexpr>)rQ   
get_clientget_integrationr,   r   r   handlersr   r   r   debugany	itertoolschainrG   r1   r0   copy)	callback_manager_clsinheritable_callbackslocal_callbacksr   rf   integrationcallbacks_listinheritable_callbacks_listsentry_handlerfr3   r4   new_configuret  sv   




	




z&_wrap_configure.<locals>.new_configure)NNr   )r
  r  r3   r	  r4   r6   q  s
   Mr6   c                       t   fdd}|S )Nc                    s6  t  t}|d u r | g|R i |S t| ||\}}t }|tj|r+d| ndtjd`}|r;|	t
j| |	t
jd |	t
jd t||  | g|R i |}|d}	|	d urrt rr|jrrt|t
j|	gdd |d}
|
d urt r|jrt|t
j|
 |W  d    S 1 sw   Y  d S )Ninvoke_agent invoke_agentrr   Finputrw   r   )rQ   r   r   r,   r   r   r   GEN_AI_INVOKE_AGENTrA   r{   r	   GEN_AI_AGENT_NAMEr   GEN_AI_RESPONSE_STREAMINGr   r`   r   r0   r   r   r   )r2   r   rf   r  r   ry   start_span_functionrD   resultr  r   r	  r3   r4   
new_invoke  sF   


$z/_wrap_agent_executor_invoke.<locals>.new_invoker   )r
  r  r3   r	  r4   r8     s   -r8   c                    r  )Nc           
         sB  t  t  d u r| g|R i |S t| ||\}}t }|tjd|  tj	d
  |r<tj| tjd tjd t| t|dkr\|d dnd }|d urrt rr jrrttj|gdd	 | g|R i |}| fd
d} fdd}	tt|dkr|	 }|S | }|S )Nr  rr   r  Tr   r   r  Frw   c                  3   sl    D ]} | V  qz|  d}W n ty   d }Y nw |d ur-t r- jr-ttj| d d d  d S Nr   r`   r   r   r0   r   r	   r   rT   eventr   r  old_iteratorrD   r3   r4   new_iterator$  s   zE_wrap_agent_executor_stream.<locals>.new_stream.<locals>.new_iteratorc                    sv   2 z	3 d H W } | V  q6 z|  d}W n ty    d }Y nw |d ur2t r2 jr2ttj| d d d  d S r  r  r  r  r3   r4   new_iterator_async7  s    zK_wrap_agent_executor_stream.<locals>.new_stream.<locals>.new_iterator_asyncz<class 'async_generator'>)rQ   r   r   r,   r   r   r   r  r   rA   rd   r{   r	   r  r   r  r   rK   r`   r   r0   r   r   r   rZ   )
r2   r   rf   r   ry   r  r  r  r  r  r	  r  r4   
new_stream  sD   
z/_wrap_agent_executor_stream.<locals>.new_streamr   )r
  r  r3   r	  r4   r:     s   Tr:   )Jr   collectionsr   	functoolsr   rQ   sentry_sdk.ai.monitoringr   sentry_sdk.ai.utilsr   r   sentry_sdk.constsr   r	   sentry_sdk.integrationsr
   r   sentry_sdk.scoper   sentry_sdk.tracing_utilsr   r   sentry_sdk.utilsr   r   typingr   r   r   r   r   r   r   r   r   uuidr   sentry_sdk.tracingr   langchain_core.agentsr   langchain_core.callbacksr   r   r   r    langchain_core.messagesr!   langchain_core.outputsr"   ImportErrorlangchain.agentsr#    GEN_AI_REQUEST_FREQUENCY_PENALTYr   GEN_AI_REQUEST_MAX_TOKENSGEN_AI_REQUEST_PRESENCE_PENALTYGEN_AI_REQUEST_TEMPERATUREGEN_AI_REQUEST_TOP_KGEN_AI_REQUEST_TOP_Pr~   r,   rC   rG   r   r   r   r   r   r   r   r6   r8   r:   r3   r3   r3   r4   <module>   sj    (

  MCT4