
    Zi0                        U d dl Zd dlZd dlmZ d dlZd dlZd dl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mZmZmZ d dlmZ d dlmZmZmZmZ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%c m&c mZ d d
l m'Z'm(Z(m)Z)m*Z*m+Z+m,Z,m-Z-m.Z.m/Z/m0Z0 d dl1m2Z2 d dl3m%c m&c m4Z4 d dl5m6Z6m7Z7m8Z8 d dl9Z:d dl;Z;d dl<m=Z=m>Z>  e2d           a?ee@eAf         eBd<   dZCde@ddfdZDdMdZE ejF         ejG        ejH                  d          ZIejJ        ee:jK        e:jL        e:jM        f                  eBd<    ejF         ejN        dd          d          ZOejJ        eA         eBd<   e G d d                      ZP ejF         ej1        d           d!          ZQ G d" d#e(          ZR G d$ d%eR          ZSd&ejT        d'e!ddfd(ZUejV        W                    e7d)*          d&ejT        d'e!ddfd+            ZXd'e!ddfd,ZYd'e!ddfd-ZZd&ejT        d'e!ddfd.Z[d'e!ddfd/Z\d'e!ddfd0Z]d'e!ddfd1Z^d'e!ddfd2Z_d'e!ddfd3Z`ejV        a                    d4*          d'e!ddfd5            Zbd'e!ddfd6ZcejV        d                    d7d8d9gid:d9d;iid<d=d;iid9gd9d;id>d9gd=d;id?d9gd9d;id=d;id@g          d'e!d7eddfdA            Ze e ejG        e:jL        e:jf        e:jg        eAehg          B          dCee         ddfdD            Zi e ejG        e:jL        e:jf        e:jg        eAehg          B          dCee         ddfdE            Zj e ejG        e@ekg          F          dGee         ddfdH            ZldMdIZmd'e!ddfdJZnd'e!ddfdKZod'e!ddfdLZpdS )N    N)	overrides)givensettingsHealthCheck)DictSetcastUnionDefaultDictAnyList)	dataclass)ID
EmbeddingsIncludeIDsvalidate_embeddingsnormalize_embeddings)System)	ClientAPI)
Collection)
BundleRuleBasedStateMachineMultipleResultsrule
initializepreconditionconsumesrun_state_machine_as_testmultiple	invariant)defaultdict)is_client_in_processNOT_CLUSTER_ONLYcreate_isolated_database)wait_for_version_increaseget_collection_versionc                      dS Nr    r*       c/root/projects/butler/venv/lib64/python3.11/site-packages/chromadb/test/property/test_embeddings.py<lambda>r-   3   s    A r+   tracesi,  keyreturnc                 0    t           | xx         dz  cc<   d S )N   )r.   r/   s    r,   tracer4   9   s    
3KKK1KKKKKr+   c                  l    t                                           D ]\  } }t          |  d|            d S )Nz: )r.   itemsprint)r/   values     r,   print_tracesr9   >   sF    llnn ! !
U    ! !r+   dtyper3   dtype_shared_st   i   )	min_value	max_value	dimensiondimension_shared_stc                   "    e Zd ZdZdZdZdZdZdS )EmbeddingStateMachineStatesr   add_embeddingsdelete_by_idsupdate_embeddingsupsert_embeddingsN)__name__
__module____qualname__r   rC   rD   rE   rF   r*   r+   r,   rB   rB   M   s+        J%N#M++r+   rB   T)with_hnsw_paramscollc            
           e Zd ZU eed<    ed          Zee         ed<   dZd'de	de
f fdZ ee          dej        fd            Zed(d            Z ee ej        e                    dej        d	ee         fd            Z e ej         ee          d                    ded	d
fd            Z ed            e ej        eedd                    dej        d	d
fd                        Z ed            e ej        e ej        eej                  dd                    dej        d	d
fd                        Z e             d(d            Z! e             d(d            Z" e             d(d            Z# e             d(d            Z$ ed            e             d(d                         Z%d	e
fd!Z&dej        d	d
fd"Z'd#e(e)         d	d
fd$Z*d%e+d	d
fd&Z, xZ-S ))EmbeddingStateMachineBase
collectionembedding_idsFclient
use_searchc                     t                                                       || _        || _        t          j                            |           | _        d S N)super__init__rP   rQ   
hypothesisstatefulRuleStrategy_rules_strategyselfrP   rQ   	__class__s      r,   rU   z"EmbeddingStateMachineBase.__init__^   sF    $)2??EEr+   rN   c                    | j                             |j        |j        |j                  | _        |j        | _        t          d           |                     t          j	                   t          j        g g g g           | _        d S )N)namemetadataembedding_functioninitids	metadatas	documents
embeddings)rP   create_collectionr_   r`   ra   rN   r4   on_state_changerB   r   
strategiesStateMachineRecordSetrecord_set_state)r[   rN   s     r,   r   z$EmbeddingStateMachineBase.initialized   s    +77()< 8 
 

 #-"?f8CDDD * @bB2!
 !
 !
r+   r0   Nc                 N    | j                             | j        j                   d S rS   )rP   delete_collectionrN   r_   r[   s    r,   teardownz"EmbeddingStateMachineBase.teardowns   s#    %%do&:;;;;;r+   target
record_setrs   c                    t          d           |                     t          j                   t	          j        |          t          d                   dk    rt          d           t          d                                       | j	        d                   }t          |          dk    rt          t          d                                       |                    }fd|D             }t          j        fd|D             d         rfd|D             nd d	         rfd
|D             nd d         rfd|D             nd           } | j        j        di  |                     t#          t          j        |                     t'          |d          S  | j        j        di  |                     t#          t          j                             t'          d          S )NrC   rd   r   add_more_embeddingsc                 F    g | ]}d                               |          S rd   )index).0idnormalized_record_sets     r,   
<listcomp>z<EmbeddingStateMachineBase.add_embeddings.<locals>.<listcomp>   s-    PPP",U399"==PPPr+   c                 ,    g | ]}d          |         S rw   r*   ry   ir{   s     r,   r|   z<EmbeddingStateMachineBase.add_embeddings.<locals>.<listcomp>   s$    JJJQ.u5a8JJJr+   re   c                 ,    g | ]}d          |         S )re   r*   r~   s     r,   r|   z<EmbeddingStateMachineBase.add_embeddings.<locals>.<listcomp>   $    VVV4[A!DVVVr+   rf   c                 ,    g | ]}d          |         S )rf   r*   r~   s     r,   r|   z<EmbeddingStateMachineBase.add_embeddings.<locals>.<listcomp>   r   r+   rg   c                 ,    g | ]}d          |         S )rg   r*   r~   s     r,   r|   z<EmbeddingStateMachineBase.add_embeddings.<locals>.<listcomp>   s$    XXX1 5l CA FXXXr+   rc   r*   )r4   ri   rB   rC   
invariantswrap_alllensetintersectionrl   list
differencerj   NormalizedRecordSetrN   add_upsert_embeddingsr	   	RecordSetr    )r[   rs   r   new_idsindicesfiltered_record_setr{   s         @r,   rC   z(EmbeddingStateMachineBase.add_embeddingsw   sE   
 	8GHHH@J@SA
 A
 $U+,,q00'(((0788EE!%(
 
 |q  34U;<<GGUUVVGPPPPPPPG.JJJJ'JJJ,[9VVVVgVVVV,[9VVVVgVVVV,\: XXXXPWXXXX      DO88"7888##D)=?R$S$STTT0788  DO88"7888##D)=?T$U$UVVV259::r+   r2   min_sizerw   rd   c                      t          d                                t          j                    fd|D             } j                            |                                t          |                     d S )Nzremove embeddingsc                 P    g | ]"}j         d                              |          #S rw   )rl   rx   )ry   rz   r[   s     r,   r|   z;EmbeddingStateMachineBase.delete_by_ids.<locals>.<listcomp>   s/    RRRT259??CCRRRr+   rw   )r4   ri   rB   rD   rN   delete_remove_embeddingsr   )r[   rd   indices_to_removes   `  r,   rD   z'EmbeddingStateMachineBase.delete_by_ids   s    !"""8FGGGRRRRcRRR3'''$5 6 677777r+   c                 >    t          | j        d                   dk    S Nrd      r   rl   ro   s    r,   r-   z"EmbeddingStateMachineBase.<lambda>       s4#8#?@@AE r+   r   collection_strategyid_strategyr   max_sizers   c                     t          d           |                     t          j                    | j        j        di | |                     |           d S )Nzupdate embeddingsr*   )r4   ri   rB   rE   rN   updater   r[   rs   s     r,   rE   z+EmbeddingStateMachineBase.update_embeddings   `     	!"""8JKKK,,,,,
+++++r+   c                 >    t          | j        d                   dk    S Nrd      r   ro   s    r,   r-   z"EmbeddingStateMachineBase.<lambda>   r   r+   c                     t          d           |                     t          j                    | j        j        di | |                     |           d S )Nzupsert embeddingsr*   )r4   ri   rB   rF   rN   upsertr   r   s     r,   rF   z+EmbeddingStateMachineBase.upsert_embeddings   r   r+   c                 t    t          j        | j        t          t          j        | j                             d S rS   )r   countrN   r	   rj   r   rl   ro   s    r,   r   zEmbeddingStateMachineBase.count   s8    OT*"68MNN	
 	
 	
 	
 	
r+   c                 8    t          j        | j                   d S rS   )r   no_duplicatesrN   ro   s    r,   r   z'EmbeddingStateMachineBase.no_duplicates   s     11111r+   c                     t          j        | j        t          t          j        | j                  d| j        | j                   d S )Ngffffff?)rN   rs   
min_recallra   rQ   )	r   ann_accuracyrN   r	   rj   r   rl   ra   rQ   ro   s    r,   r   z&EmbeddingStateMachineBase.ann_accuracy   sL    J0$2GHH#6	
 	
 	
 	
 	
 	
r+   c                 J   |                                  r!| j                                        dk    sJ d S t          t          j        | j                  }t          j        | j        |           t          j	        | j        |           t          j
        | j        |           d S r)   )_is_state_emptyrN   r   r	   rj   r   rl   r   embeddings_matchmetadatas_matchdocuments_match)r[   rl   s     r,   fields_matchz&EmbeddingStateMachineBase.fields_match   s    !! 		J?((**a//////  $J$8$:OPP'9IJJJ&t8HIII&t8HIIIIIr+   c                 *    t          | j                  S rS   )r#   rP   ro   s    r,   r-   z"EmbeddingStateMachineBase.<lambda>   s    )$+66 r+   c                 `    | j         j        }t          j        || j        g| j                   d S rS   )rP   _systemr   log_size_below_maxrN   has_collection_mutated)r[   systems     r,   r   z,EmbeddingStateMachineBase.log_size_below_max   s<    
 ,%T_%t'B	
 	
 	
 	
 	
r+   c                 H    | j                                         D ]}|r dS dS )NFT)rl   values)r[   fields     r,   r   z)EmbeddingStateMachineBase._is_state_empty   s9    *1133 	 	E uutr+   c           	         t          j        |          }t          |d                   D ]\  }}|| j        d         v r]| j        d                             |          }|d         |d         |         | j        d         |<   nI|d         J | j        J |                     |d         |         g          d         | j        d         |<   |d         | j        d         |         }|pt          t          t          t          t          t          t          f         f         |          }|d         |         "|                    |d         |                    nn|d         |         | j        d         |<   |d         |d         |         | j        d         |<   s| j        d                             |           |d         -| j        d                             |d         |                    nY| j        J |d         J | j        d                             |                     |d         |         g          d                    |d         -| j        d                             |d         |                    n | j        d                             d            |d         .| j        d                             |d         |                    | j        d                             d            d S )Nrd   rg   rf   r   re   )r   r   	enumeraterl   rx   ra   r	   r   strr
   intfloatr   append)r[   rs   r{   idxrz   
target_idxrl   s          r,   r   z,EmbeddingStateMachineBase._upsert_embeddings  s   @J@SA
 A
 !!6u!=>> C	D C	DGCT*5111!259??CC
(6B .l;C@ ),7"  1=III2>>> //.{;C@A   ),7" )5A (,'<['I*'U$'3+/ eCeO&<!<=?O, ,( 1=cBN,33 5k B3 G   
 ! 2+>sC -k:& )5A .k:3? )+6"
 %e,33B777(6B),7>>-l;C@     2>>>0=III),7>>//2;?DE   
 )5A)+6==-k:3?    )+6==dCCC(5A)+6==-k:3?    )+6==dCCCCGC	D C	Dr+   r   c                     t          |          }|                    d           |D ]:}| j        d         |= | j        d         |= | j        d         |= | j        d         |= ;d S )NT)reverserd   rg   re   rf   )r   sortrl   )r[   r   indices_listr   s       r,   r   z,EmbeddingStateMachineBase._remove_embeddingsJ  s    -..$''' 	6 	6A%e,Q/%l3A6%k215%k2155		6 	6r+   	new_statec                 8    |t           j        k    r	d| _        d S d S )NT)rB   r   r   )r[   r   s     r,   ri   z)EmbeddingStateMachineBase.on_state_changeT  s'    3>>>*.D''' ?>r+   Fr0   N).rG   rH   rI   r   __annotations__r   rO   r   r   r   boolrU   r   collection_strj   r   rp   r   
recordsetsr   r   rC   stlistsr   r   rD   r   rE   one_of	safe_textrF   r!   r   r   r   r   r   r   r   r   r   r   r   ri   __classcell__r\   s   @r,   rM   rM   Y   s&         & 7 7M6":777"F Fy Fd F F F F F F Z=)))
Z%: 
 
 
 *)
 < < < Y< 
T(:(77  ';)= ';/RTBU '; '; ';	 ';R 
Thbhxx..;;;<<<8 8 8 8 8 =<8 \EEFF	T(:( -%	
 
 
  ,J,@ ,T , , ,  GF, \EEFF	T(:( -!	-1EFF	
 
 
  ,J,@ ,T , , ,  GF, Y[[
 
 
 [

 Y[[2 2 2 [2 Y[[
 
 
 [
 Y[[
J 
J 
J [
J \66  Y[[
 
 
 [ 
    GDZ-A GDd GD GD GD GDR6CH 6 6 6 6 6/ / / / / / / / / /r+   rM   c            
           e Zd ZU  ed          Zee         ed<   ddedef fdZ	 e
e          dej        f fd            Z
 ed	            e            dd                        Z ee ej        e                    dej        d
ee         f fd            Z e ej         ee          d                    ded
df fd            Z ed            e ej        eedd                    dej        d
df fd                        Z ed            e ej        e ej        eej                  dd                    dej        d
df fd                        Z xZS )EmbeddingStateMachinerO   FrP   rQ   c                 L    t                                          ||           d S rS   )rT   rU   rZ   s      r,   rU   zEmbeddingStateMachine.__init__\  s#    ,,,,,r+   r]   rN   c                    t                                          |           t          d| j        j        d         d|j                   d| _        t                      | _        | j        	                                d         | _
        d S )Nz7[test_embeddings][initialize] Initialize collection id rz   z$ hypothesis generated collection id r   version)rT   r   r7   rN   _modelrz   log_operation_countr   unique_ids_in_log	get_modelcollection_version)r[   rN   r\   s     r,   r   z EmbeddingStateMachine.initialize_  s|    :&&&EO"4(2M		
 	
 	
 $% *-%%"&/";";"="=i"Hr+   c                 X    t            o"| j        dk    ot          | j                  dk    S )N
   r   )r$   r   r   r   ro   s    r,   r-   zEmbeddingStateMachine.<lambda>m  s5    )) ,$r),&''!+ r+   r0   Nc                    t          | j        | j        j                  }|| j        k    sJ || j        k    r2t          d           || _        d| _        t                      | _        d S t          d|           t          | j        | j        j        |t                    }d| _        t                      | _        || _        d S )NzT[test_embeddings][wait_for_compaction] collection version has changed, so reset to 0r   zY[test_embeddings][wait_for_compaction] wait for version to increase from current version )additional_time)r'   rP   rN   r_   r   r7   r   r   r   r&   VERSION_INCREASE_WAIT_TIME)r[   current_versionnew_versions      r,   wait_for_compactionz)EmbeddingStateMachine.wait_for_compactionl  s     1do>RSS$"99999 T444f   '6D#'(D$%(UUD"""k   4$ :	  K ()D$%(UUD"&1D###r+   rq   rs   c           	      j   t                                          |          }t          j        |          }t	          d|d         dt          |d                              | xj        t          |d                   z  c_        |d         D ]%}|| j        vr| j                            |           &|S )Nz,[test_embeddings][add] Non Intersection ids rd    len )	rT   rC   r   r   r7   r   r   r   r   )r[   rs   resr{   rz   r\   s        r,   rC   z$EmbeddingStateMachine.add_embeddings  s    
 gg$$Z00@J@SA
 A
 	:!%(%e,--		
 	
 	
 	  C(=e(D$E$EE  '. 	/ 	/B///&**2...
r+   r2   r   rw   rd   c                    t                                          |           t          d|dt          |                     | xj        t          |          z  c_        |D ]%}|| j        v r| j                            |           &d S )Nz[test_embeddings][delete] ids r   )rT   rD   r7   r   r   r   remove)r[   rd   rz   r\   s      r,   rD   z#EmbeddingStateMachine.delete_by_ids  s    c""".Wc#hhGGG  CHH,   	2 	2BT+++&--b111	2 	2r+   c                 >    t          | j        d                   dk    S r   r   ro   s    r,   r-   zEmbeddingStateMachine.<lambda>  r   r+   r   r   r   c                 .   t                                          |           t          d|d         dt          t	          j        |d                                        | xj        t          t	          j        |d                             z  c_        d S )Nz[test_embeddings][update] ids rd   r   )rT   rE   r7   r   r   wrapr   )r[   rs   r\   s     r,   rE   z'EmbeddingStateMachine.update_embeddings  s     	!!*---,u

5 12233		
 	
 	
 	  C

58I(J(J$K$KK    r+   c                 >    t          | j        d                   dk    S r   r   ro   s    r,   r-   zEmbeddingStateMachine.<lambda>  r   r+   c                    t                                          |           t          d|d         dt          t	          j        |d                                        | xj        t          t	          j        |d                             z  c_        t	          j        |d                   D ]%}|| j        vr| j                            |           &d S )Nz[test_embeddings][upsert] ids rd   r   )	rT   rF   r7   r   r   r   r   r   r   )r[   rs   rz   r\   s      r,   rF   z'EmbeddingStateMachine.upsert_embeddings  s     	!!*---,u

5 12233		
 	
 	
 	  C

58I(J(J$K$KK  /*U"344 	/ 	/B///&**2...	/ 	/r+   r   r   )rG   rH   rI   r   rO   r   r   r   r   rU   r   r   rj   r   r   r   r   r   r   r   rC   r   r   r   r   rD   rE   r   r   rF   r   r   s   @r,   r   r   Y  s         & 7 7M6":777- -y -d - - - - - - Z=)))
IZ%: 
I 
I 
I 
I 
I *)
I \	, 	, 
 
TVV2 2 2 V 2: 
T(:(77  )= /RTBU     	 " 
Thbhxx..;;;<<<2 2 2 2 2 2 2 =<2 \EEFF	T(:( -%	
 
 
  LJ,@ LT L L L L L  GFL \EEFF	T(:( -!	-1EFF	
 
 
  /J,@ /T / / / / /  GF/ / / / /r+   r   caplogrP   c                     t                     |                     t          j                   t	          fdt          dt          j        g                     t                       d S )Nc                  "    t                     S rS   r   rP   s   r,   r-   z'test_embeddings_state.<locals>.<lambda>  s    %f-- r+   _ deadlinesuppress_health_checkr   	r%   	set_levelloggingERRORr   r   r   filter_too_muchr9   r   rP   s    `r,   test_embeddings_stater    sv    V$$$
W]###----;3N2O
 
 
    NNNNNr+   z-Search API only available in distributed mode)reasonc                     t                     |                     t          j                   t	          fdt          dt          j        g                     t                       dS )z@Test embeddings state machine using search API instead of query.c                  &    t           d          S )NT)rQ   r   r   s   r,   r-   z3test_embeddings_state_with_search.<locals>.<lambda>  s    %f>>> r+   r   r   r  Nr  r  s    `r,   !test_embeddings_state_with_searchr    sx     V$$$
W]###>>>>;3N2O
 
 
    NNNNNr+   c                    t          |            t          |           }|                    t          j        dddddd t          j                    dt          j        i g dd
  
                   |	                                 |
                                 |                                 |                                 |                    g d	d
dgddgddgddgddgddggg dd d          \  }}}}}}|	                                 |
                                 |                                 |                                 |                    |||||g           t          s|                                 |	                                 |
                                 |                                 |                                 |                                 d S )NA00   hnsw:construction_efhnsw:search_efhnsw:Mr<   FT
r_   r`   ra   rz   r?   r:   known_metadata_keysknown_document_keywordshas_documentshas_embeddingsr]   )012345      ?     ?     P?     ?     ÿ     ?          ?     ?     Ϳ)NNNNNNrd   rg   re   rf   r   rw   )r%   r   r   rj   r   uuiduuid4npfloat16r   r   r   r   rC   rD   r$   r   rp   )rP   statev1v2v3v4v5v6s           r,   test_add_then_delete_n_minus_1r5    s$   V$$$!&))E	((+"% 
  $z||* "$&
 
 
    $ 
	KKMMM		"11111^, 01!>2"M200 >==
 
 2  BBB 
	KKMMM			RRR0111 $!!###		KKMMM			NNr+   c                    t          |            t          |           }|                    t          j        ddddddddddd t          j        d	          d
t          j        i g dd
  
                   |	                                 |
                                 |                                 |                                 |                    g ddgd
z  dgd
z  dgd
z  dgd
z  dgd
z  dgd
z  gddddddddd dddddddddddddd dd!d"d#
d$d%d&d'd(d)d*dd+dd,
gd d-.          \  }}}}}}|	                                 |
                                 |                                 |                                 |                    d/dgd
z  gd d d-.          \  }|	                                 |
                                 |                                 |                                 |                    d0d1gdgd
z  dgd
z  gd d2dd3dd4gd d-.          \  }	}
|	                                 |
                                 |                                 |                                 |                    g d5dgd
z  dgd
z  dgd
z  gd6d7d8d9d:d;d<dd=d>d?id@ddAdBdddCdDdddEdFgd d-.           |	                                 |
                                 |                                 |                                 |                    dGdHdI||	gdgd
z  dgd
z  dgd
z  dgd
z  dgd
z  gdJdKddLdMdNdOdPdQdNdRdSdTdUddVdWdXdYdZddd[d\d gd d-.           |	                                 |
                                 |                                 |                                 |                    |||
||gdgd
z  dgd
z  dgd
z  dgd
z  dgd
z  gi d]dd^d_d`ddaddbdcdddedfddgdhdiddjddedkdldmdnddodpdqddrdsdtduddvdwdxdyd dzd{dd|d}dNd~dddddddd*d7ddgd d-.           |	                                 |
                                 |                                 |                                 |                    |||
|gdgd
z  dgd
z  dgd
z  dgd
z  gddddd|dddddiddidNddddddSddgd d-.           |	                                 |
                                 |                                 |                                 |                    |gdgd
z  gddlddddddd d-.           |	                                 |
                                 |                                 |                                 |                    g ddgd
z  dgd
z  dgd
z  dgd
z  gdddddddiddhdddigd d-.          \  }}}}|	                                 |
                                 |                                 |                                 |                    ddd|gdgd
z  dgd
z  dgd
z  dgd
z  gddid(dpdddddd7dddkddddddd
djdigd d-.           |	                                 |
                                 |                                 |                                 |                    ||gdgd
z  dgd
z  gi ddddddddddddddddNddddhddddd|ddddCdddddēddddddddddʜ	ddddddkd͜gd d-.           |	                                 |
                                 |                                 |                                 |                    |gdgd
z  gd d d-.           |	                                 |
                                 |                                 |                                 |                    |||||gdgd
z  dgd
z  dgd
z  dgd
z  dgd
z  gddidddќd dddddddd؜ddddddddޜgd d-.           |	                                 |
                                 |                                 |                                 |                    ||	dd|gdgd
z  dgd
z  dgd
z  dgd
z  dgd
z  gd ddddddddddddkdddddddddddddddgd d-.           |	                                 |
                                 |                                 |                                 |                                 d S )NfOIByFHtmYRTipiU5KdoaYsUQ_EWStV4r  )z-7n92WhVE_zJ-sW
wPGA8hY7uX4rAr  r  r  z$ff006990-82c3-494b-97d5-cbb05092c861i  Tr  r]   )kgaT4dzC2h2YoNSgUqRyE-Tmxf3MTzODI-yOtbvCr   i{5DkqA9o6           g   `UUտVT)szd1wQJV-9hWf7gwQrbyHgPe0r6qQ5XYxeqPzXpiqBi|E)hqTZ6Ok767eCSwyvGEuig8aTRGxN1h8Ii)ATRsKF0Piig      vReDiig      <)
PcFwuPSWCgxEQtupcOfhueUzI4tOMfHM7uU58tW_muctZfSvOyiuDCOcywKhsii;[g      p>-g    .i)
iuTAKznMg6IdUKxaPioyzc0Zb3VTUktBu-uW6innbsomMXnpsEEEFul8JRjHi7lrYr)  r   uaK_yFsHiO5g      )
RiaaN9MNpqg9WxuexOHh2)SCeelWyLAWG_oHalYr  iObg   >   ii     g    >A)0ZbYq40POT9sTxkMz-jrqsBEfrELJctJoVeLqtsPZp5M4XFtiilyDy6q7LZUWzKuQG8m-Ti  xGwkig   ;*A)fPAPtmtxKb6CG18VrF  zc-G92nYTTBWs31rbM_L_PQDSCuxOGzFeGgh7cuT_ruA3mnzO3m3-X1ZNt2PF6M5_qIj0Yh6iKiM  34oWsr93EUl    ~.Ag   Xc80eze-QXuDdnxYMd)2fDAuv74Etz5hqGH60G-yZ6PWyM1BOkMjjGyTkwZvpmS5HbTAIbvqzPhLgg   `ff?)WKyGNZphPCKay88gsh3x_iJDcig)z_zVO2i-NtWHxoltu_E_fg9yGpikZnQ3ALwA2   iL)otM8EGeKm
skf71O0UKTzS8Kc8-l95Rpc4bGz1QmzbKVySN1yrXFl56CmDS08FhBFXAIAWx4dcB5i8w8mwQ5c7G9g2JVYzVQGb_r-hzoAg    M0lMigFg   ?dRDugrcoZv45Bg   @  zUG-sSVcri4cT1GidIii
xM0jUV4U2rmsuFYMwj_SXc)FqFWR__r  vwo4WERBljY3aWjWnqLWFib@iHr  )mf4UtV_NVCw   )VWuP7uF8r  g      iL   )a1ML_Zl2Ir85KolESaXiJvAO8orMSr  5AeE2q	6T9tEEC84i  )YxlGe5GMXUvsAljL5V5ELRv	yeLTrhAIq      ?iP      r  Rii&  ip )pzHn2YfdftMEd0C5ekByb7mhdbLJCViuLT5YOoVwEvq8Wf6iŭebEKOxi	g   uA)C1KbOOlKkzzLo9CGU2NHMfpu77F70Iclz	dz6fI-GppqVVWQrcq645F)F7Rig1RXiz_nC8-ii  g   @̼+A)FBtaPcQWV24vddLq1My3mbUL9IfIHLxuosTATUP1)AhCwm9AwPKim}4GJrAunf5
J8O0R8VGaYnK67gg    .Ag   `\A)K2cCgoObAcpaxafzpxnRCSPJUF4JZ2sKOIRDcznBaQ6F1O38etVMhss2angu-g   g    #
AiSNcMrT   )
UwbDWM2_r  nfoovt
elaMLbhEvWr  UOr  3ssGS4rSKXsKqRPFTBGrRPPsuGwVtTwei_qscyRBC_TPIXd	N0FG7Nta1g   >&woDB   IV0L3xImGgz9N--JBl0uH_au_KVmhtcAqrNfL6taIVpCXJXl
8YaEynJznBiONg   $AebBi  g   
@	cRl59jW_Og    ~.g      8g   @w)	fJz
QK8l3l4yP-r  z	-XP899RRnA9UlxNwmcG1NoCdWRn5cD	HExeVM0iMiSh)zAbCKkEvE4shnFNUc9vXWYz-iRt88MxYq)55m28krQsTFdqMhYjhFiHg      ?VG7E7xmio#  1xXUHLklmIVSVgQd7EHUCu5wakl6)9SnviLf
Y0Jw4pLTwrv3E
Fx3jsbcdqyHz-USisBi
xtLNrNyRo2g    ),)rY  Qfm_6duLSh0LkduZt5qsRJrF8DMMZrq  7OiaiD4eKi٨ii])tVsBzTRzbq6VslBBo2_12hgyKNPddxify34-np-F7FcZpODwCTHg91o4mKTjBL1Zjfyslg3z_TG-YhAQQ)1qrTKV8tLuQHsxar  UepiV6K_g    YR)prRUDZ_uRfFG6cZvICaGczunTbxz0qd2-AV1iC9OePgaBYi%D)EXXVBZUtJMO4oF8g8nr  hv3i)r%   r   r   rj   r   r*  UUIDr,  r-  r   r   r   r   rC   rF   rE   rp   )rP   r.  embedding_ids_0embedding_ids_1embedding_ids_2embedding_ids_3embedding_ids_4embedding_ids_5embedding_ids_6embedding_ids_7embedding_ids_8embedding_ids_9embedding_ids_10embedding_ids_11embedding_ids_12s                  r,   test_embeddings_flake1r[  -  su   V$$$!&))E	(!")(+"%	 	  $y?@@* "$&'
 
 
    . 
	KKMMM		 	QQQc	c	c	c	c	c	  .) #6#  /90 
 -f==#" #!*-1  +1'1 /$ $)" ;)T i5
 5
 	 7	 7	r 
	KKMMM		--39+	
 
 .  _ 
	KKMMM		',';';&>39qcCi0", 3	  
 
 (< ( ($O_  
	KKMMM			11139qcCi!s; !* 5/3 #	 	 !$"*!!,1#%. 4 ;
 
       B 
	KKMMM			  39qcCi!sQC#IsSyQ  & */+#$:"*	 	 1*; 
 !5etTT&"&$(5B  14 G$
 $
  & & &N 
	KKMMM			   39qcCi!sQC#IsSyQus % 	
 D C $ "#6 d  *   . e  !" #$  %+7(+  . IUcJJ/5II*). (#1	 	7%L _0
 0
  2 2 2f 
	KKMMM			#_oW39qcCi!sQC#IF"t44$d#U$	 d3-1'"+!	 	( /
 
    6 
	KKMMM			#$39+&4#!(  
 
      
	KKMMM		 	11139qcCi!sQC#IF %+&*$$	  $(;<<f
 
 
 	 	 	& 
	KKMMM			3939qcCi!sQC#IFt$21!%,0/<  !6$("3"1:  e-0 7
 
    > 
	KKMMM			#_539qcCi0* %  	
 2 * %d 5 % / d 3  !$   !" 5#$ &"'$!*0#"+5  : $(' &* ;%L S*
 *
  , , ,Z 
	KKMMM			#$39+	
 
     
	KKMMM			 !  39qcCi!sQC#IsSyQg!U;;/"&*"+5   %(,!&# . A!
 !
  # # #H 
	KKMMM			#_eT?S39qcCi!sQC#IsSyQ7=/4$%  &%.  $"$(&0   -$&" 3!D K&
 &
  ( ( (R 
	KKMMM			NNr+   c                 "   t          |           t          |          }|                    t          j        dddddd t          j        d          dt          j        i g dd
  
        	           |	                                 |
                                 |                                 |                                 |                    g d
ddgddgddgddgddggg dd d          \  }}}}}|	                                 |
                                 |                                 |                                 |                    |gddgg||igd d           |	                                 |                                 d S )Nr  r  r  z$2fb0c945-b877-42ab-9417-bfe0f6b172afr<   FTr  r]   )r  r  r  r  r  r  r   r!  r"  r#  r$  r%  r&  r'  r(  )NNNNNr)  r   g     ?g     ?)r%   r   r   rj   r   r*  rM  r,  r-  r   r   r   r   rC   rE   rp   )r   rP   r.  r/  r0  r1  r2  r3  s           r,   test_update_noner]  I  s   V$$$!&))E	((+"% 
  $y?@@* "$&
 
 
    $ 
	KKMMM		--,,,^, 01!>2"M20 877
 
 .  BBB 
	KKMMM			4)+;<=r(	
 
     
	NNr+   c                 P   t          |            t          |           }|                    t          j        di dddddddd	d
dddddddddddddddddddddddd t          j        d          dt          j        i g dd
  
                    |	                                 |
                                 |                                 |                                 |                    g d!g d"g d#g d$g d%g d&g d'gd d d()          }d*}i }|D ]}|||<   |d+z  }|	                                 |
                                 |                                 |                                 |                    |d*         |d,         d-|d.         d/gg d0g d1g d2g d3g d4gd d d5)           |	                                 |
                                 |                                 |                                 t          s|                                 |	                                 |
                                 |                                 |                                 |                    d6|d*         |d,         gg d7g d8g d9gd d d()           |	                                 |
                                 |                                 |                                 |                    |d:         |d+         |d.         |d,         |d         g;           |	                                 |
                                 |                                 |                                 |                    d<d=d>d?|d+         gg d@g dAg dBg dCg dDgd d d5)          }dE}|D ]}|||<   |d+z  }|	                                 |
                                 |                                 t          s|                                 d S d S )FNKR3cfUfmxsi3r  bMMvvrqM4MKmp5CJB8Ai  r_  T37PNiVkn5KZfkpod3ND5soL_KA4zcZL9lRN9   Oc8G7ysXmE8lp4Hos_POQe8Unz1uJBI930U   teFtyMr  R0ZiZr  IOwir  r  r  r  z$284b6e99-b19e-49b2-96a4-a2a93a95447dr   r  r]   )255r  z3-r   Nk9yPvT)g333333?gffffff@r  )      @      @r<   )r2   r<   r   )   r      )g!@g      "@   )rr  rs  ffffff@r)  r   r   r2   rt  KWcDaHUVD6MxEiJr   PdlP1d6w)r2      rt  )r   r   	   )r{  r   r   )r   r{     )r2   r   rt  )rd   rg   rf   re   TpjiboLSuYWBJDbRW1zeNmC)rt  ru  rv  )rv  r{  r   )r2   r   ru  r<   rw   oD3V84RtTDwlc9C8_evn)r{  g@g(\	@)gHzG?gQ@gp=
ף@)g#@g#@g?)gffffff#@rw  rr  )g333333@rw  gL#@ru  )r%   r   r   rj   r   r*  rM  r,  float32r   r   r   r   rC   rF   r$   r   rD   )rP   r.  rg   r   emb_list	embeddings         r,   test_add_delete_addr    sy   V$$$!&))E	(8%u T 	
 #D  %m " e t  T v ' !#  #!$  $y?@@* "$&7
 
 
    > 
	KKMMM		%%999				 
 
 &  J 	
AH  		Q		KKMMM			 ! &::yyy)))YYY			R
 
     
	KKMMM		 $!!###		KKMMM			-x{HQKH$99iii;	
 
     
	KKMMM			a[(1+x{HQK!M     
	KKMMM		%%$D"""       
 
 &  J 	
A  		Q		KKMMM	 $!!#####$ $r+   c                    t          |            |                     d          }|                    dgdgg           |                                dk    sJ |                    dgdgg           |                                dk    sJ |                                }|d         dgk    sJ |                    dg           |                                d	k    sJ d S )
Nfoor_   arC  rd   rg   r2   rd   rw   r   )r%   rh   r   r   getr   )rP   rK   resultss      r,   test_multi_addr    s    V$$$###//DHH#SE7H+++::<<1 	HH#SE7H+++::<<1hhjjG5>cU""""KKSEK::<<1r+   c                    t          |            |                     d          }t          j        t          j                  5  |                    ddgdgdgg           d d d            n# 1 swxY w Y   t          j        t          j                  5  |                    ddgdgdgg           d d d            d S # 1 swxY w Y   d S )Nr  r  r  rC  g?r  )r%   rh   pytestraiseserrorsDuplicateIDErrorr   r   rP   rK   s     r,   test_dup_addr    sI   V$$$###//D	v.	/	/ < <c3ZcUSEN;;;< < < < < < < < < < < < < < <	v.	/	/ ? ?c
u~>>>? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ?s$   A..A25A2CC	Cc                    t          |            |                     d          }g d}t          j                            dd          }t          j                            dd          }t          j                            ||f                                          }|                    t          t          |          |          }|D ])}||         }|J t          d |D                       sJ *d S )	Nr  r  )rf   re   rg   	distancesr2   i  d   )query_embeddingsincludec                 4    g | ]}t          |          d k    S )r   )r   )ry   results     r,   r|   z*test_query_without_add.<locals>.<listcomp>+  s#    AAACKK1$AAAr+   )
r%   rh   r,  randomrandinttolistqueryr	   r   all)	rP   rK   fieldsr  r  r  r  r   field_resultss	            r,   test_query_without_addr    s   V$$$###//DKKKF
	!T""A
	!S!!Ay''A//6688jjj*:;;V   G  C C(((AA=AAABBBBBBC Cr+   c                 \   t          |            |                     d          }|                    dgg d          }t          |d                   dk    sJ t          |d                   dk    sJ t          |d	                   dk    sJ t          |d
                   dk    sJ d S )Nr  r  r  )rf   re   rg   )rd   r  rd   r   re   rf   rg   )r%   rh   r  r   )rP   rK   r  s      r,   test_get_non_existentr  .  s    V$$$###//DXX3%)Q)Q)QXRRFve}""""vk"##q((((vk"##q((((vl#$$))))))r+   z9We don't properly escape SQL internally, causing problemsc                    t          |            d}|                     d          }|                    |gdgg           |                                dk    sJ |                    |g           |                                dk    sJ d S )	Nr  r  rC  r  r2   rw   r   )r%   rh   r   r   r   )rP   rz   rK   s      r,   test_escape_chars_in_idsr  9  s    V$$$	B###//DHH"C5'H***::<<1KKRDK::<<1r+   c                     t          |            |                     d          }t          j        t                    5  |                                 d d d            d S # 1 swxY w Y   d S )Nr  r  )r%   rh   r  r  
ValueErrorr   r  s     r,   test_delete_empty_failsr  D  s    V$$$###//D	z	"	"                   s   A!!A%(A%kwargsrd   r  wherebarwhere_documentz	$contains)rd   r  )rd   r  )rd   r  r  c                 j    t          |            |                     d          } |j        di | d S )Nr  r  r*   )r%   rh   r   )rP   r  rK   s      r,   test_delete_successr  K  sF      V$$$###//DDK&r+   )supported_typesr  c                     t          j        dd|           }t          t          t          t          |                              }t          d |D                       sJ d S )Nr   c                     g | ]p}t          |t          j                  oS|j        t          j        k    p>|j        t          j        k    p)|j        t          j        k    p|j        t          j        k    qS r*   
isinstancer,  ndarrayr:   r  float64int32int64ry   r8   s     r,   r|   zMtest_autocasting_validate_embeddings_for_compatible_types.<locals>.<listcomp>m  |     		
 		
 		
  ubj)) rz) +;"*,+;"(*+ ;"(*		
 		
 		
r+   )rj   create_embeddingsr   r	   r   r   r  r  embdsvalidated_embeddingss      r,   9test_autocasting_validate_embeddings_for_compatible_typesr  a  s     (RAAE. ''	
 	
  		
 		
 .		
 		
 		
      r+   c                     t          j        dd|           }t          t          t          t          |                              }t          d |D                       sJ d S )Nr   c                     g | ]p}t          |t          j                  oS|j        t          j        k    p>|j        t          j        k    p)|j        t          j        k    p|j        t          j        k    qS r*   r  r  s     r,   r|   zEtest_autocasting_validate_embeddings_with_ndarray.<locals>.<listcomp>  r  r+   )rj   create_embeddings_ndarrayr   r	   r   r   r  r  s      r,   1test_autocasting_validate_embeddings_with_ndarrayr  z  s     0RIIE.Z-e4455  		
 		
 .		
 		
 		
      r+   )unsupported_typesr  c           	      "   t          j        dd|           }t          j        t                    5 }t          t          t          t          |                               d d d            n# 1 swxY w Y   dt          |j
                  v sJ d S )Nr   zsExpected embeddings to be a list of floats or ints, a list of lists, a numpy array, or a list of numpy arrays, got )rj   r  r  r  r  r   r	   r   r   r   r8   )r  r  rX  s      r,   7test_autocasting_validate_embeddings_incompatible_typesr    s     (R1BCCE	z	"	" KaD-A%-H-HIIJJJK K K K K K K K K K K K K K K 	~qw<<	 	 	 	 	 	s   0A,,A03A0c                      t          j        g           g} t          j        t                    5 }t          |            d d d            n# 1 swxY w Y   dt          |          v sJ d S )NzExpected each embedding in the embeddings to be a 1-dimensional numpy array with at least 1 int/float value. Got a 1-dimensional numpy array with no values at pos)r,  arrayr  r  r  r   r   )r  rX  s     r,   test_0dim_embedding_validationr    s    "E	z	"	" #aE"""# # # # # # # # # # # # # # # 	mq66	 	 	 	 	 	s   AAAc           	      N   t          |            |                     d          }t          | |j                  }t	          ddd          D ]3}|                    d t	          ||dz             D                        4t          st          | |j        |t                     d S d S )Nnoopr  r     r  c                 ,    g | ]}t          |          S r*   r   ry   r   s     r,   r|   z)test_no_op_compaction.<locals>.<listcomp>  s    CCCAQCCCr+   rw   )	r%   rh   r'   r_   ranger   r$   r&   r   )rP   rK   initial_versionbatchs       r,   test_no_op_compactionr    s    V$$$###00D,VTY??Oq$$$ E ECCueck)B)BCCCDDDD 
!DI0J	
 	
 	
 	
 	

 
r+   c                    t          |            d}d}|                     d          }t          | |j                  }t	          d||          D ]L}d t	          |||z             D             }d |D             }|                    |d |D             	           Mt          st          | |j        |t                     t          | |j                  }t	          d||          D ]A}d
 t	          |||z             D             }d |D             }|	                    |           Bt          st          | |j        |t                     t          |                                d                   dk    sJ d S )Nr  r  add_then_purger  r   c                     g | ]}|S r*   r*   r  s     r,   r|   z'test_add_then_purge.<locals>.<listcomp>      GGG!GGGr+   c                 ,    g | ]}t          |          S r*   r  r  s     r,   r|   z'test_add_then_purge.<locals>.<listcomp>      555c!ff555r+   c                 (    g | ]}d |z  d |z  dz   gS )r<   r2   r*   r  s     r,   r|   z'test_add_then_purge.<locals>.<listcomp>  s)    'S'S'SqQA	(:'S'S'Sr+   r  c                     g | ]}|S r*   r*   r  s     r,   r|   z'test_add_then_purge.<locals>.<listcomp>  r  r+   c                 ,    g | ]}t          |          S r*   r  r  s     r,   r|   z'test_add_then_purge.<locals>.<listcomp>  r  r+   rw   rd   )r%   rh   r'   r_   r  r   r$   r&   r   r   r   r  )rP   record_countbatch_countrK   witness_versionr  record_id_vals
record_idss           r,   test_add_then_purger    s   V$$$LK##)9#::D,VTY??O q,44 
 
GGU5%+2E%F%FGGG55n555
'S'SN'S'S'S 	 	
 	
 	
 	
  
!DI0J	
 	
 	

 -VTY??Oq,44 $ $GGU5%+2E%F%FGGG55n555

#### 
!DI0J	
 	
 	

 txxzz% !!Q&&&&&&r+   c                    t          |            |                     d          }t          | |j                  }d}d t	          ||dz             D             }d t	          ||dz             D             }|dz   }|                    ||           |                    |           t          st          | |j        |t                     t          | |j                  }d}d	 t	          |||z             D             }d
 t	          |||z             D             }|                    ||           |                    |           |                    ||           t          st          | |j        |t                     |
                                }t          |d                   |k    sJ |D ]}||d         v sJ d S )Nencompassing_deleter   c                 ,    g | ]}t          |          S r*   r  r  s     r,   r|   z,test_encompassing_delete.<locals>.<listcomp>  s    
9
9
9a3q66
9
9
9r+   ru  c                 "    g | ]}|d z  |d z  gS g      ?r*   r  s     r,   r|   z,test_encompassing_delete.<locals>.<listcomp>  s%    LLL1s7AG$LLLr+   r  rw      c                 ,    g | ]}t          |          S r*   r  r  s     r,   r|   z,test_encompassing_delete.<locals>.<listcomp>  s    
B
B
Ba3q66
B
B
Br+   c                 "    g | ]}|d z  |d z  gS r  r*   r  s     r,   r|   z,test_encompassing_delete.<locals>.<listcomp>  s%    UUU1s7AG$UUUr+   rd   )r%   rh   r'   r_   r  r   r   r$   r&   r   r  r   )	rP   colr  id_startrd   rg   
len_to_addget_resultsrz   s	            r,   test_encompassing_deleter    s   V$$$

"
"#8
9
9C,VSX>>OH
9
958a<88
9
9
9CLLeHhl.K.KLLLJ!|HGG
G+++JJ3J 
!CHo/I	
 	
 	
 -VSX>>O J
B
B58j+@AA
B
B
BCUUeHh>S.T.TUUUJGG
G+++JJ3JGG
G+++ 
!CHo/I	
 	
 	

 ''))K{5!""j0000 ( ([''''''( (r+   r   )qhypothesis.statefulrV   hypothesis.strategiesr   r  r  rj   r   r   r   r   typingr   r   r	   r
   r   r   r   dataclassesr   chromadb.api.typesr   r   r   r   r   r   chromadb.configr   chromadb.errorsr  chromadb.apir   chromadb.api.models.Collectionr   !chromadb.test.property.strategiestestpropertyr   r   r   r   r   r   r   r   r    r!   collectionsr"   !chromadb.test.property.invariantsr   chromadb.test.conftestr#   r$   r%   numpyr,  r*  -chromadb.test.utils.wait_for_version_increaser&   r'   r.   r   r   r   r   r4   r9   sharedsampled_fromfloat_typesr;   SearchStrategyr-  r  r  integersr@   rB   r   rM   r   LogCaptureFixturer  markskipifr  r5  r[  r]  r  r  r  r  r  xfailr  r  parametrizer  r  r  r   r  r  r   r  r  r  r  r  r*   r+   r,   <module>r     s                        " " " " " " 3 3 3 3 3 3 3 3 3 3 A A A A A A A A A A A A A A A A A A ! ! ! ! ! !                # " " " " "             " " " " " " 5 5 5 5 5 5 6 6 6 6 6 6 6 6 6 6 6 6                        $ # # # # # 6 6 6 6 6 6 6 6 6 6 6 6         
             !,II 6 6CH 6 6 6 ! s t    
! ! ! ! BIoboj4557CCC "	"*bj"*
,- D D D /8biBK!t,,,+/ / / R&s+   
 , , , , , , , , 	0*0$GGGVTTT}/ }/ }/ }/ }/ 5 }/ }/ }/@~/ ~/ ~/ ~/ ~/5 ~/ ~/ ~/B	&": 	I 	RV 	 	 	 	 :   
f.F 
PY 
^b 
 
 
	 
39 3 3 3 3 3lY9 Y Y Y Y Yx4V5 4y 4T 4 4 4 4n$	 $d $ $ $ $D9     &? ?t ? ? ? ?C9 C C C C C *) * * * * * UVVY 4    WVI $     		5%.!	K/05%.11K+?@@7U^*E2	
 	
 	 3 4     r
BHbhU'STTUUU#Y	   VU0 r
BHbhU'STTUUU#Y	   VU* #t55666
Cy
	
 
 
 76
   	
) 	
 	
 	
 	
 	
'	 'd ' ' ' 'D&(Y &(4 &( &( &( &( &( &(r+   