
    3fi                        d Z ddlmZ ddlZddlZddlmZmZmZm	Z	m
Z
mZmZ ddlmZ erddlmZ ddlmZ ddlmZ dd	lmZmZmZ d
ZdZddZ G d de      Zy)z6Cassandra-based chat message history, based on cassIO.    )annotationsN)TYPE_CHECKINGAnyDictIterableListOptionalSequence)	SetupMode)Session)RowType)BaseChatMessageHistory)BaseMessagemessage_to_dictmessages_from_dictmessage_storec                    | D cg c]  }|d   	 c}d d d   }|D cg c]  }t        j                  |       }}t        |      }|S c c}w c c}w )N	body_blob)jsonloadsr   )rowsrowmessage_blobsmessage_blobitemsmessagess         r/var/www/auto_recruiter/arenv/lib/python3.12/site-packages/langchain_community/chat_message_histories/cassandra.py_rows_to_messagesr      sS    156#S%6tt<M:GH,TZZ%HEH!%(HO 7Hs
   A	Ac                      e Zd ZdZddeefej                  d	 	 	 	 	 	 	 	 	 	 	 	 	 ddZe	dd       Z
ddZddZddZdd	Zdd
Zy)CassandraChatMessageHistoryz1Chat message history that is backed by Cassandra.N)
setup_modec                   	 ddl m} || _        || _        i }|t        j                  k(  rd|d<    |d||||ddgd|t        j                  k(  d	|| _	        y
# t        t        f$ r t        d      w xY w)a  
        Initialize a new instance of CassandraChatMessageHistory.

        Args:
            session_id: arbitrary key that is used to store the messages
                of a single chat session.
            session: Cassandra driver session.
                If not provided, it is resolved from cassio.
            keyspace: Cassandra key space. If not provided, it is resolved from cassio.
            table_name: name of the table to use.
            ttl_seconds: time-to-live (seconds) for automatic expiration
                of stored entries. None (default) for no expiration.
            setup_mode: mode used to create the Cassandra table (SYNC, ASYNC or OFF).
        r   )ClusteredCassandraTablezTCould not import cassio python package. Please install it with `pip install cassio`.Tasync_setupTEXTTIMEUUIDDESC)sessionkeyspacetablettl_secondsprimary_key_typeordering_in_partitionskip_provisioningN )
cassio.tabler$   ImportErrorModuleNotFoundError
session_idr,   r   ASYNCOFFr+   )	selfr4   r)   r*   
table_namer,   r"   r$   kwargss	            r   __init__z$CassandraChatMessageHistory.__init__$   s    0	< %&!#($(F=!, 	
#$j1"((IMM9	
 	

 01 	? 	s   A A4c                d    | j                   j                  | j                        }t        |      S )%Retrieve all session messages from DBpartition_id)r+   get_partitionr4   r   r7   r   s     r   r   z$CassandraChatMessageHistory.messagesS   s1     zz'' ( 
 !&&    c                   K   | j                   j                  | j                         d{   }t        |      S 7 w)r<   r=   N)r+   aget_partitionr4   r   r@   s     r   aget_messagesz)CassandraChatMessageHistory.aget_messages\   s?      ZZ.. / 
 
 !&&
s   *><>c           	         t        j                         }| j                  j                  | j                  |t        j                  t        |            | j                         y)z]Write a message to the table

        Args:
            message: A message to write.
        r>   row_idr   r,   N)	uuiduuid4r+   putr4   r   dumpsr   r,   )r7   messagethis_row_ids      r   add_messagez'CassandraChatMessageHistory.add_messaged   sH     jjl

jj!9:((	 	 	
rA   c           
        K   |D ]m  }t        j                         }| j                  j                  | j                  |t        j                  t        |            | j                         d {    o y 7 w)NrF   )	rH   rI   r+   aputr4   r   rK   r   r,   )r7   r   rL   rM   s       r   aadd_messagesz)CassandraChatMessageHistory.aadd_messagesr   sh      	G**,K**//!__"**_W%=> ,,	 "   	s   A,A8.A6/A8c                N    | j                   j                  | j                         yzClear session memory from DBN)r+   delete_partitionr4   r7   s    r   clearz!CassandraChatMessageHistory.clear|   s    

##DOO4rA   c                j   K   | j                   j                  | j                         d{    y7 wrS   )r+   adelete_partitionr4   rU   s    r   aclearz"CassandraChatMessageHistory.aclear   s"     jj**4??;;;s   )313)r4   strr)   zOptional[Session]r*   zOptional[str]r8   rZ   r,   zOptional[int]r"   r   returnNone)r[   List[BaseMessage])rL   r   r[   r\   )r   zSequence[BaseMessage]r[   r\   )r[   r\   )__name__
__module____qualname____doc__DEFAULT_TABLE_NAMEDEFAULT_TTL_SECONDSr   SYNCr:   propertyr   rD   rN   rQ   rV   rY   r0   rA   r   r!   r!   !   s    ;
 &*"&,%8-
 !*-
-
 #-
  	-

 -
 #-
 -
 
-
^ ' ''
5<rA   r!   )r   zIterable[RowType]r[   r]   )ra   
__future__r   r   rH   typingr   r   r   r   r   r	   r
   'langchain_community.utilities.cassandrar   cassandra.clusterr   cassio.table.table_typesr   langchain_core.chat_historyr   langchain_core.messagesr   r   r   rb   rc   r   r!   r0   rA   r   <module>rm      sU    < "   O O O =)0 >  %  a<"8 a<rA   