Ë
    ümxi�G  ã                  óÜ  — d Z ddlmZ ddlZddlZd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 ddlmZ ddlmZ dd	lmZ dd
lmZ ddlmZ eej.                  ej.                  gef   Z	 	 	 	 	 	 dd„Z	 	 	 	 dd„Zej6                   G d„ d«      «       Z	 	 	 	 	 	 dd„Z G d„ deeej<                  df   gej>                  f   ej@                  ¬«      Z! G d„ de!«      Z" G d„ de!«      Z#y)zLLMFunction.é    )ÚannotationsN)ÚAbstractSetÚAnyÚCallableÚIterableÚMappingÚOptionalÚSequenceÚUnion)Úllmfn_input_utils)Úllmfn_output_row)Úllmfn_outputs)Úllmfn_post_process)Úllmfn_post_process_cmds)Úmodel)Úprompt_utilsc                óD   — | j                  «       |j                  «       k(  S )z-Default function used when comparing outputs.)Úresult_value)ÚlhsÚrhss     úd/home/htdocs/ttos/venv/lib/python3.12/site-packages/google/generativeai/notebook/lib/llm_function.pyÚ_is_equal_fnr   4   s!   € ð
 ×ÑÓ ×!1Ñ!1Ó!3Ñ3Ð3ó    c                ó   ‡ — 	 	 	 	 dˆ fd„}|S )z1Vectorize a single-row-based comparison function.c                óH   •— | D ��cg c]  \  }} ‰||«      ‘Œ c}}S c c}}w ©N© )Úlhs_and_rhs_rowsr   r   Úfns      €r   Ú_fnz0_convert_compare_fn_to_batch_add_fn.<locals>._fnG   s#   ø€ ð .>×>¡  c‘�3˜•Ó>Ð>ùÓ>s   ‡)r   zYSequence[tuple[llmfn_output_row.LLMFnOutputRowView, llmfn_output_row.LLMFnOutputRowView]]ÚreturnzSequence[Any]r   )r   r    s   ` r   Ú#_convert_compare_fn_to_batch_add_fnr"   <   s#   ø€ ð?ð
ð?ð 
õ?ð €Jr   c                  ó@   — e Zd ZU ded<   ded<   ded<   ded<   ded<   y	)
Ú_PromptInfoÚintÚ
prompt_numÚstrÚpromptÚ	input_numzMapping[str, str]Úprompt_varsÚmodel_inputN)Ú__name__Ú
__module__Ú__qualname__Ú__annotations__r   r   r   r$   r$   T   s   … àƒOØƒKØƒNØ"Ó"ØÔr   r$   c           	   #  óÜ   K  — g }|�t        j                  |«      }|si g}t        | «      D ]<  \  }}t        |«      D ])  \  }} |j                  di |¤Ž}t	        |||||¬«      –— Œ+ Œ> y­w)a  Generate a tuple of fields needed for processing prompts.

    Args:
      prompts: A list of prompts, with optional keyword placeholders.
      inputs: A list of key/value pairs to substitute into placeholders in
        `prompts`.

    Yields:
      A _PromptInfo instance.
    N)r&   r(   r)   r*   r+   r   )r   Úto_normalized_inputsÚ	enumerateÚformatr$   )ÚpromptsÚinputsÚnormalized_inputsr&   r(   r)   r*   r+   s           r   Ú_generate_promptsr7   ]   s”   è ø€ ð 68ÐØÐÜ-×BÑBÀ6ÓJÐñ Ø˜DÐä'¨Ó0ò 
Ñˆ
�FÜ&/Ð0AÓ&Bò 		Ñ"ˆI�{à'˜&Ÿ-™-Ñ6¨+Ñ6ˆKÜØ%ØØ#Ø'Ø'ôó ñ		ñ
ùs   ‚A*A,c                  óÀ   — e Zd ZdZ	 d
	 dd„Zej                  dd„«       Zej                  	 	 	 	 dd„«       Z	 d
	 	 	 dd„Z		 	 	 	 	 	 dd„Z
	 	 	 	 	 	 dd„Z	 	 	 	 	 	 dd	„Zy)ÚLLMFunctionz6Base class for LLMFunctionImpl and LLMCompareFunction.Nc                ó    — g | _         || _        y)a  Constructor.

        Args:
          outputs_ipython_display_fn: Optional function that will be used to
            override how the outputs of this LLMFunction will be displayed in a
            notebook (See further documentation in LLMFnOutputs.__init__().)
        N)Ú_post_process_cmdsÚ_outputs_ipython_display_fn)ÚselfÚoutputs_ipython_display_fns     r   Ú__init__zLLMFunction.__init__ˆ   s   € ð Z\ˆÔØ+EˆÕ(r   c                 ó   — y)zLReturns the placeholders that should be present in inputs for this function.Nr   ©r=   s    r   Úget_placeholderszLLMFunction.get_placeholders–   ó   � r   c                 ó   — y)z&Concrete implementation of __call__().Nr   )r=   r5   s     r   Ú
_call_implzLLMFunction._call_implš   rC   r   c                óf   — | j                  |«      }t        j                  || j                  ¬«      S )z+Runs and returns results based on `inputs`.)ÚoutputsÚipython_display_fn)rE   r   ÚLLMFnOutputsr<   )r=   r5   rG   s      r   Ú__call__zLLMFunction.__call__    s0   € ð —/‘/ &Ó)ˆä×)Ñ)Ø°×0PÑ0Pô
ð 	
r   c                óf   — | j                   j                  t        j                  ||¬«      «       | S ©N©Únamer   )r;   Úappendr   ÚLLMFnPostProcessReorderCommand©r=   rN   r   s      r   Úadd_post_process_reorder_fnz'LLMFunction.add_post_process_reorder_fnª   s0   € ð 	×Ñ×&Ñ&Ü#×BÑBÈÐQSÔTô	
ð ˆr   c                óf   — | j                   j                  t        j                  ||¬«      «       | S rL   )r;   rO   r   ÚLLMFnPostProcessAddCommandrQ   s      r   Úadd_post_process_add_fnz#LLMFunction.add_post_process_add_fn²   s/   € ð
 	×Ñ×&Ñ&Ü#×>Ñ>ÀDÈRÔPô	
ð ˆr   c                óf   — | j                   j                  t        j                  ||¬«      «       | S rL   )r;   rO   r   ÚLLMFnPostProcessReplaceCommandrQ   s      r   Úadd_post_process_replace_fnz'LLMFunction.add_post_process_replace_fn¼   s0   € ð
 	×Ñ×&Ñ&Ü#×BÑBÈÐQSÔTô	
ð ˆr   r   )r>   ú3Callable[[llmfn_outputs.LLMFnOutputs], None] | None©r!   zAbstractSet[str]©r5   ú*llmfn_input_utils.LLMFunctionInputs | Noner!   z(Sequence[llmfn_outputs.LLMFnOutputEntry])r5   r\   r!   zllmfn_outputs.LLMFnOutputs)rN   r'   r   z1llmfn_post_process.LLMFnPostProcessBatchReorderFnr!   r9   )rN   r'   r   z-llmfn_post_process.LLMFnPostProcessBatchAddFnr!   r9   )rN   r'   r   z1llmfn_post_process.LLMFnPostProcessBatchReplaceFnr!   r9   )r,   r-   r.   Ú__doc__r?   ÚabcÚabstractmethodrB   rE   rJ   rR   rU   rX   r   r   r   r9   r9      sß   „ ñ Að [_ðFà$WóFð 	×Ñò[ó ð[ð 	×Ñð5Ø@ð5à	1ò5ó ð5ð DHð
Ø@ð
à	#ó
ðØðØNðà	óðàðð :ðð 
ó	ðàðð >ðð 
ô	r   r9   )Ú	metaclassc                  ó`   ‡ — e Zd ZdZ	 	 d	 	 	 	 	 	 	 dˆ fd„Z	 	 	 	 dd„Zd	d„Z	 	 	 	 d
d„Zˆ xZS )ÚLLMFunctionImpla»  Callable class that executes the contents of a Magics cell.

    An LLMFunction is constructed from the Magics command line and cell contents
    specified by the user. It is defined by:
    - A model instance,
    - Model arguments
    - A prompt template (e.g. "the opposite of hot is {word}") with an optional
      keyword placeholder.

    The LLMFunction takes as its input a sequence of dictionaries containing
    values for keyword replacement, e.g. [{"word": "hot"}, {"word": "tall"}].

    This will cause the model to be executed with the following prompts:
      "The opposite of hot is"
      "The opposite of tall is"

    The results will be returned in a LLMFnOutputs instance.
    c                ó$  •— t         ‰| �  |¬«       || _        || _        |€t	        j
                  «       n|| _        t        i «      | _        | j                  D ]5  }| j                  j                  t        j                  |«      «      | _        Œ7 y)aÌ  Constructor.

        Args:
          model: The model that the prompts will execute on.
          prompts: A sequence of prompt templates with optional placeholders. The
            placeholders will be replaced by the inputs passed into this function.
          model_args: Optional set of model arguments to configure how the model
            executes the prompts.
          outputs_ipython_display_fn: See documentation in LLMFunction.__init__().
        ©r>   N)Úsuperr?   Ú_modelÚ_promptsÚ	model_libÚModelArgumentsÚ_model_argsÚ	frozensetÚ_placeholdersÚunionr   rB   )r=   r   r4   Ú
model_argsr>   r(   Ú	__class__s         €r   r?   zLLMFunctionImpl.__init__Û   s…   ø€ ô" 	‰ÑÐ4NÐÔOØˆŒØˆŒØ9CÐ9Kœ9×3Ñ3Ô5ÐQ[ˆÔô ' r›]ˆÔØ—m‘mò 	aˆFØ!%×!3Ñ!3×!9Ñ!9¼,×:WÑ:WÐX^Ó:_Ó!`ˆDÕñ	ar   c           
     ó²  — | j                   D ]\  }	 t        |t        j                  «      r|j	                  |«      }n-t        j                  dj                  t        |«      «      «      ‚Œ^ |S # t
        j                  $ r ‚ t        $ rL}t        j                  dj                  |j                  «       t        |«      j                  |«      «      ‚d}~ww xY w©z-Runs post-processing commands over `results`.z)Unsupported post-process command type: {}z Error executing "{}", got {}: {}N)r;   Ú
isinstancer   ÚLLMFnImplPostProcessCommandÚrunr   ÚPostProcessExecutionErrorr3   ÚtypeÚRuntimeErrorrN   r,   )r=   ÚresultsÚcmdÚes       r   Ú_run_post_processing_cmdsz)LLMFunctionImpl._run_post_processing_cmdsö   sÈ   € ð ×*Ñ*ò 	ˆCðÜ˜cÔ#:×#VÑ#VÔWØ!Ÿg™g gÓ.‘Gä,×FÑFØC×JÑJÌ4ÐPSË9ÓUóð ð ð	ð ˆøô &×?Ñ?ò ØÜò Ü(×BÑBØ6×=Ñ=¸c¿h¹h»jÌ$ÈqË'×JZÑJZÐ\]Ó^óð ûðús   ‘AA.Á.CÂ
ACÃCc                ó   — | j                   S r   ©rl   rA   s    r   rB   z LLMFunctionImpl.get_placeholders
  ó   € Ø×!Ñ!Ð!r   c                óz  — g }t        | j                  |¬«      D �]  }| j                  j                  |j                  | j
                  ¬«      }g }t        |j                  «      D ]c  \  }}|j                  t        j                  t        j                  j                  |t        j                  j                  |it        ¬«      «       Œe |j                  t        j                   |j"                  |j$                  |j&                  |j(                  |j                  || j+                  |«      ¬«      «       �Œ  |S )N)r4   r5   )r+   rn   ©ÚdataÚresult_type)r&   r)   r(   r*   r+   Úmodel_resultsÚoutput_rows)r7   rg   rf   Ú
call_modelr+   rj   r2   Útext_resultsrO   r   ÚLLMFnOutputRowr   ÚColumnNamesÚ
RESULT_NUMÚTEXT_RESULTr'   ÚLLMFnOutputEntryr&   r)   r(   r*   r{   )r=   r5   rx   Úinforƒ   r„   Ú
result_numÚtext_results           r   rE   zLLMFunctionImpl._call_impl  s  € ð 9;ˆÜ%¨d¯m©mÀFÔKó 	ˆDØ ŸK™K×2Ñ2Ø ×,Ñ,¸×9IÑ9Ið 3ó ˆMð BDˆKÜ+4°]×5OÑ5OÓ+Pò 	Ñ'�
˜KØ×"Ñ"Ü$×3Ñ3ä)×5Ñ5×@Ñ@À*Ü)×5Ñ5×AÑAÀ;ðô %(ôõð	ð �N‰NÜ×.Ñ.Ø#Ÿ™Ø"Ÿn™nØŸ;™;Ø $× 0Ñ 0Ø $× 0Ñ 0Ø"/Ø $× >Ñ >¸{Ó Kôö
ð	ð4 ˆr   ©NN)r   zmodel_lib.AbstractModelr4   úSequence[str]rn   zmodel_lib.ModelArguments | Noner>   rY   )rx   ú)Sequence[llmfn_output_row.LLMFnOutputRow]r!   r‘   rZ   r[   )	r,   r-   r.   r]   r?   r{   rB   rE   Ú__classcell__©ro   s   @r   rb   rb   Ç   sp   ø„ ñð. 7;ØZ^ðaà&ðað ðað 4ð	að
 %Xõað6Ø@ðà	2óó("ðØ@ðà	1÷r   rb   c                  ó|   ‡ — e Zd ZdZ	 	 d	 	 	 	 	 	 	 dˆ fd„Z	 	 	 	 	 	 	 	 d	d„Zd
d„Z	 	 	 	 dd„Z	 	 	 	 	 	 dd„Zˆ xZ	S )ÚLLMCompareFunctionz€LLMFunction for comparisons.

    LLMCompareFunction runs an input over a pair of LLMFunctions and compares the
    result.
    c                ó¼  •— t         ‰| �  |¬«       |d   | _        |d   | _        |d   | _        |d   | _        t        | j                  j                  «       «      j                  | j
                  j                  «       «      | _	        |sd| _
        t        | _        y|d   \  }}|| _
        || _        |dd D ]!  \  }}| j                  |t        |«      ¬«       Œ# y)aÄ  Constructor.

        Args:
          lhs_name_and_fn: Name and function for the left-hand side of the
            comparison.
          rhs_name_and_fn: Name and function for the right-hand side of the
            comparison.
          compare_name_and_fns: Optional names and functions for comparing the
            results of the left- and right-hand sides.
          outputs_ipython_display_fn: See documentation in LLMFunction.__init__().
        rd   r   é   Úis_equaléÿÿÿÿNrM   )re   r?   Ú	_lhs_nameÚ_lhs_fnÚ	_rhs_nameÚ_rhs_fnrk   rB   rm   rl   Ú_result_namer   Ú_result_compare_fnÚadd_compare_post_process_add_fnr"   )	r=   Úlhs_name_and_fnÚrhs_name_and_fnÚcompare_name_and_fnsr>   rN   r   Úcmp_fnro   s	           €r   r?   zLLMCompareFunction.__init__5  sê   ø€ ô$ 	‰ÑÐ4NÐÔOØ-¨aÑ0ˆŒØ$3°AÑ$6ˆŒØ-¨aÑ0ˆŒØ$3°AÑ$6ˆŒÜ& t§|¡|×'DÑ'DÓ'FÓG×MÑMØ�L‰L×)Ñ)Ó+ó
ˆÔñ $Ø *ˆDÔÜ&2ˆDÕ#ð ,¨BÑ/‰HˆD�"Ø $ˆDÔØ&(ˆDÔ#ð !5°S°bÐ 9ò ‘��fØ×4Ñ4ØÔ"EÀfÓ"Mð 5õ ñr   c           
     ó  — | j                   D ]’  }	 t        |t        j                  «      r|j	                  |«      }nct        |t        j
                  «      r&|j	                  t        t        |||«      «      «      }n#t        dj                  t        |«      «      «      ‚Œ” |S # t        j                  $ r ‚ t        $ rL}t        j                  dj                  |j                  «       t        |«      j                  |«      «      ‚d}~ww xY wrq   )r;   rr   r   rs   rt   ÚLLMCompareFnPostProcessCommandÚlistÚziprw   r3   rv   r   ru   rN   r,   )r=   Úlhs_output_rowsÚrhs_output_rowsrx   ry   rz   s         r   r{   z,LLMCompareFunction._run_post_processing_cmds`  sì   € ð ×*Ñ*ò 	ˆCðÜ˜cÔ#:×#VÑ#VÔWØ!Ÿg™g gÓ.‘GÜ Ô%<×%[Ñ%[Ô\Ø!Ÿg™g¤d¬3¨ÀÐQXÓ+YÓ&ZÓ[‘Gä&ØC×JÑJÌ4ÐPSË9ÓUóð øð	ð  ˆøô &×?Ñ?ò ØÜò Ü(×BÑBØ6×=Ñ=¸c¿h¹h»jÌ$ÈqË'×JZÑJZÐ\]Ó^óð ûðús   ‘BB$Â$DÃ ADÄDc                ó   — | j                   S r   r}   rA   s    r   rB   z#LLMCompareFunction.get_placeholdersy  r~   r   c                óî  — | j                  |«      }| j                  |«      }g }t        ||«      D �]À  \  }}|j                  |j                  k7  r/t	        dj                  |j                  |j                  «      «      ‚|j                  |j                  k7  r/t	        dj                  |j                  |j                  «      «      ‚|j                  |j                  k7  r/t	        dj                  |j                  |j                  «      «      ‚t        t        |j                  «      t        |j                  «      «      }|j                  d | }|j                  d | }	g }
t        t        ||	«      «      D �]  \  }}|\  }}t        j                  j                  || j                  | j!                  ||«      i}t#        j$                  |t&        ¬«      }|j)                  t        j                  j*                  |j                  «       | j,                  |f| j.                  |ffD ]]  \  }}|j1                  «       D ]E  \  }}|t        j                  j                  k7  sŒ$|j)                  dj                  ||«      |«       ŒG Œ_ |
j3                  |«       �Œ" |j3                  t        j4                  |j                  |j                  |j                  | j7                  ||	|
¬«      ¬«      «       �ŒÃ |S )NzPrompt num mismatch: {} vs {}zInput num mismatch: {} vs {}zPrompt vars mismatch: {} vs {}r€   z{}_{})r©   rª   rx   )r&   r)   r*   r„   )r›   r�   r¨   r&   rw   r3   r)   r*   ÚminÚlenr„   r2   r   rˆ   r‰   rž   rŸ   r   r‡   r   ÚaddÚPROMPT_VARSrš   rœ   ÚitemsrO   r‹   r{   )r=   r5   Úlhs_resultsÚrhs_resultsrG   Ú	lhs_entryÚ	rhs_entryÚnum_output_rowsr©   rª   r„   r�   Úlhs_and_rhs_output_rowÚlhs_output_rowÚrhs_output_rowÚrow_dataÚ
output_rowrN   ÚrowÚkÚvs                        r   rE   zLLMCompareFunction._call_impl|  sØ  € ð —l‘l 6Ó*ˆØ—l‘l 6Ó*ˆð 9;ˆÜ$'¨°[Ó$Aó E	Ñ ˆI�yØ×#Ñ# y×';Ñ';Ò;Ü"Ø3×:Ñ:Ø!×,Ñ,¨i×.BÑ.Bóóð ð
 ×"Ñ" i×&9Ñ&9Ò9Ü"Ø2×9Ñ9¸)×:MÑ:MÈy×ObÑObÓcóð ð ×$Ñ$¨	×(=Ñ(=Ò=Ü"Ø4×;Ñ;Ø!×-Ñ-¨y×/DÑ/Dóóð ô "¤# i×&;Ñ&;Ó"<¼cÀ)×BWÑBWÓ>XÓYˆOØ'×3Ñ3Ð4D°_ÐEˆOØ'×3Ñ3Ð4D°_ÐEˆOØACˆKÜ6?Ü�O _Ó5ó7ó /Ñ2�
Ð2ð 2HÑ.� ô "×-Ñ-×8Ñ8¸*Ø×%Ñ% t×'>Ñ'>¸~È~Ó'^ð,�ô .×<Ñ<À(ÔX[Ô\�
ð —‘œ}×8Ñ8×DÑDÀi×F[ÑF[Ô\ð —^‘^ ^Ð4Ø—^‘^ ^Ð4ð"ò G‘I�D˜#ð !$§	¡	£ò G™˜˜1Ø¤× 9Ñ 9× DÑ DÓDð 'ŸN™N¨7¯>©>¸$ÀÓ+BÀAÕFñ	Gð	Gð ×"Ñ" :Ö.ð?/ðB �N‰NÜ×.Ñ.Ø(×3Ñ3Ø'×1Ñ1Ø )× 5Ñ 5Ø $× >Ñ >Ø(7Ø(7Ø +ð !?ó !ô		öðuE	ðL ˆr   c                óf   — | j                   j                  t        j                  ||¬«      «       | S rL   )r;   rO   r   Ú!LLMCompareFnPostProcessAddCommandrQ   s      r   r    z2LLMCompareFunction.add_compare_post_process_add_fnÌ  s0   € ð
 	×Ñ×&Ñ&Ü#×EÑEÈ4ÐTVÔWô	
ð ˆr   r�   )r¡   útuple[str, LLMFunction]r¢   rÁ   r£   z&Sequence[tuple[str, CompareFn]] | Noner>   rY   )r©   r‘   rª   r‘   rx   r‘   r!   r‘   rZ   r[   )rN   r'   r   ú4llmfn_post_process.LLMCompareFnPostProcessBatchAddFnr!   r9   )
r,   r-   r.   r]   r?   r{   rB   rE   r    r’   r“   s   @r   r•   r•   .  s¬   ø„ ñð HLØZ^ð)à0ð)ð 1ð)ð Eð	)ð
 %Xõ)ðVàBðð Cðð ;ð	ð
 
3óó2"ðNØ@ðNà	1óNð`àðð Aðð 
÷	r   r•   )r   ú#llmfn_output_row.LLMFnOutputRowViewr   rÃ   r!   Úbool)r   zYCallable[[llmfn_output_row.LLMFnOutputRowView, llmfn_output_row.LLMFnOutputRowView], Any]r!   rÂ   )r4   r�   r5   r\   r!   zIterable[_PromptInfo])$r]   Ú
__future__r   r^   ÚdataclassesÚtypingr   r   r   r   r   r	   r
   r   Ú google.generativeai.notebook.libr   r   r   r   r   r   rh   r   ÚLLMFnOutputRowViewÚ	CompareFnr   r"   Ú	dataclassr$   r7   ÚLLMFunctionInputsrI   ÚABCMetar9   rb   r•   r   r   r   ú<module>rÎ      s7  ðñ Ý "ã 
Û ÷	÷ 	ó 	õ ?Ý =Ý :Ý ?Ý DÝ ?Ý 9ð Ø×(Ñ(Ð*:×*MÑ*MÐNØð	ñ€	ð4Ø	,ð4à	,ð4ð 
ó4ðð	ðð :óð0 ×Ñ÷ð ó ððØðØ$NðàóôDEØØ	Ð ×2Ñ2°DÐ8Ñ	9Ð:Ø×"Ñ"ð	$ñð �k‰kõEôPd�kô dôNf˜õ fr   