
    Ig#                        d Z ddlmZ ddlmZmZmZmZ ddl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 ddlmZ ddlmZ ddlmZ  G d dee      Zy)zDHypothetical Document Embeddings.

https://arxiv.org/abs/2212.10496
    )annotations)AnyDictListOptionalN)CallbackManagerForChainRun)
Embeddings)BaseLanguageModel)StrOutputParser)BasePromptTemplate)Runnable)
ConfigDict)Chain)
PROMPT_MAP)LLMChainc                      e Zd ZU dZded<   ded<    edd      Zedd	       Zedd
       Z	ddZ
ddZddZ	 d	 	 	 	 	 ddZe	 	 d	 	 	 	 	 	 	 	 	 	 	 dd       Zedd       Zy)HypotheticalDocumentEmbedderzrGenerate hypothetical document for query, and then embed that.

    Based on https://arxiv.org/abs/2212.10496
    r	   base_embeddingsr   	llm_chainTforbid)arbitrary_types_allowedextrac                P    | j                   j                  j                         d   S )z Input keys for Hyde's LLM chain.required)r   input_schemamodel_json_schemaselfs    W/var/www/html/answerous/venv/lib/python3.12/site-packages/langchain/chains/hyde/base.py
input_keysz'HypotheticalDocumentEmbedder.input_keys&   s"     ~~**<<>zJJ    c                h    t        | j                  t              r| j                  j                  S dgS )z!Output keys for Hyde's LLM chain.text)
isinstancer   r   output_keysr   s    r   r%   z(HypotheticalDocumentEmbedder.output_keys+   s)     dnnh/>>---8Or!   c                8    | j                   j                  |      S )zCall the base embeddings.)r   embed_documents)r   textss     r   r'   z,HypotheticalDocumentEmbedder.embed_documents3   s    ##33E::r!   c                ^    t        t        j                  |      j                  d            S )z)Combine embeddings into final embeddings.r   )axis)listnparraymean)r   
embeddingss     r   combine_embeddingsz/HypotheticalDocumentEmbedder.combine_embeddings7   s$    BHHZ(--1-566r!   c                    | j                   d   }| j                  j                  ||i      }t        | j                  t              r|| j
                  d      g}n|g}| j                  |      }| j                  |      S )z1Generate a hypothetical document and embedded it.r   )r    r   invoker$   r   r%   r'   r0   )r   r#   var_nameresult	documentsr/   s         r   embed_queryz(HypotheticalDocumentEmbedder.embed_query;   sw    ??1%&&$'78dnnh/ 0 0 345II)))4
&&z22r!   Nc                    |xs t        j                         }| j                  j                  |d|j	                         i      S )zCall the internal llm chain.	callbacks)config)r   get_noop_managerr   r2   	get_child)r   inputsrun_manager_run_managers       r   _callz"HypotheticalDocumentEmbedder._callF   sG     #S&@&Q&Q&S~~$$K)?)?)AB % 
 	
r!   c                    ||}n>||t         v r
t         |   }n*t        dt        t        j                                d      ||z  t	               z  } | d||d|S )zILoad and use LLMChain with either a specific prompt key or custom prompt.zHMust specify prompt_key if custom_prompt not provided. Should be one of .)r   r    )r   
ValueErrorr+   keysr   )clsllmr   
prompt_keycustom_promptkwargspromptr   s           r   from_llmz%HypotheticalDocumentEmbedder.from_llmQ   sy     $"F#
j(@
+F:??,-.a1 
 SL?#44	R?iR6RRr!   c                     y)N
hyde_chainrB   r   s    r   _chain_typez(HypotheticalDocumentEmbedder._chain_typeh   s    r!   )return	List[str])r(   rP   rO   List[List[float]])r/   rQ   rO   List[float])r#   strrO   rR   )N)r<   zDict[str, Any]r=   z$Optional[CallbackManagerForChainRun]rO   zDict[str, str])NN)rF   r
   r   r	   rG   zOptional[str]rH   zOptional[BasePromptTemplate]rI   r   rO   r   )rO   rS   )__name__
__module____qualname____doc____annotations__r   model_configpropertyr    r%   r'   r0   r6   r?   classmethodrK   rN   rB   r!   r   r   r      s    
   $L
 K K  ;7	3 =A	
	
 :	
 
		
 
 %)6:SS $S "	S
 4S S 
&S S,  r!   r   )rW   
__future__r   typingr   r   r   r   numpyr,   langchain_core.callbacksr   langchain_core.embeddingsr	   langchain_core.language_modelsr
   langchain_core.output_parsersr   langchain_core.promptsr   langchain_core.runnablesr   pydanticr   langchain.chains.baser   langchain.chains.hyde.promptsr   langchain.chains.llmr   r   rB   r!   r   <module>ri      sF   
 # , ,  ? 0 < 9 5 -  ' 4 )R5* Rr!   