
    q9iA                     H   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	 d dl
mZ d dlmZmZmZ d dlmZmZmZmZ d d	lmZ d d
lmZ d dlmZmZmZmZmZmZm Z m!Z!m"Z"m#Z#m$Z$m%Z%m&Z&m'Z'm(Z(m)Z)m*Z* d dl+m,Z,m-Z-m.Z.m/Z/ d dl0m1Z1 d dl2m3Z3m4Z4m5Z5m6Z6  G d dee      Z7 G d dee      Z8y)    N)OptionalSequence)UUID)override)UserIdentity)maybe_set_tenant_and_database)AsyncAdminAPIAsyncClientAPIAsyncServerAPI)CreateCollectionConfigurationUpdateCollectionConfiguration.validate_embedding_function_conflict_on_create+validate_embedding_function_conflict_on_get)AsyncCollection)SharedSystemClient)CollectionMetadata
DataLoader	Documents
EmbeddableEmbeddingFunction
Embeddings	GetResultIDsIncludeIncludeMetadataDocuments!IncludeMetadataDocumentsDistancesLoadable	MetadatasQueryResultSchemaURIsDefaultEmbeddingFunction)DEFAULT_DATABASEDEFAULT_TENANTSettingsSystem)ChromaError)DatabaseTenantWhereWhereDocumentc                   b   e Zd ZU dZeed<   eZeed<   e	Z
eed<   eed<   eee	 e       fdedededd fd       Zeee	fd	edededd fd
       Zeed	eddfd              Zedefd       Zee	fdededdfd       Zededdfd       ZdededdfdZedefd       Ze	 dAdee   dee   dee   fd       Zedefd       Zeddd e       ddfdedee    dee!   dee"   dee#e$      dee%e&      de'defd       Z(e e       dfdedee#e$      dee%e&      defd        Z)eddd e       dfdedee    dee!   dee"   dee#e$      dee%e&      defd!       Z*e	 	 	 dBd"e+d#ee   d$ee"   d%ee,   ddf
d&       Z-ededdfd'       Z.e	 	 	 dBd(e/d)e+d*e0d+ee1   d,ee2   d-ee3   de'fd.       Z4e	 	 	 	 dCd)e+d(e/d*ee0   d+ee1   d,ee2   d-ee3   de'fd/       Z5e	 	 	 dBd)e+d(e/d*e0d+ee1   d,ee2   d-ee3   de'fd0       Z6ed)e+defd1       Z7edDd)e+d3ede8fd4       Z9eddddde:fd)e+d(ee/   d5ee;   dee   dee   d6ee<   d7e=de8fd8       Z>	 	 dAd)e+d(ee/   d5ee;   d6ee<   ddf
d9Z?edd2dde@fd)e+d:e0d(ee/   d;ed5ee;   d6ee<   d7e=deAfd<       ZBede'fd=       ZCedefd>       ZDedefd?       ZEedefd@       ZFy)EAsyncClienta  A client for Chroma. This is the main entrypoint for interacting with Chroma.
    A client internally stores its tenant and database and proxies calls to a
    Server API instance of Chroma. It treats the Server API and corresponding System
    as a singleton, so multiple clients connecting to the same resource will share the
    same API instance.

    Client implementations should be implement their own API-caching strategies.
    _admin_clienttenantdatabase_serversettingsreturnc                   K    | |      }t        j                  |j                         ||_        ||_        |j                  j                  t              |_        |j                          d {   }t        ||j                  ||      \  }}|r||_        |r||_        t        j                  |j                        |_        |j                  |j                  |j                         d {    |j                          |S 7 7 w)Nr2   )4overwrite_singleton_tenant_database_access_from_authuser_provided_tenantuser_provided_databaser/   r0   )r   _populate_data_from_system_systemr/   r0   instancer   r1   get_user_identityr   ;chroma_overwrite_singleton_tenant_database_access_from_authAsyncAdminClientfrom_systemr.   _validate_tenant_database_submit_client_start_event)clsr/   r0   r2   selfuser_identitymaybe_tenantmaybe_databases           b/var/www/html/leadgen/airagagent/rag_env/lib/python3.12/site-packages/chromadb/api/async_client.pycreatezAsyncClient.create;   s      H%55dllC  ||,,^<"4466'DAI  BF  BF!'#+	(
$n &DK*DM-99$,,G,,DKK$--,XXX'')% 7 	Ys%   A.D0D1B D1D
2D
Dsystemc                 `   K   t         j                  |||j                         d{   S 7 w)zRCreate a client from an existing system. This is useful for testing and debugging.N)r-   rI   r2   )rC   rJ   r/   r0   s       rH   from_system_asynczAsyncClient.from_system_async`   s'      !''&//JJJJs   %.,.r   c                     t        d       )zNAsyncClient cannot be created synchronously. Use .from_system_async() instead.)NotImplementedError)rC   rJ   s     rH   r@   zAsyncClient.from_systemk   s     "\
 	
    c                 R   K   | j                   j                          d {   S 7 wN)r1   r=   rD   s    rH   r=   zAsyncClient.get_user_identityv   s     \\335555   '%'Nc                 b   K   | j                  ||       d {    || _        || _        y 7 wNr9   rA   r/   r0   )rD   r/   r0   s      rH   
set_tenantzAsyncClient.set_tenantz   s3     ,,FX,NNN  	Os   /-/c                 h   K   | j                  | j                  |       d {    || _        y 7 wrU   rV   )rD   r0   s     rH   set_databasezAsyncClient.set_database   s0     ,,DKK(,SSS  	Ts   !202c                 |  K   	 | j                   j                  |       d {    	 | j                   j                  ||       d {    y 7 +# t        j                  $ r t	        d      t
        $ r}|d }~wt        $ r t	        d| d      w xY w7 N# t        j                  $ r t	        d      w xY ww)NnamezACould not connect to a Chroma server. Are you sure it is running?zCould not connect to tenant z. Are you sure it exists?r\   r/   )r.   
get_tenanthttpxConnectError
ValueErrorr'   	Exceptionget_database)rD   r/   r0   es       rH   rA   z%AsyncClient._validate_tenant_database   s     	$$//V/<<<	$$11x1OOO =!! 	S   	G 	.vh6OP 	 P!! 	S 	s[   B<A AA  B 	B
B B<A &B7A99BB<B B99B<c                 R   K   | j                   j                          d {   S 7 wrQ   )r1   	heartbeatrR   s    rH   rf   zAsyncClient.heartbeat   s     \\++----rS   limitoffsetc                    K   | j                   j                  ||| j                  | j                         d {   }|D cg c]  }t	        | j                   |       c}S 7 )c c}w w)Nr9   )clientmodel)r1   list_collectionsr/   r0   r   )rD   rg   rh   modelsrk   s        rH   rl   zAsyncClient.list_collections   sb      ||446$++ 5 
 
 PVVet||5AVV
 Ws!   7A*A#A*A% A*%A*c                    K   | j                   j                  | j                  | j                         d {   S 7 wrU   )r1   count_collectionsr/   r0   rR   s    rH   ro   zAsyncClient.count_collections   s9     \\33;; 4 
 
 	
 
s   5><>Fr\   schemaconfigurationmetadataembedding_functiondata_loaderget_or_createc           	        K   |i }|j                  d      }t        ||       ||||d<   | j                  j                  ||||| j                  | j
                  |       d {   }	t        | j                  |	||      S 7 w)Nrs   )r\   rp   rq   rr   r/   r0   ru   rj   rk   rs   rt   )getr   r1   create_collectionr/   r0   r   )
rD   r\   rp   rq   rr   rs   rt   ru   configuration_efrk   s
             rH   ry   zAsyncClient.create_collection   s       M(,,-AB6 0	
 ).>.F2DM./ll44';;]]' 5 
 
 <<1#	
 	

s   A$B&B'Bc                   K   | j                   j                  || j                  | j                         d {   }|j                  j                  d      }t        ||       t        | j                   |||      S 7 Dw)Nr\   r/   r0   rs   rw   )r1   get_collectionr/   r0   configuration_jsonrx   r   r   )rD   r\   rs   rt   rk   persisted_ef_configs         rH   r}   zAsyncClient.get_collection   s      ll11;;]] 2 
 

 $66::;OP3 3	
 <<1#	
 	

s   6A?A=AA?c                 \  K   |i }|j                  d      }t        ||       ||||d<   | j                  j                  ||||| j                  | j
                         d {   }|j                  j                  d      }	t        ||	       t        | j                  |||      S 7 Dw)Nrs   )r\   rp   rq   rr   r/   r0   rw   )	rx   r   r1   get_or_create_collectionr/   r0   r~   r   r   )
rD   r\   rp   rq   rr   rs   rt   rz   rk   r   s
             rH   r   z$AsyncClient.get_or_create_collection   s       M(,,-AB6 0	
 ).>.F2DM./ll;;';;]] < 
 
 $66::;OP3 3	
 <<1#	
 	

s   A#B,%B*&AB,idnew_namenew_metadatanew_configurationc                    K   | j                   j                  ||||| j                  | j                         d {   S 7 w)N)r   r   r   r   r/   r0   )r1   _modifyr/   r0   )rD   r   r   r   r   s        rH   r   zAsyncClient._modify#  sH      \\))%/;;]] * 
 
 	
 
s   9AA Ac                    K   | j                   j                  || j                  | j                         d {   S 7 w)Nr|   )r1   delete_collectionr/   r0   rD   r\   s     rH   r   zAsyncClient.delete_collection4  s?     
 \\33;;]] 4 
 
 	
 
s   6?=?idscollection_id
embeddings	metadatas	documentsurisc           
         K   | j                   j                  ||||||| j                  | j                         d {   S 7 w)N)r   r   r   r   r   r   r/   r0   )r1   _addr/   r0   )rD   r   r   r   r   r   r   s          rH   r   zAsyncClient._addC  sN      \\&&'!;;]] ' 	
 	
 		
 	
   ;AAAc           
         K   | j                   j                  ||||||| j                  | j                         d {   S 7 wN)r   r   r   r   r   r   r/   r0   )r1   _updater/   r0   rD   r   r   r   r   r   r   s          rH   r   zAsyncClient._updateX  N      \\))'!;;]] * 	
 	
 		
 	
r   c           
         K   | j                   j                  ||||||| j                  | j                         d {   S 7 wr   )r1   _upsertr/   r0   r   s          rH   r   zAsyncClient._upsertm  r   r   c                 V   K   | j                   j                  |       d {   S 7 w)N)r   )r1   _count)rD   r   s     rH   r   zAsyncClient._count  s/     \\((' ) 
 
 	
 
    )')
   nc                 X   K   | j                   j                  ||       d {   S 7 w)N)r   r   )r1   _peek)rD   r   r   s      rH   r   zAsyncClient._peek  s2     \\''' ( 
 
 	
 
   !*(*wherewhere_documentincludec                    K   | j                   j                  |||||||| j                  | j                  	       d {   S 7 w)N)	r   r   r   rg   rh   r   r   r/   r0   )r1   _getr/   r0   )rD   r   r   r   rg   rh   r   r   s           rH   r   zAsyncClient._get  sQ      \\&&');;]] ' 

 

 
	
 

   <AAAc                    K   | j                   j                  ||||| j                  | j                         d {    y 7 w)N)r   r   r   r   r/   r0   )r1   _deleter/   r0   )rD   r   r   r   r   s        rH   r   zAsyncClient._delete  sC      ll""');;]] # 
 	
 	
s   9AAAquery_embeddings	n_resultsc                    K   | j                   j                  |||||||| j                  | j                  	       d {   S 7 w)N)	r   r   r   r   r   r   r   r/   r0   )r1   _queryr/   r0   )rD   r   r   r   r   r   r   r   s           rH   r   zAsyncClient._query  sQ      \\(('-);;]] ) 

 

 
	
 

r   c                 R   K   | j                   j                          d {   S 7 wrQ   )r1   resetrR   s    rH   r   zAsyncClient.reset  s     \\''))))rS   c                 R   K   | j                   j                          d {   S 7 wrQ   )r1   get_versionrR   s    rH   r   zAsyncClient.get_version  s     \\--////rS   c                 6    | j                   j                         S rQ   )r1   get_settingsrR   s    rH   r   zAsyncClient.get_settings  s    ||((**rO   c                 R   K   | j                   j                          d {   S 7 wrQ   )r1   get_max_batch_sizerR   s    rH   r   zAsyncClient.get_max_batch_size  s     \\446666rS   )NN)NNN)NNNN)r   )G__name__
__module____qualname____doc__r	   __annotations__r$   r/   strr#   r0   r   classmethodr%   rI   r&   rL   r   r@   r   r=   rW   rY   rA   intrf   r   r   r   rl   ro   r"   r    r   r   r   r   r   r   boolry   r}   r   r   r   r   r   r   r   r   r   r!   r   r   r   r   r   r   r   r*   r+   r   r   r   r   r   r   r   r   r   r    rO   rH   r-   r-   )   sc    !  FC $Hc$ %(%Z	"" " 	"
 
" "H 
 %(	KK K 	K
 
K K 

 

  
 6 6 6 <L !s !c !QU ! !
 !3 !4 ! !c S T 0 . . . CGWc]W3;C=W	/	"W W 
 
 

  $(AE15 %&6:#(
(
  (
   =>	(

 -.(
 %j)
(
 j23(
 (
 
(
 (
T  %&6:

 %j)

 j23
 

 
4  $(AE15 %&6:*
*
  *
   =>	*

 -.*
 %j)
*
 j23*
 
*
 *
X  #'59EI

 3-
 12	

 $$AB
 

 
  

 

 
  *.)-#

 
 	

 I&
 I&
 tn
 

 
( 
 ,0)-)-#

 
 Z(	

 I&
 I&
 tn
 

 
(  *.)-#

 
 	

 I&
 I&
 tn
 

 
( 
$ 
3 
 

 
 
# 
y 
 
  "!%# $263

 c]
 	

 }
 
 !/
 
 

 
4 "&26

 c]
 	

 !/
 

  
 "!%26<

 %
 c]	

 
 
 !/
 
 

 
, *T * * 03 0 0 +h + + 7# 7 7rO   r-   c                   <    e Zd ZU eed<    e       fdeddf fdZeefde	de	ddfd       Z
eefde	de	defd	       Zeefde	de	ddfd
       Zeddefdee   dee   de	dee   fd       Zede	ddfd       Zede	defd       Zeededd fd              Z xZS )r?   r1   r2   r3   Nc                 l    t         |   |       | j                  j                  t              | _        y rQ   )super__init__r;   r<   r   r1   )rD   r2   	__class__s     rH   r   zAsyncAdminClient.__init__  s&    "||,,^<rO   r\   r/   c                 X   K   | j                   j                  ||       d {   S 7 wNr]   )r1   create_databaserD   r\   r/   s      rH   r   z AsyncAdminClient.create_database  &     \\11tF1KKKKr   c                 X   K   | j                   j                  ||       d {   S 7 wr   )r1   rc   r   s      rH   rc   zAsyncAdminClient.get_database  s&     \\..D.HHHHr   c                 X   K   | j                   j                  ||       d {   S 7 wr   )r1   delete_databaser   s      rH   r   z AsyncAdminClient.delete_database  r   r   rg   rh   c                 Z   K   | j                   j                  |||       d {   S 7 w)N)rg   rh   r/   )r1   list_databases)rD   rg   rh   r/   s       rH   r   zAsyncAdminClient.list_databases  s5      \\00v 1 
 
 	
 
s   "+)+c                 V   K   | j                   j                  |       d {   S 7 wNr[   )r1   create_tenantr   s     rH   r   zAsyncAdminClient.create_tenant  s$     \\//T/::::r   c                 V   K   | j                   j                  |       d {   S 7 wr   )r1   r^   r   s     rH   r^   zAsyncAdminClient.get_tenant  s$     \\,,$,7777r   rJ   c                 V    t        j                  |        | |j                        }|S )Nr5   )r   r:   r2   )rC   rJ   r<   s      rH   r@   zAsyncAdminClient.from_system  s%     	55f=0rO   )r   r   r   r   r   r%   r   r   r$   r   r   r(   rc   r   r   r   r   r   r   r)   r^   r   r&   r@   __classcell__)r   s   @rH   r?   r?     sz   ,4J = =4 = =K L# Ls LPT L L :H Is IC IX I I =K L# Ls LPT L L   $ $$	
}
 
 	

 
(	
 
 ; ; ; ; 8S 8V 8 8  
  rO   r?   )9r_   typingr   r   uuidr   	overridesr   chromadb.authr   chromadb.auth.utilsr   chromadb.apir	   r
   r   %chromadb.api.collection_configurationr   r   r   r   #chromadb.api.models.AsyncCollectionr   !chromadb.api.shared_system_clientr   chromadb.api.typesr   r   r   r   r   r   r   r   r   r   r   r   r   r   r    r!   r"   chromadb.configr#   r$   r%   r&   chromadb.errorsr'   chromadb.typesr(   r)   r*   r+   r-   r?   r   rO   rH   <module>r      s|     %   & = F F  @ @    & O N ' A Ar7$n r7n.)= .rO   