§
    šŠtjS"  ã                   ó®   — d dl Z d dlZd dlZd dlZd dl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 d dlmZ ded	ed
dfd„Z G d„ de¦  «        ZdS )é    N)Údeepcopy)ÚAnyÚDictÚListÚOptional)ÚAgentActionÚAgentFinish)ÚBaseCallbackHandler)Ú	LLMResult)Úflatten_dictÚdataÚ	file_pathÚreturnc                 ó‚   — t          |d¦  «        5 }t          j        | |¦  «         ddd¦  «         dS # 1 swxY w Y   dS )zŽSave dict to local file path.

    Parameters:
        data (dict): The dictionary to be saved.
        file_path (str): Local file path.
    ÚwN)ÚopenÚjsonÚdump)r   r   Úoutfiles      ún/var/www/html/CA-Chatbot/venv/lib/python3.11/site-packages/langchain_community/callbacks/sagemaker_callback.pyÚ	save_jsonr      s‡   € õ 
ˆi˜Ñ	Ô	ð ! ÝŒ	�$˜Ñ Ô Ð ð!ð !ð !ñ !ô !ð !ð !ð !ð !ð !ð !ð !øøøð !ð !ð !ð !ð !ð !s   ‘4´8»8c                   óä  ‡ — e Zd ZdZdeddfˆ fd„Zd(d„Zdeeef         de	e         d	eddfd
„Z
ded	eddfd„Zded	eddfd„Zded	eddfd„Zdeeef         deeef         d	eddfd„Zdeeef         d	eddfd„Zded	eddfd„Zdeeef         ded	eddfd„Zded	eddfd„Zded	eddfd„Zded	eddfd„Zded	eddfd„Zded	edefd „Z	 d)d"eeef         d#ed$ed%ee         ddf
d&„Zd(d'„Zˆ xZS )*ÚSageMakerCallbackHandlerz¿Callback Handler that logs prompt artifacts and metrics to SageMaker Experiments.

    Parameters:
        run (sagemaker.experiments.run.Run): Run object where the experiment is logged.
    Úrunr   Nc                 ó°   •— t          ¦   «                              ¦   «          || _        ddddddddddddddœ| _        t	          j        ¦   «         | _        dS )zInitialize callback handler.r   )ÚstepÚstartsÚendsÚerrorsÚtext_ctrÚchain_startsÚ
chain_endsÚ
llm_startsÚllm_endsÚllm_streamsÚtool_startsÚ	tool_endsÚ
agent_endsN)ÚsuperÚ__init__r   ÚmetricsÚtempfileÚmkdtempÚtemp_dir)Úselfr   Ú	__class__s     €r   r*   z!SageMakerCallbackHandler.__init__#   sm   ø€ å‰Œ×ÒÑÔÐàˆŒð ØØØØØØØØØØØØð
ð 
ˆŒõ" !Ô(Ñ*Ô*ˆŒˆˆó    c                 óX   — | j                              ¦   «         D ]\  }}d| j         |<   Œd S )Nr   )r+   Úitems)r/   ÚkÚvs      r   Ú_resetzSageMakerCallbackHandler._reset<   s:   € Ø”L×&Ò&Ñ(Ô(ð 	 ð 	 ‰DˆAˆqØˆDŒL˜‰OˆOð	 ð 	 r1   Ú
serializedÚpromptsÚkwargsc           	      óÞ  — | j         dxx         dz  cc<   | j         dxx         dz  cc<   | j         dxx         dz  cc<   | j         d         }i }|                     ddi¦  «         |                     t          |¦  «        ¦  «         |                     | j         ¦  «         t          |¦  «        D ];\  }}t	          |¦  «        }||d<   |                      || j        d|› d	|› �¦  «         Œ<d
S )zRun when LLM starts.r   é   r#   r   ÚactionÚon_llm_startÚpromptÚ
llm_start_Ú_prompt_N)r+   Úupdater   Ú	enumerater   Újsonfr.   )	r/   r7   r8   r9   r#   ÚrespÚidxr>   Úprompt_resps	            r   r=   z%SageMakerCallbackHandler.on_llm_start@   s$  € ð 	Œ�VÐÐÔ Ñ!ÐÐÑØŒ�\Ð"Ð"Ô" aÑ'Ð"Ð"Ñ"ØŒ�XÐÐÔ !Ñ#ÐÐÑà”\ ,Ô/ˆ
à!ˆØ�Š�X˜~Ð.Ñ/Ô/Ð/Ø�Š•L Ñ,Ô,Ñ-Ô-Ð-Ø�Š�D”LÑ!Ô!Ð!å$ WÑ-Ô-ð 	ð 	‰KˆC�Ý" 4™.œ.ˆKØ$*ˆK˜Ñ!Ø�JŠJØØ”Ø6˜ZÐ6Ð6°Ð6Ð6ñô ð ð ð	ð 	r1   Útokenc                 ó  — | j         dxx         dz  cc<   | j         dxx         dz  cc<   | j         d         }i }|                     d|dœ¦  «         |                     | j         ¦  «         |                      || j        d|› �¦  «         dS )z#Run when LLM generates a new token.r   r;   r%   Úon_llm_new_token)r<   rG   Úllm_new_tokens_N©r+   rA   rC   r.   )r/   rG   r9   r%   rD   s        r   rI   z)SageMakerCallbackHandler.on_llm_new_tokenX   s¥   € àŒ�VÐÐÔ Ñ!ÐÐÑØŒ�]Ð#Ð#Ô# qÑ(Ð#Ð#Ñ#à”l =Ô1ˆà!ˆØ�ŠÐ1¸EÐBÐBÑCÔCÐCØ�Š�D”LÑ!Ô!Ð!à�
Š
�4˜œÐ(G¸+Ð(GÐ(GÑHÔHÐHÐHÐHr1   Úresponsec           
      ó^  — | j         dxx         dz  cc<   | j         dxx         dz  cc<   | j         dxx         dz  cc<   | j         d         }i }|                     ddi¦  «         |                     t          |j        pi ¦  «        ¦  «         |                     | j         ¦  «         |j        D ]|}t          |¦  «        D ]j\  }}t          |¦  «        }|                     t          |                     ¦   «         ¦  «        ¦  «         |                      || j	        d|› d|› �¦  «         ŒkŒ}d	S )
zRun when LLM ends running.r   r;   r$   r   r<   Ú
on_llm_endÚllm_end_Ú_generation_N)
r+   rA   r   Ú
llm_outputÚgenerationsrB   r   ÚdictrC   r.   )	r/   rL   r9   r$   rD   rR   rE   Ú
generationÚgeneration_resps	            r   rN   z#SageMakerCallbackHandler.on_llm_ende   s]  € àŒ�VÐÐÔ Ñ!ÐÐÑØŒ�ZÐ Ð Ô  AÑ%Ð Ð Ñ ØŒ�VÐÐÔ Ñ!ÐÐÑà”< 
Ô+ˆà!ˆØ�Š�X˜|Ð,Ñ-Ô-Ð-Ø�Š•L Ô!4Ð!:¸Ñ;Ô;Ñ<Ô<Ð<à�Š�D”LÑ!Ô!Ð!à#Ô/ð 		ð 		ˆKÝ#,¨[Ñ#9Ô#9ð ð ‘��ZÝ"*¨4¡.¤.�Ø×&Ò&¥|°J·O²OÑ4EÔ4EÑ'FÔ'FÑGÔGÐGà—
’
ØØ”MØ:˜xÐ:Ð:°SÐ:Ð:ñô ð ð ð	ð		ð 		r1   Úerrorc                 óZ   — | j         dxx         dz  cc<   | j         dxx         dz  cc<   dS )zRun when LLM errors.r   r;   r   N©r+   ©r/   rV   r9   s      r   Úon_llm_errorz%SageMakerCallbackHandler.on_llm_error~   óD   € àŒ�VÐÐÔ Ñ!ÐÐÑØŒ�XÐÐÔ !Ñ#ÐÐÑÐÐr1   Úinputsc                 ó  — | j         dxx         dz  cc<   | j         dxx         dz  cc<   | j         dxx         dz  cc<   | j         d         }i }|                     ddi¦  «         |                     t          |¦  «        ¦  «         |                     | j         ¦  «         d                     d„ |                     ¦   «         D ¦   «         ¦  «        }t          |¦  «        }||d	<   |                      || j        d
|› �¦  «         dS )zRun when chain starts running.r   r;   r!   r   r<   Úon_chain_startú,c                 ó"   — g | ]\  }}|› d |› �‘ŒS ©ú=© ©Ú.0r4   r5   s      r   ú
<listcomp>z;SageMakerCallbackHandler.on_chain_start.<locals>.<listcomp>’   s&   € ÐFÐFÐF©t¨q°! 1 
 
 q 
 
ÐFÐFÐFr1   r\   Úchain_start_N)r+   rA   r   Újoinr3   r   rC   r.   )r/   r7   r\   r9   r!   rD   Úchain_inputÚ
input_resps           r   r^   z'SageMakerCallbackHandler.on_chain_startƒ   s  € ð 	Œ�VÐÐÔ Ñ!ÐÐÑØŒ�^Ð$Ð$Ô$¨Ñ)Ð$Ð$Ñ$ØŒ�XÐÐÔ !Ñ#ÐÐÑà”| NÔ3ˆà!ˆØ�Š�XÐ/Ð0Ñ1Ô1Ð1Ø�Š•L Ñ,Ô,Ñ-Ô-Ð-Ø�Š�D”LÑ!Ô!Ð!à—h’hÐFÐF°v·|²|±~´~ÐFÑFÔFÑGÔGˆÝ˜d‘^”^ˆ
Ø*ˆ
�8Ñà�
Š
�:˜tœ}Ð.K¸\Ð.KÐ.KÑLÔLÐLÐLÐLr1   Úoutputsc                 ó¦  — | j         dxx         dz  cc<   | j         dxx         dz  cc<   | j         dxx         dz  cc<   | j         d         }i }d                     d„ |                     ¦   «         D ¦   «         ¦  «        }|                     d|dœ¦  «         |                     | j         ¦  «         |                      || j        d	|› �¦  «         d
S )zRun when chain ends running.r   r;   r"   r   r_   c                 ó"   — g | ]\  }}|› d |› �‘ŒS ra   rc   rd   s      r   rf   z9SageMakerCallbackHandler.on_chain_end.<locals>.<listcomp>¡   s&   € Ð HÐ HÐ H±°°1 A  ¨  Ð HÐ HÐ Hr1   Úon_chain_end)r<   rk   Ú
chain_end_N)r+   rh   r3   rA   rC   r.   )r/   rk   r9   r"   rD   Úchain_outputs         r   rn   z%SageMakerCallbackHandler.on_chain_end˜   sê   € àŒ�VÐÐÔ Ñ!ÐÐÑØŒ�\Ð"Ð"Ô" aÑ'Ð"Ð"Ñ"ØŒ�VÐÐÔ Ñ!ÐÐÑà”\ ,Ô/ˆ
à!ˆØ—x’xÐ HÐ H¸¿º¹¼Ð HÑ HÔ HÑIÔIˆØ�Š˜~¸,ÐGÐGÑHÔHÐHØ�Š�D”LÑ!Ô!Ð!à�
Š
�4˜œÐ(A°ZÐ(AÐ(AÑBÔBÐBÐBÐBr1   c                 óZ   — | j         dxx         dz  cc<   | j         dxx         dz  cc<   dS )zRun when chain errors.r   r;   r   NrX   rY   s      r   Úon_chain_errorz'SageMakerCallbackHandler.on_chain_error§   r[   r1   Ú	input_strc                 óˆ  — | j         dxx         dz  cc<   | j         dxx         dz  cc<   | j         dxx         dz  cc<   | j         d         }i }|                     d|dœ¦  «         |                     t          |¦  «        ¦  «         |                     | j         ¦  «         |                      || j        d|› �¦  «         dS )	zRun when tool starts running.r   r;   r&   r   Úon_tool_start)r<   rs   Útool_start_N)r+   rA   r   rC   r.   )r/   r7   rs   r9   r&   rD   s         r   ru   z&SageMakerCallbackHandler.on_tool_start¬   sÞ   € ð 	Œ�VÐÐÔ Ñ!ÐÐÑØŒ�]Ð#Ð#Ô# qÑ(Ð#Ð#Ñ#ØŒ�XÐÐÔ !Ñ#ÐÐÑà”l =Ô1ˆà!ˆØ�Š˜¸YÐGÐGÑHÔHÐHØ�Š•L Ñ,Ô,Ñ-Ô-Ð-Ø�Š�D”LÑ!Ô!Ð!à�
Š
�4˜œÐ(C°kÐ(CÐ(CÑDÔDÐDÐDÐDr1   Úoutputc                 ób  — t          |¦  «        }| j        dxx         dz  cc<   | j        dxx         dz  cc<   | j        dxx         dz  cc<   | j        d         }i }|                     d|dœ¦  «         |                     | j        ¦  «         |                      || j        d|› �¦  «         dS )	zRun when tool ends running.r   r;   r'   r   Úon_tool_end)r<   rw   Ú	tool_end_N)Ústrr+   rA   rC   r.   )r/   rw   r9   r'   rD   s        r   ry   z$SageMakerCallbackHandler.on_tool_end½   sÍ   € å�V‘”ˆØŒ�VÐÐÔ Ñ!ÐÐÑØŒ�[Ð!Ð!Ô! QÑ&Ð!Ð!Ñ!ØŒ�VÐÐÔ Ñ!ÐÐÑà”L Ô-ˆ	à!ˆØ�Š˜}¸Ð?Ð?Ñ@Ô@Ð@Ø�Š�D”LÑ!Ô!Ð!à�
Š
�4˜œÐ(?°IÐ(?Ð(?Ñ@Ô@Ð@Ð@Ð@r1   c                 óZ   — | j         dxx         dz  cc<   | j         dxx         dz  cc<   dS )zRun when tool errors.r   r;   r   NrX   rY   s      r   Úon_tool_errorz&SageMakerCallbackHandler.on_tool_errorÌ   r[   r1   Útextc                 ó  — | j         dxx         dz  cc<   | j         dxx         dz  cc<   | j         d         }i }|                     d|dœ¦  «         |                     | j         ¦  «         |                      || j        d|› �¦  «         dS )z+
        Run when agent is ending.
        r   r;   r    Úon_text)r<   r~   Úon_text_NrK   )r/   r~   r9   r    rD   s        r   r€   z SageMakerCallbackHandler.on_textÑ   s¦   € ð 	Œ�VÐÐÔ Ñ!ÐÐÑØŒ�ZÐ Ð Ô  AÑ%Ð Ð Ñ à”< 
Ô+ˆà!ˆØ�Š˜y°$Ð7Ð7Ñ8Ô8Ð8Ø�Š�D”LÑ!Ô!Ð!à�
Š
�4˜œÐ(=°8Ð(=Ð(=Ñ>Ô>Ð>Ð>Ð>r1   Úfinishc                 óf  — | j         dxx         dz  cc<   | j         dxx         dz  cc<   | j         dxx         dz  cc<   | j         d         }i }|                     d|j        d         |j        dœ¦  «         |                     | j         ¦  «         |                      || j        d|› �¦  «         d	S )
zRun when agent ends running.r   r;   r(   r   Úon_agent_finishrw   )r<   rw   ÚlogÚagent_finish_N)r+   rA   Úreturn_valuesr…   rC   r.   )r/   r‚   r9   r(   rD   s        r   r„   z(SageMakerCallbackHandler.on_agent_finishà   sÝ   € àŒ�VÐÐÔ Ñ!ÐÐÑØŒ�\Ð"Ð"Ô" aÑ'Ð"Ð"Ñ"ØŒ�VÐÐÔ Ñ!ÐÐÑà”\ ,Ô/ˆ
Ø!ˆØ�Šà+Ø Ô.¨xÔ8Ø”zðð ñ	
ô 	
ð 	
ð 	�Š�D”LÑ!Ô!Ð!à�
Š
�4˜œÐ(D¸
Ð(DÐ(DÑEÔEÐEÐEÐEr1   r<   c                 óf  — | j         dxx         dz  cc<   | j         dxx         dz  cc<   | j         dxx         dz  cc<   | j         d         }i }|                     d|j        |j        |j        dœ¦  «         |                     | j         ¦  «         |                      || j        d|› �¦  «         dS )	zRun on agent action.r   r;   r&   r   Úon_agent_action)r<   ÚtoolÚ
tool_inputr…   Úagent_action_N)r+   rA   rŠ   r‹   r…   rC   r.   )r/   r<   r9   r&   rD   s        r   r‰   z(SageMakerCallbackHandler.on_agent_actionó   sÝ   € àŒ�VÐÐÔ Ñ!ÐÐÑØŒ�]Ð#Ð#Ô# qÑ(Ð#Ð#Ñ#ØŒ�XÐÐÔ !Ñ#ÐÐÑà”l =Ô1ˆØ!ˆØ�Šà+ØœØ$Ô/Ø”zð	ð ñ	
ô 	
ð 	
ð 	�Š�D”LÑ!Ô!Ð!Ø�
Š
�4˜œÐ(E¸Ð(EÐ(EÑFÔFÐFÐFÐFr1   Tr   Údata_dirÚfilenameÚ	is_outputc                 ó¦   — t           j                             ||› d�¦  «        }t          ||¦  «         | j                             |||¬¦  «         dS )z,To log the input data as json file artifact.z.json)Únamer�   N)ÚosÚpathrh   r   r   Úlog_file)r/   r   r�   rŽ   r�   r   s         r   rC   zSageMakerCallbackHandler.jsonf  sU   € õ ”G—L’L ¨hÐ+=Ð+=Ð+=Ñ>Ô>ˆ	Ý�$˜	Ñ"Ô"Ð"ØŒ×Ò˜)¨(¸iÐÑHÔHÐHÐHÐHr1   c                 ó`   — |                       ¦   «          t          j        | j        ¦  «         dS )z9Reset the steps and delete the temporary local directory.N)r6   ÚshutilÚrmtreer.   )r/   s    r   Úflush_trackerz&SageMakerCallbackHandler.flush_tracker  s'   € à�Š‰ŒˆÝŒ�d”mÑ$Ô$Ð$Ð$Ð$r1   )r   N)T) Ú__name__Ú
__module__Ú__qualname__Ú__doc__r   r*   r6   r   r{   r   r=   rI   r   rN   ÚBaseExceptionrZ   r^   rn   rr   ru   ry   r}   r€   r	   r„   r   r‰   r   ÚboolrC   r˜   Ú__classcell__)r0   s   @r   r   r      sv  ø€ € € € € ðð ð+˜Cð + Dð +ð +ð +ð +ð +ð +ð2 ð  ð  ð  ðØ˜s C˜xœ.ðØ37¸´9ðØHKðà	ðð ð ð ð0I cð I°Sð I¸Tð Ið Ið Ið Ið 9ð ¸ð Àð ð ð ð ð2$ -ð $¸3ð $À4ð $ð $ð $ð $ð
MØ˜s C˜xœ.ðMØ26°s¸C°x´.ðMØLOðMà	ðMð Mð Mð Mð*C D¨¨c¨¤Nð C¸cð CÀdð Cð Cð Cð Cð$ Mð $¸Sð $ÀTð $ð $ð $ð $ð
EØ˜s C˜xœ.ðEØ58ðEØDGðEà	ðEð Eð Eð Eð"A #ð A°ð A¸ð Að Að Að Að$ =ð $¸Cð $ÀDð $ð $ð $ð $ð
?˜Cð ?¨3ð ?°4ð ?ð ?ð ?ð ?ðF kð F¸Sð FÀTð Fð Fð Fð Fð&G kð G¸Sð GÀSð Gð Gð Gð Gð0 %)ð
Ið 
Ià�3˜�8Œnð
Ið ð
Ið ð	
Ið
 ˜D”>ð
Ið 
ð
Ið 
Ið 
Ið 
Ið%ð %ð %ð %ð %ð %ð %ð %r1   r   )r   r’   r–   r,   Úcopyr   Útypingr   r   r   r   Úlangchain_core.agentsr   r	   Úlangchain_core.callbacksr
   Úlangchain_core.outputsr   Ú#langchain_community.callbacks.utilsr   rS   r{   r   r   rc   r1   r   ú<module>r¦      s!  ðØ €€€Ø 	€	€	€	Ø €€€Ø €€€Ø Ð Ð Ð Ð Ð Ø ,Ð ,Ð ,Ð ,Ð ,Ð ,Ð ,Ð ,Ð ,Ð ,Ð ,Ð ,à :Ð :Ð :Ð :Ð :Ð :Ð :Ð :Ø 8Ð 8Ð 8Ð 8Ð 8Ð 8Ø ,Ð ,Ð ,Ð ,Ð ,Ð ,ðð ð ð ð ð ð
!�Dð ! Sð !¨Tð !ð !ð !ð !ðy%ð y%ð y%ð y%ð y%Ð2ñ y%ô y%ð y%ð y%ð y%r1   