
    dh                         S SK JrJrJrJrJrJr  S SKJr  S SK	J
r
  S\\\\4      S\4S jr " S S\5      r " S	 S
\
5      rg)    )AnyCallableIteratorListOptionalTuple)Document)
BaseLoaderdocsreturnc                 R    SR                  U  Vs/ sH  oS   PM	     sn5      $ s  snf )z#Default joiner for content columns.
   )join)r   docs     f/var/www/html/shao/venv/lib/python3.13/site-packages/langchain_community/document_loaders/rocksetdb.pydefault_joinerr      s%    99-!f-..-s   $c                   8   ^  \ rS rSrSrS\S\4U 4S jjrSrU =r$ )ColumnNotFoundError   zColumn not found error.missing_keyqueryc                 0   > [         TU ]  SU SU 35        g )NzColumn "z" not selected in query:
)super__init__)selfr   r   	__class__s      r   r   ColumnNotFoundError.__init__   s    8K=0J5'RS     )	__name__
__module____qualname____firstlineno____doc__strr   __static_attributes____classcell__)r   s   @r   r   r      s"    !TC T T Tr   r   c                       \ rS rSrSrS\4S\S\S\\   S\	\\      S\
\\\\4      /\4   4
S	 jjrS
\\   4S jrSrg)RocksetLoader   a  Load from a `Rockset` database.

To use, you should have the `rockset` python package installed.

Example:
    .. code-block:: python

        # This code will load 3 records from the "langchain_demo"
        # collection as Documents, with the `text` column used as
        # the content

        from langchain_community.document_loaders import RocksetLoader
        from rockset import RocksetClient, Regions, models

        loader = RocksetLoader(
            RocksetClient(Regions.usw2a1, "<api key>"),
            models.QueryRequestSql(
                query="select * from langchain_demo limit 3"
            ),
            ["text"]
        )
    )
Nclientr   content_keysmetadata_keyscontent_columns_joinerc                     SSK JnJn  SSKJn  [        X5      (       d  [        S[        U5       35      e[        X(5      (       d  [        S[        U5       35      eXl	        X l
        X0l        XPl        X@l        X`l        Xl         U R                  R!                  S5        g! [
         a    [        S5      ef = f! ["         a     gf = f)	a  Initialize with Rockset client.

Args:
    client: Rockset client object.
    query: Rockset query object.
    content_keys: The collection columns to be written into the `page_content`
        of the Documents.
    metadata_keys: The collection columns to be written into the `metadata` of
        the Documents. By default, this is all the keys in the document.
    content_columns_joiner: Method that joins content_keys and its values into a
        string. It's method that takes in a List[Tuple[str, Any]]],
        representing a list of tuples of (column name, column value).
        By default, this is a method that joins each column value with a new
        line. This method is only relevant if there are multiple content_keys.
r   )QueryPaginatorRocksetClient)QueryRequestSqlz]Could not import rockset client python package. Please install it with `pip install rockset`.z;client should be an instance of rockset.RocksetClient, got zBquery should be an instance of rockset.model.QueryRequestSql, got 	langchainN)rocksetr1   r2   rockset.modelsr3   ImportError
isinstance
ValueErrortyper,   r   r-   r/   r.   	paginatorrequest_modelset_applicationAttributeError)	r   r,   r   r-   r.   r/   r1   r2   r3   s	            r   r   RocksetLoader.__init__-   s    .	=6 &00F|n& 
 %11E{m% 
 
(&<#*',	KK''47  	@ 	8  		s   B% 	B> %B;>
C
Cr   c              #     #    U R                   R                  R                  U R                  S9R                  nU Hg  n [	        U R                  U R                   Vs/ sH  o3X#   4PM
     sn5      U R                  b  U R                   Vs0 sH  o3X#   _M	     snOUS9v   Mi     g s  snf s  snf ! [         a(  n[        UR                  S   U R                  5      eS nAff = f7f)N)sql)page_contentmetadatar   )r,   Queriesr   resultsr	   r/   r-   r.   KeyErrorr   args)r   query_resultsr   coles        r   	lazy_loadRocksetLoader.lazy_loadg   s     ++11

 2 

' 	 !C!%!<!<484E4EF4ESsx4EF" ))5 8<7I7IJ7I38m7IJ  ! GJ
  *FF1Itzz sG   =C' B2B(,!B2B-	B2#C'(
B22
C$<#CC$$C')r,   r/   r-   r.   r;   r   r<   )r!   r"   r#   r$   r%   r   r   r   r&   r   r   r   r   r   r	   rK   r'   r    r   r   r*   r*      s}    : .2IW88 8 3i	8
  S	*8 !)$uS#X*?)@#)E F8t8H- r   r*   N)typingr   r   r   r   r   r   langchain_core.documentsr	   )langchain_community.document_loaders.baser
   r&   r   	Exceptionr   r*   r    r   r   <module>rQ      sM    A A - @/eCHo. /3 /
T) TfJ fr   