a
    bg                     @  st   d dl mZ d dlmZmZ d dlZd dlZd dlmZm	Z	 d dl
mZ d dlmZ d dlmZ G dd	 d	eZdS )
    )annotations)ListOptionalN)#AsyncCallbackManagerForRetrieverRunCallbackManagerForRetrieverRun)Document)BaseRetriever)
ConfigDictc                   @  s   e Zd ZU dZded< ded< dZded< dZd	ed
< dZded< eddZ	ddddddZ
ddddddZdddddZdS )ChatGPTPluginRetrieverz`ChatGPT plugin` retriever.strurlbearer_token   inttop_kNzOptional[dict]filterzOptional[aiohttp.ClientSession]
aiosessionT)Zarbitrary_types_allowedr   zList[Document])queryrun_managerreturnc                C  s   |  |\}}}tj|||d}| d d d }g }|D ]D}	|	d}
|	d|	}|drn|d|d< |t|
|d q<|S )	N)jsonheadersresultsr   textmetadata	source_idsourceZpage_contentr   )_create_requestrequestspostr   popgetappendr   )selfr   r   r   r   r   responser   docsdcontentr    r)   /var/www/html/cobodadashboardai.evdpl.com/venv/lib/python3.9/site-packages/langchain_community/retrievers/chatgpt_plugin_retriever.py_get_relevant_documents"   s    

z.ChatGPTPluginRetriever._get_relevant_documentsr   c             
     s`  |  |\}}}| jst 4 I d H h}|j|||d4 I d H $}| I d H }W d   I d H  qz1 I d H sp0    Y  W d   I d H  q1 I d H s0    Y  nT| jj|||d4 I d H $}| I d H }W d   I d H  q1 I d H s0    Y  |d d d }	g }
|	D ]H}|d}|d|}|drF|d|d< |
t	||d q|
S )	N)r   r   r   r   r   r   r   r   r   )
r   r   aiohttpZClientSessionr    r   r!   r"   r#   r   )r$   r   r   r   r   r   sessionr%   resr   r&   r'   r(   r   r)   r)   r*   _aget_relevant_documents1   s&    d8
z/ChatGPTPluginRetriever._aget_relevant_documentsztuple[str, dict, dict])r   r   c                 C  s>   | j  d}d|| j| jdgi}dd| j d}|||fS )Nz/queryZqueries)r   r   r   zapplication/jsonzBearer )zContent-TypeAuthorization)r   r   r   r   )r$   r   r   r   r   r)   r)   r*   r   J   s    

z&ChatGPTPluginRetriever._create_request)__name__
__module____qualname____doc____annotations__r   r   r   r	   Zmodel_configr+   r/   r   r)   r)   r)   r*   r
      s   
r
   )
__future__r   typingr   r   r,   r   Zlangchain_core.callbacksr   r   Zlangchain_core.documentsr   Zlangchain_core.retrieversr   Zpydanticr	   r
   r)   r)   r)   r*   <module>   s   