o
    Zh                     @   sR   d dl Z d dlmZmZmZ d dlmZ d dlmZ e 	e
ZG dd deZdS )    N)IteratorListOptional)Document)
BaseLoaderc                   @   s`   e Zd ZdZddddededededeee  d	eee  d
dfddZd
ee	 fddZ
dS )CouchbaseLoaderaU  Load documents from `Couchbase`.

    Each document represents one row of the result. The `page_content_fields` are
    written into the `page_content`of the document. The `metadata_fields` are written
    into the `metadata` of the document. By default, all columns are written into
    the `page_content` and none into the `metadata`.
    N)page_content_fieldsmetadata_fieldsconnection_stringdb_usernamedb_passwordqueryr   r	   returnc             
   C   s   zddl m} ddlm} ddlm}	 W n ty% }
 ztd|
d}
~
ww |s,td|s2td|s8td	|||}|||	|| _|| _	|| _
|| _dS )
a  Initialize Couchbase document loader.

        Args:
            connection_string (str): The connection string to the Couchbase cluster.
            db_username (str): The username to connect to the Couchbase cluster.
            db_password (str): The password to connect to the Couchbase cluster.
            query (str): The SQL++ query to execute.
            page_content_fields (Optional[List[str]]): The columns to write into the
                `page_content` field of the document. By default, all columns are
                written.
            metadata_fields (Optional[List[str]]): The columns to write into the
                `metadata` field of the document. By default, no columns are written.
        r   )PasswordAuthenticator)Cluster)ClusterOptionsz]Could not import couchbase package.Please install couchbase SDK with `pip install couchbase`.Nz#connection_string must be provided.zdb_username must be provided.zdb_password must be provided.)Zcouchbase.authr   Zcouchbase.clusterr   Zcouchbase.optionsr   ImportError
ValueErrorclusterr   r   r	   )selfr
   r   r   r   r   r	   r   r   r   eauth r   e/var/www/html/lang_env/lib/python3.10/site-packages/langchain_community/document_loaders/couchbase.py__init__   s4   
zCouchbaseLoader.__init__c                 #   s    ddl m} | j|dd | j| j}|D ]2| j}| j  s)t  |s-g }fdd|D }d	 fdd	
 D }t||d
V  qdS )z1Load Couchbase data into Document objects lazily.r   )	timedelta   )secondsc                    s   i | ]}| | qS r   r   ).0field)rowr   r   
<dictcomp>Z   s    z-CouchbaseLoader.lazy_load.<locals>.<dictcomp>
c                 3   s*    | ]\}}| v r| d | V  qdS )z: Nr   )r   kv)r   r   r   	<genexpr>\   s     z,CouchbaseLoader.lazy_load.<locals>.<genexpr>)Zpage_contentmetadataN)datetimer   r   Zwait_until_readyr   r	   r   listkeysjoinitemsr   )r   r   resultr	   r&   Zdocumentr   )r   r    r   	lazy_loadG   s"   zCouchbaseLoader.lazy_load)__name__
__module____qualname____doc__strr   r   r   r   r   r-   r   r   r   r   r      s(    

	
3r   )loggingtypingr   r   r   Zlangchain_core.documentsr   Z)langchain_community.document_loaders.baser   	getLoggerr.   loggerr   r   r   r   r   <module>   s    
