
    LKi^                        S SK r S SKrS SKrS SKJr  S SKJr  S SKJ	r	  S SK
Jr  S SKJr  S SKJr  S SKJr  S S	KJr  S S
KJr  S SKrS SKJr  S SKJrJr  S SKJrJrJrJ r J!r!J"r"J#r#J$r$J%r%  S SK&J'r'  S SK(r)S SK*Jr+  S SK,J-r-J.r.  S SK/J0r0  S SK1J2r2  S SK3J4r4  S SK5J6r6J7r7  S SK8J9r9  S SK:J;r;  S SK<J=r=J>r>J?r?J@r@  \"" S5      rAS\ \A\!\A   4   S\!\A   4S jrBS\S\4S jrCS\\-R                     S\\-R                     SS4S jrES \0S\SS4S! jrFS \0S"\S#\ \'S$   \'S%   \'S&   4   SS4S' jrGS \0S\SS4S( jrHS \0S\SS4S) jrIS \0S\SS4S* jrJS \0S\SS4S+ jrKS \0SS4S, jrL\9R                  4S-\-R                  S.\-R                  S/\\+R                  \+R                  /\P4   S\\!\!\Q      \!\!\P      4   4S0 jjrRS1\QSS4S2 jrSS \04S3 jrT      SHS \0S\S4\QS5\PS6\\-R                     S7\\!\Q      S8\\-R                     S9\VSS4S: jjrWS;\-R                  S4\QSS4S< jrYS=\S\Q4S> jrZS?\S@\!\0   SA\VSS4SB jr[S?\S@\!\0   S\%\\Q4   4SC jr\S?\S@\!\0   SA\VSS4SD jr]\SE\.SF\^4SG j5       r_g)I    N)UUID)contextmanager)
SegmentAPI)SysDB)create_topic_name)System)get_sql)SqliteDB)sleep)SegmentType)NormalizedRecordSet	RecordSet)	CallableOptionalTupleUnionListTypeVarcastAnyDict)Literal)types	ClientAPI)
Collection)note)InvalidArgument)Table	functions)distance_functions)Search)KnnSelectLimitKeyTvaluereturnc                 R    U c  [        S5      e[        U [        5      (       a  U $ U /$ )z*Wrap a value in a list if it is not a listzvalue cannot be None)r   
isinstancer   )r'   s    c/var/www/html/dynamic-report/venv/lib/python3.13/site-packages/chromadb/test/property/invariants.pywrapr,   $   s,    }455	E4	 	 w    
record_setc                    U S   c  SnGO[        U S   [        5      (       a  U S   c   e[        U S   5      S:  a  [        S U S    5       5      (       a  [	        [
        R                  U S   5      nO[        S U S    5       5      (       a  [	        [
        R                  U S   5      nOm[        S U S    5       5      (       a  [	        [
        R                  U S   /5      nO4[        S5      e[	        [
        R                  U S   5      nO[        S5      e[        [        U S	   5      U S
   b  [        U S
   5      OSU S   b  [        U S   5      US9$ SUS9$ )z9Ensure that an embedding set has lists for all its values
embeddingsNr   c              3   B   #    U  H  n[        U[        5      v   M     g 7fN)r*   list.0	embeddings     r+   	<genexpr>wrap_all.<locals>.<genexpr>7   s      =U	
9d++=Us   c              3   V   #    U  H  n[        U[        R                  5      v   M!     g 7fr2   )r*   npndarrayr4   s     r+   r7   r8   ;   s%      !9I 9bjj11!9s   ')c              3      #    U  H9  n[        U[        [        [        R                  [        R
                  45      v   M;     g 7fr2   )r*   intfloatr:   integerfloatingr5   es     r+   r7   r8   A   s2      5 q3rzz2;;"GHH5s   AAz-an embedding must be a list of floats or intszVembeddings must be a list of lists, a list of numpy arrays, a list of numbers, or Noneids	documents	metadatas)rC   rD   rE   r0   )
r*   r3   lenallr   r   
Embeddingsr   r   r,   )r.   embedding_lists     r+   wrap_allrJ   .   s    ,'	J|,d	3	3,'333z,'(1, =G=U   "&e&6&6
<8P!Q !+L!9   "&e&6&6
<8P!Q '5   &*%*:*:Z=U<V%WN)G  "%"2"2J|4LMNd
 	
 E"#k". z+./k". z+./ "	 	 !	 	r-   expectedgotc           	         U c  Ub  U b  Uc   eU b  Ub  [        U 5      [        U5      :X  d   eU R                  5        H|  u  p#X!;   d   e[        X   [        5      (       aM  [        X   [        5      (       a6  [	        [        [        X   5      [        [        X   5      -
  5      S:  d   eMp  X   X   :X  a  M|   e   g g g )Nư>)rF   itemsr*   r>   absr   )rK   rL   keyvals       r+   check_metadatarS   ]   s       8}C((( (HC::(-//Jsx4O4O4x}5UCH8MMNQUUUU}000 ) !0r-   
collectionc           	          U R                  5       n[        U5      nU[        US   5      :w  a  [        SUS[        US   5      5        U[        US   5      :X  d   eg)z?The given collection count is equal to the number of embeddingsrC   zcount mismatch:z=!N)countrJ   rF   print)rT   r.   rV   normalized_record_sets       r+   rV   rV   m   sa    E$Z0)%011c2G2N.OPC-e45555r-   rX   
field_namerD   rE   r0   c                   ^^ U R                  US   U/S9m[        US   5       VVs0 s H  u  p4XC_M	     snnmTU   n[        US   5      S:X  aB  US:X  a3  [        [        R
                  [           U5      R                  S:X  d   e gU/ :X  d   egUc   e[        [        U5      UU4S jS9nU VVs/ s H  u  pxUPM	     n	nnX   n
U
c  S/[        US   5      -  n
US:X  aG  [        R                  " [        R                  " U	5      [        R                  " U
5      5      (       d   eg[        U	5      [        U
5      :X  d   e[        X5       H_  u  p[        U
[        5      (       a>  [        [        [         R"                  U5      [        [         R"                  U
5      5        MX  X:X  a  M_   e   gs  snnf s  snnf )ze
The actual embedding field is equal to the expected field
field_name: one of [documents, metadatas]
rC   rC   includer   r0   Nc                     > TTS   U S         $ )NrC   r    )index_and_field_valueembedding_id_to_indexresults    r+   <lambda> _field_matches.<locals>.<lambda>   s    *?5M/23+
r-   rQ   )get	enumeraterF   r   nptNDArrayr   sizesortedr:   allclosearrayzipr*   dictrS   r   Metadata)rT   rX   rY   iidactual_fieldsorted_field_field_valuefield_valuesexpected_fieldr`   ra   s              @@r+   _field_matchesrx   v   s    ^^ 5e <zl^SF 1::OPU:V0WX0WuqRU0WX*%L
 '(A-%C(,7<<AAA 	  2%%% ###,
L 7CClNAKlLC*6N #&;E&B"CC\!{{288L1288N3KLLLL< C$7777+.|+L'K.$//58
 #444 ,M? Y& Ds   F55F;c                    ^ [        U5      nU R                  US   / S9S   n[        US   5       VVs0 s H  u  pEXT_M	     snnm[        UU4S jS9nX2S   :X  d   egs  snnf )z5The actual embedding ids is equal to the expected idsrC   r[   c                    > TU    $ r2   r^   )rq   r`   s    r+   rb   ids_match.<locals>.<lambda>   s    3H3Lr-   rd   N)rJ   re   rf   rj   )rT   r.   rX   
actual_idsrp   rq   r`   s         @r+   	ids_matchr}      sx    $Z0$9%$@"MeTJ 1::OPU:V0WX0WuqRU0WX
(LMJu5555 Ys   Ac                 2    [        U5      n[        XS5        g)z?The actual embedding metadata is equal to the expected metadatarE   NrJ   rx   rT   r.   rX   s      r+   metadatas_matchr          $Z0:kBr-   c                 2    [        U5      n[        XS5        g)AThe actual embedding documents is equal to the expected documentsrD   Nr   r   s      r+   documents_matchr      r   r-   c                 2    [        U5      n[        XS5        g)r   r0   Nr   r   s      r+   embeddings_matchr      s    $Z0:lCr-   c                 p    U R                  5       S   n[        U5      [        [        U5      5      :X  d   eg )NrC   )re   rF   set)rT   rC   s     r+   no_duplicatesr      s-    
..
5
!Cs8s3s8}$$$r-   querytargetsdistance_fnc                 4  ^^ [         R                  " U [         R                  S9n[         R                  " U[         R                  S9m[         R                  " UU4S jSU5      n[         R                  " U5      R                  5       UR                  5       4$ )zGReturn the ordered indices and distances from each query to each target)dtypec                 6   > [         R                  " TSTU 5      $ )N   )r:   apply_along_axis)r   r   
np_targetss    r+   rb   "_exact_distances.<locals>.<lambda>   s    b))+q*eLr-   r   )r:   rl   float32r   argsorttolist)r   r   r   np_query	distancesr   s     `  @r+   _exact_distancesr      sp     xxRZZ0H'4J ##L	I ::i '')9+;+;+===r-   threadpool_sizec                 t   [         R                  " 5       nUR                  5       nSnSn[        U Vs/ s H!  nSUR                  ;   d  M  UR                  PM#     sn5      S-
  U :  a  XC:  a{  [
        R                  " 5         UR                  5       nUS-  n[        S5        [        U Vs/ s H!  nSUR                  ;   d  M  UR                  PM#     sn5      S-
  U :  a  XC:  a  M{  [        U Vs/ s H!  nSUR                  ;   d  M  UR                  PM#     sn5      S-
  U ::  d   egs  snf s  snf s  snf )zg
Checks that the open file descriptors are not exceeding the threadpool size
works only for SegmentAPI
   r   sqlite3r   N)psutilProcess
open_filesrF   pathgccollectr   )r   current_processr   max_retriesretry_countps         r+    fd_not_exceeding_threadpool_sizer      s   
 nn&O ++-JKK 	Z?Z9+>VQVVZ?@1DV%


$//1
qa 	Z?Z9+>VQVVZ?@1DV% 	Z?Z9+>VQVVZ?@1DWW @? @s#   D+D+,D0D02D5
D5c                    S nSU R                   ;   a  U R                   S   nU R                  R                  c  U$ SU R                  R                  ;   a  U R                  R                  R                  S5      b_  SU R                  R                  R                  S5      ;   a6  U R                  R                  R                  S5      R                  S5      nU$ SU R                  R                  ;   a  U R                  R                  R                  S5      b`  SU R                  R                  R                  S5      ;   a7  Uc4  U R                  R                  R                  S5      R                  S5      nU$ )Nz
hnsw:spacespannspacehnsw)metadata_modelconfiguration_jsonre   )rT   r   s     r+   	get_spacer      s?    Ez***##L1++3:$$7770044W=Iz((;;??HH!!4488AEEgN L 	*##6660044V<Hz((;;??GG=%%88<<VDHHQELr-   	n_results
min_recallembedding_functionquery_indicesquery_embeddings
use_searchc           
      D	   [        U5      n[        US   5      S:X  a  gUS   n	U	SL=(       a    [        U	5      S:  n
U
(       d2  Uc   eUS   c   e[        US   [        5      (       d   eU" US   5      n	[	        U 5      nUc  [
        R                  nODUS:X  a  [
        R                  nO-US:X  a  [
        R                  nOUS:X  a  [
        R                  nS	nU R                  c   eU	c   e[        S
 U	 5       5      (       d   e[        U	S   5      nU[        R                  " S[        [        R                  " U5      5      5      -  nUc  Uc  U	OU Vs/ s H  oU   PM	     snnUS   nUb  Ub  U Vs/ s H  nUU   PM
     nn[        XiWS9u  nnU(       a  / nU Hl  n[        U[         R"                  5      (       a  UR%                  5       nOUn['        [)        US9[+        US9S9R-                  5       nUR/                  U5        Mn     U R1                  US9n[3        [4        R6                  US   US   US   US   US   S.5      nO&U R9                  U
(       a  UOSU
(       d  UOSU/ SQS9n[;        UU5        US   c   eUS   c   eUS   c   eUS   c   eUS   c   e[=        US   5       VVs0 s H	  u  nnUU_M     nnnSn[=        [?        UU5      5       GH"  u  nu  nn[         R@                  " US   5      USU    nU[        [C        U5      [C        US   U   5      -
  5      -  n[=        US   U   5       H  u  nnUU;  n UU   n![         RD                  " UU!   US   U   U   US9n"U (       a  U"(       a  US-  nOMH  U"(       d   e[         RD                  " U	U!   US   U   U   5      (       d   eUS   b  US   U!   US   U   U   :X  d   eUS   c  M  [G        US   U!   US   U   U   5        M     GM%     [        US   5      n#U#U-
  U#-  n$ [I        SU$ SU SU# SU 35        U$U:  d   eUS    H5  n%[         RD                  " [         RL                  " U%5      U%5      (       a  M5   e   gs  snf s  snf s  snnf ! [J         a     Ncf = f)zBValidate that the API performs nearest_neighbor searches correctlyrC   r   Nr0   rD   cosineipl2rN   c              3   b   #    U  H%  n[        U[        [        R                  45      v   M'     g 7fr2   )r*   r3   r:   r;   rA   s     r+   r7   ann_accuracy.<locals>.<genexpr>>  s#     E*Qz!dBJJ/00*s   -/
   )r   )r   )limit)rankr   )searchesscoresrE   )rC   r   r0   rD   rE   )r0   rD   rE   r   )r   query_textsr   r\   r   )atolr   z
# recall: z
, missing z out of z, accuracy threshold )'rJ   rF   r*   r3   r   r    r   r   r   r   rG   mathpowr=   log10r   r:   r;   r   r!   r"   r$   
select_allappendsearchr   r   QueryResultr    _query_results_are_correct_shaperf   rm   rl   r   rk   rS   r   r   sort)&rT   r.   r   r   r   r   r   r   rX   r0   have_embeddingsr   distance_functionaccuracy_thresholddimrp   query_documentsindicesr   search_requestsquery_embeddingquery_embedding_listr   search_resultsquery_resultsrq   id_to_indexmissing	indices_idistances_iexpected_idsjunexpected_idindexcorrect_distanceri   recalldistance_results&                                         r+   ann_accuracyr     s    %Z0
 '(A--B<-PJ ,DZ11DO!---$[1===/<dCCCC'(=k(JK
j!E}.11	(	.55	$.11	$.11***!!!
 E*EEEEE
jm
C+dhhr3tzz#;O.PP  $ )67AQ-7 	
 ,K8O _%@7DE}!?1-}E)2CGY /O/2::66'6'='='?$'6$34), jl  ""6*  0 $**$ + 

 %e,+ -\:+K8+K8
 #((1@-d/>DI	 ) 
 %]I> +++%111&222%111%111 '00Ee0L&MN&MUQ2q5&MKNG'0Wi1H'I##I{xx 5e <=i
>ST3s<(3}U/CA/F+GGHH }U3A67EArl2MOE!{{E"k*1-a0' 
  $qLG''';;z%0-2Ma2PQR2STTTT$[1=)+6u=$[1!4Q788 %[1=)+6u=!+.q1!47 8 (JL $U+,DWn$F
7)8D6AVWiVjk	
 Z )5{{277?3_EEEE 6e 8 Fr O^  s$   R9R<R$R 
RRr   c                 b   ^ S H(  nX   c   e[        U4S jX    5       5      (       a  M(   e   g )N)r   r0   rD   rE   c              3   @   >#    U  H  n[        U5      T:H  v   M     g 7fr2   )rF   )r5   ra   r   s     r+   r7   3_query_results_are_correct_shape.<locals>.<genexpr>  s      
3MCK9$3Ms   )rG   )r   r   result_types    ` r+   r   r     sG     M)555 
3@3M
 
 
 	
 
 Mr-   sqlitec                    [        S5      nU R                  5       R                  U5      nU R                  5        n[	        UR                  [        R                  " UR                  5      5      U R                  5       5      u  pEUR                  XE5      n[        [        UR                  5       S   5      sS S S 5        $ ! , (       d  f       g = f)Nembeddings_queuer   )r   querybuilderfrom_txr	   selectr   Countseq_idparameter_formatexecuter   r=   fetchone)r   tqcursqlparamsra   s          r+   _total_embedding_queue_log_sizer     s     !A##A&A	HHY__QXX./1H1H1J
 S)C*1-. 
s   A;C  
Csystemcollectionshas_collection_mutatedc                    U R                  [        5      nU R                  R                  (       d  U R                  R                  S:X  a  g U(       aj  [        U5      S:  d   e[        S U 5       5      n[        S U 5       5      nU R                  R                  S:X  a  UOXE-   n[        U5      S-
  U::  d   eg [        U5      S:X  d   eg )N!chromadb.api.rust.RustBindingsAPIr   c              3   x   #    U  H0  nUR                   b  UR                   R                  SS5      OSv   M2     g 7f)Nhnsw:sync_threshold  r   re   r5   rT   s     r+   r7   %log_size_below_max.<locals>.<genexpr>  sE      !
 *
 "". ##$94@ *   8:c              3   x   #    U  H0  nUR                   b  UR                   R                  SS5      OSv   M2     g 7f)Nhnsw:batch_sized   r  r  s     r+   r7   r    sE      
 *
 "". ##$5s; *r  r   )instancer
   settingsis_persistentchroma_api_implr   sum)r   r   r   r   sync_threshold_sumbatch_size_sumr   s          r+   log_size_below_maxr    s     __X&F OO))OO++/RR.v6!;;; ! !
 *	!
 
  
 *	
 
 ..2UU #4 	 /v6:eCCC.v6!;;;r-   c           	         U R                  [        5      n[        S5      nUR                  5       R	                  U5      nU R
                  R                  S5      nU R
                  R                  S5      nU Vs0 s H  n[        XVUR                  5      U_M     nnUR                  5        n	[        UR                  UR                  [        R                  " UR                  5      5      R!                  S5      UR#                  5       5      u  pU	R%                  X5      n0 nUR'                  5        H  nUS   XUS      R                  '   M     UsS S S 5        $ s  snf ! , (       d  f       g = f)Nr   	tenant_idtopic_namespacetopicr   r   )r
  r
   r   r   r   r  requirer   rq   r   r	   r   r  r   r   r   groupbyr   r   fetchall)r   r   r   r   r   _tenant_topic_namespacerT   topic_mappingsr   r   r   ra   outress                  r+   ._total_embedding_queue_log_size_per_collectionr    s9    __X&F !A##A&Aoo%%k2G../@A &%J 	'Z]]CZO%   
HHQWWiooahh78@@I##%
 S)??$C-0VCs1v&))* % 
	 
s   :!E-B&E""
E0c                   ^^ U R                   R                  S:X  a  g U R                  [        5      nU(       a  [	        U5      S:  d   eU Vs0 s H9  nUR
                  UR                  b  UR                  R                  SS5      OS_M;     nnU Vs0 s H0  nUR
                  UR                  5       XTR
                     -  S-   _M2     snm[        X5      m[        TR                  5       5      [        TR                  5       5      :X  d   e[        UU4S jU 5       5      (       d   eg [	        U5      S:X  d   eg s  snf s  snf )Nr   r   r  r	  c              3   b   >#    U  H$  nTUR                      TUR                      :H  v   M&     g 7fr2   )rq   )r5   rT   actual_sizesexpected_sizess     r+   r7   :log_size_for_collections_match_expected.<locals>.<genexpr>3  s/      
)
 '>*--+HH)s   ,/r   )r  r  r
  r
   r   rq   r   re   rV   r  r   keysrG   )r   r   r   r   rT   r  r!  r"  s         @@r+   'log_size_for_collections_match_expectedr%    sX    &&*MM__X&F.v6!;;; *	
 *
 MM"". &..223DcJ *	 	 
 *
)
 MM:++-}}0MMPQQQ)

 F
 <$$&'3~/B/B/D+EEEE 
)
 
 
 	
 
 /v6!;;;+

s   A E 7Eclientcollection_namec           
   #     #    U R                  5        Vs/ s H  o"R                  PM     nnX;   d   eU R                  U5      n/ n[        U R                  [
        5      (       Gac  U R                  R                  nUR                  UR                  S9n0 nSnU H  n	SXyS   '   U	S   [        R                  R                  :X  d  M-  UR                  b  SUR                  ;   a  UR                  S   OSn
UR                  5       U
:  d  Mq  Sn[        R                  R!                  [        R                  R#                  U R%                  5       R&                  [)        U	S   5      5      5      (       a  M   e   U(       a$  U[        R                  R                     (       d   eU[        R*                  R                     (       d   eS v   U R                  5        Vs/ s H  o"R                  PM     nnX;  d   e[-        U5      S:  a  U R                  R                  nUR                  UR                  S9n[-        U5      S:X  d   eU H  n	U	S   [        R                  R                  :X  d  M&  [        R                  R!                  [        R                  R#                  U R%                  5       R&                  [)        U	S   5      5      5      (       d  M   e   g g s  snf s  snf 7f)	N)rT   FTtyper  r  rq   r   )list_collectionsnameget_collectionr*   _serverr   _sysdbget_segmentsrq   r   HNSW_LOCAL_PERSISTEDr'   r   rV   osr   existsjoinget_settingspersist_directorystrSQLITErF   )r&  r'  ccollection_namesrT   segmentssysdbsegment_typesshould_have_hnswsegmentsync_thresholdsegments_afters               r+   collection_deletedrA  <  s     )/(?(?(AB(A1(AB...&&7JH&..*--~~,,%%%? G-1M&/*v+"B"B"H"HH "**6-1D1DD ''(=> 	  $$&7'+$77>>"//1CCSQUEW     $  !A!A!G!GHHH[//55666	 )/(?(?(AB(A1(AB222
8}q~~,,++z}}+E>"a'''Gv+"B"B"H"HH77>>GGLL++-??WT]AS     	 I CD Cs?   K,K"BK,A K,A&K,6A)K,K'4A<K,4A$K,K,)r   gffffff?NNNF)`r   r   os.pathr1  uuidr   
contextlibr   chromadb.api.segmentr   chromadb.db.systemr   chromadb.ingest.impl.utilsr   chromadb.configr   chromadb.db.baser	   chromadb.db.impl.sqliter
   timer   r   chromadb.segmentr   !chromadb.test.property.strategiesr   r   typingr   r   r   r   r   r   r   r   r   typing_extensionsr   numpyr:   numpy.typingrg   chromadb.apir   r   chromadb.api.models.Collectionr   
hypothesisr   hypothesis.errorsr   pypikar   r   chromadb.utilsr    "chromadb.execution.expression.planr!   &chromadb.execution.expression.operatorr"   r#   r$   r%   r&   r,   rJ   ro   rS   rV   rx   r}   r   r   r   r   r   rH   	ArrayLiker>   r=   r   r   r   EmbeddingFunctionboolr   r   r   r   r  r  r%  r6  rA  r^   r-   r+   <module>r]     s   	    % + $ 8 " $ ,   ( L S S S %   ) 5  - # - 5 J JCLaaj! d1g , ,': ,^1u~~&1-5enn-E1	1 6j 6i 6D 64545.45 gk2GL4II45 
45n6* 6) 6 6C
 C	 Cd CC
 C	 Cd CD D Dt D%j %T % 	>>> 	&-> 4S	?De--.>*c d .* 8 <@)-37fFfFfF fF 	fF
 !!8!89fF DI&fF u//0fF fF 
fFR
$$
14
	
	/H 	/ 	/(<(<!%j!1(<KO(<	(<Vj! 
$)_2"<"<!%j!1"<KO"<	"<J 0y 03 0 0r-   