
    Ng                     ~    d Z ddlZddlmZmZmZm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Z G d de          ZdS )	zUtil that calls Tavily Search API.

In order to set this up, follow instructions at:
https://docs.tavily.com/docs/tavily-api/introduction
    N)AnyDictListOptionalget_from_dict_or_env)	BaseModel
ConfigDict	SecretStrmodel_validatorzhttps://api.tavily.comc                   Z   e Zd ZU dZeed<    ed          Z ed          e	de
defd	                        Zd
dg g dddfdedee         dee         deee                  deee                  dee         dee         dee         de
fdZd
dg g dddfdedee         dee         deee                  deee                  dee         dee         dee         dee
         fdZd
dg g dddfdedee         dee         deee                  deee                  dee         dee         dee         de
fdZd
dg g dddfdedee         dee         deee                  deee                  dee         dee         dee         dee
         fdZdee
         dee
         fdZdS )TavilySearchAPIWrapperzWrapper for Tavily Search API.tavily_api_keyforbid)extrabefore)modevaluesreturnc                 2    t          |dd          }||d<   |S )z9Validate that api key and endpoint exists in environment.r   TAVILY_API_KEYr   )clsr   r   s      g/var/www/html/ai-engine/env/lib/python3.11/site-packages/langchain_community/utilities/tavily_search.pyvalidate_environmentz+TavilySearchAPIWrapper.validate_environment   s.     .$&6
 
 $2        advancedFquerymax_resultssearch_depthinclude_domainsexclude_domainsinclude_answerinclude_raw_contentinclude_imagesc	           
          | j                                         ||||||||d	}	t          j        t           d|	          }
|
                                 |
                                S )N	api_keyr   r   r    r!   r"   r#   r$   r%   /searchjson)r   get_secret_valuerequestspostTAVILY_API_URLraise_for_statusr+   )selfr   r   r    r!   r"   r#   r$   r%   paramsresponses              r   raw_resultsz"TavilySearchAPIWrapper.raw_results&   s     *;;==&(..,#6,

 

 =&&&
 
 

 	!!###}}r   c	           
      r    |                      ||||||||          }	|                     |	d                   S )a  Run query through Tavily Search and return metadata.

        Args:
            query: The query to search for.
            max_results: The maximum number of results to return.
            search_depth: The depth of the search. Can be "basic" or "advanced".
            include_domains: A list of domains to include in the search.
            exclude_domains: A list of domains to exclude from the search.
            include_answer: Whether to include the answer in the results.
            include_raw_content: Whether to include the raw content in the results.
            include_images: Whether to include images in the results.
        Returns:
            query: The query that was searched for.
            follow_up_questions: A list of follow up questions.
            response_time: The response time of the query.
            answer: The answer to the query.
            images: A list of images.
            results: A list of dictionaries containing the results:
                title: The title of the result.
                url: The url of the result.
                content: The content of the result.
                score: The score of the result.
                raw_content: The raw content of the result.
        )r   r    r!   r"   r#   r$   r%   results)r4   clean_results)
r1   r   r   r    r!   r"   r#   r$   r%   raw_search_resultss
             r   r6   zTavilySearchAPIWrapper.resultsD   sT    F "--#%++) 3) . 	
 	
 !!"4Y"?@@@r   c	           
          K   dt           f f	d}	 |	             d{V }
t          j        |
          S )z6Get results from the Tavily Search API asynchronously.r   c            
        	K   j                                         
	d	} t          j                    4 d {V }|                    t
           d|           4 d {V }|j        dk    r@|                                 d {V }|cd d d           d {V  cd d d           d {V  S t          d|j         d|j	                   # 1 d {V swxY w Y   	 d d d           d {V  d S # 1 d {V swxY w Y   d S )Nr'   r)   r*      zError z: )
r   r,   aiohttpClientSessionr.   r/   statustext	Exceptionreason)r2   sessionresdatar"   r#   r!   r%   r$   r   r   r    r1   s       r   fetchz7TavilySearchAPIWrapper.raw_results_async.<locals>.fetch   s     .??AA* ,#2#2"0':"0
 
F ,.. M M M M M M M'"<<>(B(B(B<PP M M M M M M MTWzS((%(XXZZ//////#M M M M M M M M M M M M MM M M M M M M M M M M M M M ((K(K(Ksz(K(KLLLM M M M M M M M M M M M M M MM M M M M M M M M M M M M M M M M M M M M M M M M M M M M Ms<    &C6&'CC62C
C	C6C	C66
D D N)strr+   loads)r1   r   r   r    r!   r"   r#   r$   r%   rE   results_json_strs   `````````  r   raw_results_asyncz(TavilySearchAPIWrapper.raw_results_asyncs   s      	MS 	M 	M 	M 	M 	M 	M 	M 	M 	M 	M 	M 	M 	M 	M( "'======z*+++r   c	           
         K   |                      ||||||||           d {V }	|                     |	d                   S )N)r   r   r    r!   r"   r#   r$   r%   r6   )rI   r7   )
r1   r   r   r    r!   r"   r#   r$   r%   results_jsons
             r   results_asyncz$TavilySearchAPIWrapper.results_async   ss       "33#%++) 3) 4 	
 	
 	
 	
 	
 	
 	
 	
 !!,y"9:::r   r6   c                 \    g }|D ]&}|                     |d         |d         d           '|S )z%Clean results from Tavily Search API.urlcontent)rN   rO   )append)r1   r6   r7   results       r   r7   z$TavilySearchAPIWrapper.clean_results   sT     	 	F  !%=%i0     r   N)__name__
__module____qualname____doc__r   __annotations__r
   model_configr   classmethodr   r   r   rF   r   intr   boolr4   r6   rI   rL   r7    r   r   r   r      s        ((:  L _(###$ 3    [ $# &'&0/1/1)..3).  c] sm	
 "$s), "$s), ! &d^ ! 
   B &'&0/1/1)..3).-A -A-A c]-A sm	-A
 "$s),-A "$s),-A !-A &d^-A !-A 
d-A -A -A -Ad &'&0/1/1)..3).#, #,#, c]#, sm	#,
 "$s),#, "$s),#, !#, &d^#, !#, 
#, #, #, #,P &'&0/1/1)..3).; ;; c]; sm	;
 "$s),; "$s),; !; &d^; !; 
d; ; ; ;.
T$Z 
DJ 
 
 
 
 
 
r   r   )rU   r+   typingr   r   r   r   r<   r-   langchain_core.utilsr   pydanticr	   r
   r   r   r/   r   r[   r   r   <module>r_      s      , , , , , , , , , , , ,   5 5 5 5 5 5 F F F F F F F F F F F F)g g g g gY g g g g gr   