
    i                     0   d dl Z d dlmZmZmZmZ d dlmZ d dlm	Z	 d dl
mZ erd dlmZmZ dee         fdZdee         fd	ZddefdZ G d de          Zeee	         eeef         f         Z G d de          Zdddeeef         fdZdefdZdS )    N)TYPE_CHECKINGAnyOptionalUnion)ChatCompletion)	BaseModel)GenerateJsonSchema)ResultWithTextSearchResponsereturnc                     | j         d         j        j        rR| j         d         j        j        D ]:}|j        j        dk    r(t          j        |j        j                  d         }|c S ;dS )z+Extract query from completion if it exists.r   searchqueryN)choicesmessage
tool_callsfunctionnamejsonloads	arguments)
completion	tool_callr   s      S/home/ubuntu/.hermes/hermes-agent/venv/lib/python3.11/site-packages/exa_py/utils.pymaybe_get_queryr      su    !$/ #+A.6A 	 	I!&(22
9#5#?@@I 3 4    c                     | j         d         j        }|j        s
J d            d |D             }|                    |dd|j        d         j        |dg           |S )zaAdd assistant message and exa result to messages list. Also remove previous exa call and results.r   z&Must use this with a tool call requestc                 D    g | ]}|                     d           dk    |S )roler   )get).0r   s     r   
<listcomp>z+add_message_to_messages.<locals>.<listcomp>   s5       F0C0Cz0Q0Q0Q0Q0Qr   toolr   )r   r   tool_call_idcontent)r   r   r   extendid)r   messages
exa_resultassistant_messages       r   add_message_to_messagesr+      s    "*1-5'QQ)QQQ' '  H OO  1 <Q ? B%	 	

 
 
 Or   max_lenc                 T    fd| j         D             }d                    |          S )z(Format exa result for pasting into chat.c           
      V    g | ]%}d |j          d|j         d|j        d          d&S )zUrl: z
Title: 
N)urltitletext)r!   resultr-   s     r   r"   z%format_exa_result.<locals>.<listcomp>0   sT        	O
NNV\NNV['5JNNN  r   r0   )resultsjoin)r)   r-   strs    ` r   format_exa_resultr8   .   s@        (  C
 99S>>r   c                   f     e Zd ZdZ	 	 dded         f fdZeded         defd            Z xZ	S )	ExaOpenAICompletionz"Exa wrapper for OpenAI completion.Nr)   zSearchResponse[ResultWithText]c	           	      f    t                                          |||||||           || _        d S )Nr'   r   createdmodelobjectsystem_fingerprintusage)super__init__r)   )
selfr)   r'   r   r=   r>   r?   r@   rA   	__class__s
            r   rC   zExaOpenAICompletion.__init__;   sH     	1 	 	
 	
 	
 %r   r   c           
      n     | ||j         |j        |j        |j        |j        |j        |j                  S )N)r)   r'   r   r=   r>   r?   r@   rA   r<   )clsr)   r   s      r   from_completionz#ExaOpenAICompletion.from_completionQ   sI     s!}&&"$)<"	
 	
 	
 		
r   )NN)
__name__
__module____qualname____doc__r   rC   classmethodr   rH   __classcell__rE   s   @r   r:   r:   8   s        ,,  % %=>% % % % % %, 
=>
 #
 
 
 [
 
 
 
 
r   r:   c                   *     e Zd ZdZd fd	Zd Z xZS )InlineJsonSchemaGeneratorzXCustom JSON schema generator that inlines all schemas without creating $defs references.
validationc                     t                                          ||          }d|v r'|d         }|                     ||          }d|v r|d= |S |S )zDGenerate JSON schema normally, then post-process to inline all refs.z$defs)rB   generate_inline_refs)rD   schemamoder4   definitionsinlined_resultrE   s         r   rT   z"InlineJsonSchemaGenerator.generatei   sg     !!&$// f /K!..v{CCN.(("7+!!r   c                 ~    t          |t                    rd|v r\t          |          dk    rI|d         }|                    d          r*|dd         }|v r                     |                   S |S  fd|                                D             S t          |t                    r fd|D             S |S )z5Recursively replace all $ref with actual definitions.z$ref   z#/$defs/   Nc                 D    i | ]\  }}|                     |          S  rU   )r!   keyvaluerX   rD   s      r   
<dictcomp>z:InlineJsonSchemaGenerator._inline_refs.<locals>.<dictcomp>   s?       "U **5+>>  r   c                 <    g | ]}                     |          S r^   r_   )r!   itemrX   rD   s     r   r"   z:InlineJsonSchemaGenerator._inline_refs.<locals>.<listcomp>   s)    IIITD%%dK88IIIr   )
isinstancedictlen
startswithrU   itemslist)rD   objrX   ref_pathdef_names   ` `  r   rU   z&InlineJsonSchemaGenerator._inline_refsy   s    c4   	}}SQv;&&z22 U'|H;..#00X1FTTT
    &)iikk    T"" 	IIIIISIIII Jr   )rR   )rI   rJ   rK   rL   rT   rU   rN   rO   s   @r   rQ   rQ   f   sR        bb            r   rQ   schema_inputJSONSchemaInputc                     t          | t                    r2t          | t                    r|                     ddt
                    S t          | t                    r| S t          dt          |            d          )zConvert various schema input types to JSON Schema dict.

    Args:
        schema_input: Either a Pydantic BaseModel class or a dict containing JSON Schema

    Returns:
        dict: JSON Schema representation (fully inlined without $defs)
    Fserialization)by_aliasrW   schema_generatorzUnsupported schema type: z#. Expected BaseModel class or dict.)re   type
issubclassr   model_json_schemarQ   rf   
ValueError)rn   s    r   _convert_schema_inputrx      s     ,%% 
*\9*M*M 
-- 6 . 
 
 	

 
L$	'	' 
_\(:(:___
 
 	
r   c                  @   	 	 ddl } n# t          $ r ddl} Y nw xY wddl}|j                            |j                            t                              }|j                            |          }|j                            |d          }|j        	                    |          rgt          |d          5 }|                     |          }|                    di                               dd          cddd           S # 1 swxY w Y   n# t          $ r Y nw xY wdS )z,Get the package version from pyproject.toml.r   Nzpyproject.tomlrbprojectversionunknown)tomllibImportErrortomliospathdirnameabspath__file__r6   existsopenloadr    	Exception)r~   r   current_dirproject_rootpyproject_pathfdatas          r   _get_package_versionr      s   	$NNNN 	$ 	$ 	$######	$ 				goobgooh&?&?@@w{33l4DEE7>>.)) 	Ind++ Iq||Axx	2..229iHHI I I I I I I I I I I I I I I I     9sL    D D BD 6?D5D DD 	D
D 
DD)r,   )r   typingr   r   r   r   openai.types.chatr   pydanticr   pydantic.json_schemar	   
exa_py.apir
   r   r7   r   rj   rf   r+   intr8   r:   rt   ro   rQ   rx   r   r^   r   r   <module>r      s    6 6 6 6 6 6 6 6 6 6 6 6 , , , , , ,       3 3 3 3 3 3 :999999998C=    d    0 3    (
 (
 (
 (
 (
. (
 (
 (
V Yc3h78) ) ) ) ) 2 ) ) )X
(9 
d38n 
 
 
 
0c      r   