
    q9i             
       $   d dl Z d dlZd dlZd dlZd dlZd dlmZmZ d dlmZ d dl	Z	d dl
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mZmZ d dlmZ d dlmZmZmZ d dlmZ ej>                  d	        Z ej>                  d
        Z!ej>                  d        Z"dde#fdZ$dde%de#fdZ&ejN                  jQ                  de!g      d        Z)ejN                  jQ                  de!g      d        Z*ejN                  jQ                  de!g      d        Z+ejN                  jQ                  de!g      d        Z,d Z-d Z.d Z/d Z0d Z1g dg dgddgdZ2d Z3d  Z4d! Z5g dg dgddgdZ6d" Z7d# Z8d$ Z9d% Z:d& Z;d' Z<d( Z=d) Z>d* Z?d+ Z@d, ZAd- ZBd. ZCd/ ZDd0 ZEd1 ZFd2 ZGd3 ZHd4 ZId5 ZJd6 ZKd7 ZLd8 ZMd9 ZNd: ZOg dg dgd;d<gd=d>d?d@dAdBigdCZPdD ZQdE ZRdF ZSdG ZTdH ZUdI ZVg dg dgd;d<gdJdKdLiidJg dMigdCZWdN ZXdO ZYdP ZZdQ Z[g dg dgd;d<gd=d>d?d@dBdRdSdTgdCZ\dU Z]dV Z^dW Z_dX Z`dY ZadZ Zbd[ Zcg d\g d]gd;d<gdZdd^g d\g d]giZeg dg dgd_d`Zfda Zgdb Zhdc Zig dg dgdddegd;d<gd=d>d?d@dBdRdSdTgdfZjdg Zkdh Zldi Zmg dg dg djg dkgg dld=d>d?dmdndBdRdSdmdodpdqdrdmdodsdtdudmdogg dvdwZndx Zody Zpg dzg dgd;d<gd=d>d?d@dAdBigd{d|gdwZqd} Zrd~ Zsd Ztd Zud Zvd Zwd Zxd Zyd Zzd Z{d Z|d Z}d Z~d Zd Zg dzg dg dgg dd=d>d?d@dAdBiddrigg ddwZg dg dgd;dgd=dd?d@dAdsigddgdwZd Zd Zd Zd Zd ZejN                  j                  dpej                  j                  d            d        ZejN                  j                  dpej                  j                  d            d        Zd Zd Zd Zd Zd Zd Zd ZddZddZ G d d      Z G d d      Z G d d      Z G d d      Z G d d      Z G d d      Zy)    N)datetime	timedelta)Any)FastAPI)DocumentEmbeddingFunctionQueryResultTYPE_KEYSPARSE_VECTOR_TYPE_VALUE)Settings)ChromaErrorNotFoundErrorInvalidArgumentError)DefaultEmbeddingFunctionc                  *    t        j                         S N)tempfilemkdtemp     _/var/www/html/leadgen/airagagent/rag_env/lib/python3.12/site-packages/chromadb/test/test_api.pypersist_dirr   !   s    r   c              #      K   t        j                  t        ddddddd|             }| |j                          t        j
                  j                  |       rt        j                  | d       y y wNzchromadb.api.segment.SegmentAPIz chromadb.db.impl.sqlite.SqliteDBz7chromadb.segment.impl.manager.local.LocalSegmentManagerT)chroma_api_implchroma_sysdb_implchroma_producer_implchroma_consumer_implchroma_segment_manager_implallow_resetis_persistentpersist_directory)ignore_errors	chromadbClientr   clear_system_cacheospathexistsshutilrmtreer   clients     r   local_persist_apir/   &   j     __=@!C!C(a)		
F L
	ww~~k"k6 #   A3A5c              #      K   t        j                  t        ddddddd|             }| |j                          t        j
                  j                  |       rt        j                  | d       y y wr   r$   r-   s     r   local_persist_api_cache_bustr3   ;   r0   r1   returnc                 $    t        | |z
        |k  S r   )absab	tolerances      r   approx_equalr;   O   s    q1u:	!!r   r:   c                     t        |       t        |      k7  ryt        t        | |      D  cg c]  \  } }t        | ||       c}}       S c c}} w )NF)lenallzipr;   r7   s      r   vector_approx_equalr@   S   sA    
1vQ#a)D$!QQ9-DEEDs   A
api_fixturec                    |j                  d      }|j                          |j                  d      }|j                  dd       |j                  d      }|j	                  d      }g d}|j                  ddg d	      }|j                         D ]H  }||v s|d
k(  rt        ||         dk(  rJ |dk(  rt        ||         t        |      k(  r@J ||   HJ  y )Nr/   testid1helloids	documentsr3   
embeddingsrH   	metadatas	distances   query_texts	n_resultsincluderG   included)	getfixturevalueresetcreate_collectionaddget_collectionquerykeysr=   set)rA   requestr.   
collectionapi2includesnnkeys           r   test_persist_index_loadingra   Y   s    $$%89F
LLN))&1JNNuN0""#ABD$$V,JDH			E 
 
B
 wwy #8Or#w<1$$$Jr#w<3x=000c7?"?#r   c                      G d dt         t                 }|j                  d      }|j                          |j	                  d |             }|j                  dd       |j                  d	      }|j                  d |             }g d
}|j                  dd|      }|j                         D ]H  }||v s|dk(  rt        ||         dk(  rJ |dk(  rt        ||         t        |      k(  r@J ||   HJ  y )Nc                   n     e Zd Zd Zdededdf fdZdefdZdeeef   ddfd	Z	deeef   fd
Z
 xZS )=test_persist_index_loading_embedding_function.<locals>.TestEFc                 z    t        t        |            D cg c]  }t        j                  g d       c}S c c}w NrM         ranger=   nparrayselfinput_s      r   __call__zFtest_persist_index_loading_embedding_function.<locals>.TestEF.__call__u   (    16s5z1BCABHHY'CCC   8argskwargsr4   Nc                 $    t        |   |i | y r   super__init__ro   ru   rv   	__class__s      r   rz   zFtest_persist_index_loading_embedding_function.<locals>.TestEF.__init__x       Gd-f-r   c                      yNrC   r   ro   s    r   namezBtest_persist_index_loading_embedding_function.<locals>.TestEF.name{       r   configc                      y r   r   ro   r   s     r   build_from_configzOtest_persist_index_loading_embedding_function.<locals>.TestEF.build_from_config~       r   c                     i S r   r   r   s    r   
get_configzHtest_persist_index_loading_embedding_function.<locals>.TestEF.get_config       Ir   __name__
__module____qualname__rr   r   rz   strr   dictr   r   __classcell__r|   s   @r   TestEFrd   t   ^    	D	.# 	. 	. 	.	# 		DcN 	t 		S#X 	r   r   r/   rC   embedding_functionrD   rE   rF   r3   rI   rM   rN   rG   rR   )r   r   rS   rT   rU   rV   rW   rX   rY   r=   rZ   )	rA   r[   r   r.   r\   client2r^   r_   r`   s	            r   -test_persist_index_loading_embedding_functionr   r   s   "8,   $$%89F
LLN))&VX)NJNNuN0%%&DEG''68'LJDH			 
 
B
 wwy #8Or#w<1$$$Jr#w<3x=000c7?"?#r   c                     G d dt         t                 }|j                  d      }|j                          |j	                  d |             }|j                  dd       |j                  d	      }|j	                  d |             }g d
}|j                  dd|      }|j                         D ]H  }||v s|dk(  rt        ||         dk(  rJ |dk(  rt        ||         t        |      k(  r@J ||   HJ  |d   dggk(  sJ |d   d   d   j                         g dk(  sJ |d   dggk(  sJ |d   dggk(  sJ y )Nc                   n     e Zd Zd Zdededdf fdZdefdZdeeef   ddfd	Z	deeef   fd
Z
 xZS )Ctest_persist_index_get_or_create_embedding_function.<locals>.TestEFc                 z    t        t        |            D cg c]  }t        j                  g d       c}S c c}w rf   rj   rn   s      r   rr   zLtest_persist_index_get_or_create_embedding_function.<locals>.TestEF.__call__   rs   rt   ru   rv   r4   Nc                 $    t        |   |i | y r   rx   r{   s      r   rz   zLtest_persist_index_get_or_create_embedding_function.<locals>.TestEF.__init__   r}   r   c                      yr   r   r   s    r   r   zHtest_persist_index_get_or_create_embedding_function.<locals>.TestEF.name   r   r   r   c                      y r   r   r   s     r   r   zUtest_persist_index_get_or_create_embedding_function.<locals>.TestEF.build_from_config   r   r   c                     i S r   r   r   s    r   r   zNtest_persist_index_get_or_create_embedding_function.<locals>.TestEF.get_config   r   r   r   r   s   @r   r   r      r   r   r   r/   rC   r   rD   rE   rF   r3   rI   rM   rN   rG   rR   rJ   r   rg   rH   rL   )r   r   rS   rT   get_or_create_collectionrV   rX   rY   r=   rZ   tolist)	rA   r[   r   apir\   r]   r^   r_   r`   s	            r   3test_persist_index_get_or_create_embedding_functionr      sv   "8,   
!
!"5
6CIIK--f-RJNNuN0""#ABD..v&(.SJDH			 
 
B wwy #8Or#w<1$$$Jr#w<3x=000c7?"?# e9%	!!!lAq!((*i777k?yk)))k?se###r   c                    |j                  | j                        }|j                          |j                  d      } |j                  di t
         |j                         dk(  sJ |j                  | j                        }|j                  d      }|j                         dk(  sJ |j                  d       |j                  | j                        }|j                         g k(  sJ y N	testspacerh   r   )
rS   r   rT   rU   rV   batch_recordscountrW   delete_collectionlist_collections)rA   r[   r.   r\   s       r   test_persistr      s    $$[%9%9:F
LLN))+6JJNN#]#"""$$[%9%9:F&&{3J"""
[)$$[%9%9:F""$***r   c                     | j                         }t        |t              sJ |dz  }t        j                  |      }|t        j
                         t        d      z
  kD  sJ y )Ni ʚ;
   )seconds)	heartbeat
isinstanceintr   fromtimestampnowr   )r.   heartbeat_nsheartbeat_sr   s       r   test_heartbeatr      sX    ##%LlC(((%'K&&{3Ix||~	"(=====r   c                 2    | j                         }|dkD  sJ y )Nr   )get_max_batch_size)r.   
batch_sizes     r   test_max_batch_sizer      s    **,J>>r   c                     t        | t              st        j                  d       | j	                          | j                         }|du sJ y )NNot a FastAPI instanceT)r   r   pytestskiprT   supports_base64_encoding)r.   r   s     r   test_supports_base64_encodingr      s?    fg&,-
LLN%>>@#t+++r   c                     t        | t              st        j                  d       | j	                          ddi| _        | j                         du sJ | j                         dk(  sJ y )Nr   max_batch_sized   F)r   r   r   r   rT   pre_flight_checksr   r   )r.   s    r   $test_supports_base64_encoding_legacyr      sc    fg&,-
LLN 	# F **,555$$&#---r   c                 `   t        | t              st        j                  d       t	        j
                  | j                   d      }|j                  dk(  sJ |j                         J d|j                         j                         v sJ d|j                         j                         v sJ y )Nr   z/pre-flight-checks   r   r   )
r   r   r   r   httpxget_api_urlstatus_codejsonrY   )r.   resps     r   test_pre_flight_checksr     s    fg&,-99((:;<Ds"""99;"""tyy{//1111%)9)9);;;;r   皙?ffffff@皙	@333333?Q@r   zhttps://example.com/1zhttps://example.com/2rJ   rG   c                     | j                          | j                  d      } |j                  di t         |j	                         dk(  sJ y r   )rT   rU   rV   r   r   r.   r\   s     r   test_addr     sC    
LLN))+6JJNN#]#"""r   c                     | j                          | j                  d      }| j                  d       t        j                  t
        d      5   |j                  di t         d d d        y # 1 sw Y   y xY wNrC   Collection .* does not existmatchr   )rT   rU   r   r   raisesr   rV   r   r   s     r   2test_collection_add_with_invalid_collection_throwsr   "  s^    
LLN))&1J
V$	},K	L (
''( ( (   A..A7c                 f   | j                          | j                  d      } |j                  di t         |j	                         dk(  sJ t        j                  t              5  | j                  d      }d d d        | j                  d      }|j	                         dk(  sJ y # 1 sw Y   0xY wr   )	rT   rU   rV   r   r   r   r   	Exceptionr   r   s     r   test_get_or_creater   +  s    
LLN))+6JJNN#]#"""	y	! ;--k:
; 00=J"""; ;s   &B''B0c                     | j                          | j                  d      } |j                  di t         |j	                         dk(  sJ y r   )rT   rU   rV   minimal_recordsr   r   s     r   test_add_minimalr   B  sC    
LLN))+6JJNN%_%"""r   c                 T   | j                          | j                  d      } |j                  di t         g d}|j	                  |      }|j                         D ]H  }||v s|dk(  rt        ||         dk(  rJ |dk(  rt        ||         t        |      k(  r@J ||   HJ  y )Nr   rJ   rH   rK   rQ   rG   rh   rR   r   )rT   rU   rV   r   r   rY   r=   rZ   )r.   r\   r^   recordsr`   s        r   test_get_from_dbr   L  s    
LLN))+6JJNN#]#7HnnXn.G||~ (8Ows|$)))Jws|$H5553<'''(r   c                     | j                          | j                  d      }| j                  d       t        j                  t
        d      5  |j                          d d d        y # 1 sw Y   y xY wNrC   r   r   )rT   rU   r   r   r   r   r   r   s     r   2test_collection_get_with_invalid_collection_throwsr   [  sW    
LLN))&1J
V$	},K	L      A((A1c                     | j                          | j                  d      } |j                  di t         |j	                         dk(  sJ | j                          t        | j                               dk(  sJ y )Nr   rh   r   r   )rT   rU   rV   r   r   r=   r   r   s     r   test_reset_dbr   d  sh    
LLN))+6JJNN#]#"""
LLNv&&()Q...r   c                 *   | j                          | j                  d      } |j                  d
i t         g d}|j	                  g dd|      }|j                         D ]H  }||v s|dk(  rt        ||         dk(  rJ |dk(  rt        ||         t        |      k(  r@J ||   HJ  |j	                  g dgd|      }|j                         D ]H  }||v s|dk(  rt        ||         dk(  rJ |dk(  rt        ||         t        |      k(  r@J ||   HJ  |j	                  g dg dgd|      }|j                         D ]H  }||v s|dk(  rt        ||         d	k(  rJ |dk(  rt        ||         t        |      k(  r@J ||   HJ  y )Nr   rI   r   rM   query_embeddingsrP   rQ   rG   rR   )皙?r         @rh   r   )rT   rU   rV   r   rX   rY   r=   rZ   r.   r\   r^   r_   r`   s        r   test_get_nearest_neighborsr   o  s   
LLN))+6JJNN#]#DH			( 
 
B
 wwy #8Or#w<1$$$Jr#w<3x=000c7?"?# 
		)* 
 
B
 wwy #8Or#w<1$$$Jr#w<3x=000c7?"?# 
		)?; 
 
B
 wwy #8Or#w<1$$$Jr#w<3x=000c7?"?#r   c                    | j                          | j                  d      } |j                  di t         |j	                         dk(  sJ t        j                  t              5  |j                          d d d        y # 1 sw Y   y xY wr   )	rT   rU   rV   r   r   r   r   r   deleter   s     r   test_deleter     sp    
LLN))+6JJNN#]#"""	y	!   s   &B  B	c                     | j                          | j                  d      } |j                  di t         |j	                         dk(  sJ |j                  t        d         J y )Nr   rh   rG   rG   r   )rT   rU   rV   r   r   r   r   s     r   test_delete_returns_noner     sb    
LLN))+6JJNN#]#"""u!56>>>r   c                     | j                          | j                  d      } |j                  di t         |j	                         dk(  sJ |j                  g dgd       y )Nr   rh   r   rM   r   rP   r   )rT   rU   rV   r   r   rX   r   s     r   test_delete_with_indexr    sY    
LLN))+6JJNN#]#"""&71Er   c                     | j                          | j                  d      }| j                  d       t        j                  t
        d      5  |j                  dg       d d d        y # 1 sw Y   y xY w)NrC   r   r   rD   r   )rT   rU   r   r   r   r   r   r   s     r   5test_collection_delete_with_invalid_collection_throwsr    s`    
LLN))&1J
V$	},K	L 'ug&' ' '   A++A4c                     | j                          | j                  d      }|j                         dk(  sJ  |j                  di t         |j                         dk(  sJ y )Nr   r   rh   r   )rT   rU   r   rV   r   r   s     r   
test_countr    sZ    
LLN))+6J"""JNN#]#"""r   c                     | j                          | j                  d      }| j                  d       t        j                  t
        d      5  |j                          d d d        y # 1 sw Y   y xY wr   )rT   rU   r   r   r   r   r   r   s     r   4test_collection_count_with_invalid_collection_throwsr
    sY    
LLN))&1J
V$	},K	L   r   c                     | j                          | j                  d      }|j                  d       |j                  dk(  sJ y Nr   
testspace2r   )rT   rU   modifyr   r   s     r   test_modifyr    s@    
LLN))+6J<( ??l***r   c                     | j                          | j                  d      }| j                  d       t        j                  t
        d      5  |j                  d       d d d        y # 1 sw Y   y xY w)NrC   r   r   test2r  )rT   rU   r   r   r   r   r  r   s     r   5test_collection_modify_with_invalid_collection_throwsr    s^    
LLN))&1J
V$	},K	L (w'( ( (s   A**A3c                     | j                          | j                  d       | j                  d      }t        j                  t              5  |j                  d       d d d        y # 1 sw Y   y xY wr  rT   rU   r   r   r   r  )r.   c2s     r   "test_modify_error_on_existing_namer    sX    
LLN
[)		!	!,	/B	y	! $
		{	#$ $ $s   A((A1c                     | j                          | j                  d      }t        j                  t        d      5  |j                  ddi       d d d        y # 1 sw Y   y xY w)Nr   znot supportedr   
hnsw:spacecosinemetadatar  )r.   caplogr\   s      r   test_modify_warn_on_DF_changer    sT    
LLN))+6J	y	8 =L(#;<= = =s   AA$c                 X   | j                          ddd}| j                  d|      }|j                  J |j                  d   dk(  sJ |j                  d   dk(  sJ | j                  d      }|j                  J |j                  d   dk(  sJ |j                  d   dk(  sJ |j	                  ddd	       |j                  d   dk(  sJ |j                  d
   dk(  sJ d|j                  vsJ | j                  d      }|j                  J |j                  d   dk(  sJ |j                  d
   dk(  sJ d|j                  vsJ | j                  d      }|j                  J |j                  d   dk(  sJ |j                  d
   dk(  sJ | j                  d      }|j                  J | j                         }|D ]f  }|j                  dk(  r7|j                  J |j                  d   dk(  sJ |j                  d
   dk(  rGJ |j                  dk(  sY|j                  fJ  y )NrM   rh   )r8   r9   r   r  r8   r9   ri   )r8   cr   r  )rT   rU   r  rW   r  r   r   r   )r.   
metadata_ar\   collectionss       r   test_metadata_crur#    s   
LLNq!J))+
)KJ***s#q(((s#q((( &&{3J***s#q(((s#q((( QQ/0s#q(((s#q(((j))))) &&{3J***s#q(((s#q(((j))))) 00=J***s#q(((s#q((( 00>J&&& ))+K! /
??k)&&222&&s+q000&&s+q000__,&&.../r   c                    | j                          | j                  d      } |j                  d	i t         |j	                         dk(  sJ g d}|j                  g dgd|      }|j                         D ]H  }||v s|dk(  rt        ||         dk(  rJ |dk(  rt        ||         t        |      k(  r@J ||   HJ  y )
Nr   rh   rI   r   rM   r   rG   rR   r   )	rT   rU   rV   r   r   rX   rY   r=   rZ   r   s        r   test_increment_index_onr%  '  s    
LLN))+6JJNN#]#"""DH			)* 
 
B
 wwy #8Or#w<1$$$Jr#w<3x=000c7?"?#r   c                    | j                          | j                  d       | j                  d      }|j                  dk(  sJ t	        j
                  t              5  | j                  d      }d d d        y # 1 sw Y   y xY w)Nr   r  )rT   rU   rW   r   r   r   r   r   s     r   test_add_a_collectionr'  =  sp    
LLN
[) &&{3J??k))) 
y	! 9**<8
9 9 9s   A88Bc                     | j                          t        j                  t              5 }| j	                  d       d d d        j
                  j                  J y # 1 sw Y   "xY w)Nr  )rT   r   r   r   rW   valuetrace_id)http_clienterrors     r   test_error_includes_trace_idr-  J  sW    	{	# 1u""<01 ;;+++1 1s   AA&c                     | j                          | j                  d       | j                  d       | j                         }t        |      dk(  sJ y )Nr   r  rh   rT   rU   r   r=   r.   r"  s     r   test_list_collectionsr1  S  sJ    
LLN
[)
\* ))+K{q   r   c                    | j                          | j                  d       | j                  d       | j                         }t        |      dk(  sJ | j                          | j                         }t        |      dk(  sJ y )Nr   r  rh   r   r/  r0  s     r   
test_resetr3  ]  st    
LLN
[)
\* ))+K{q   
LLN))+K{q   r   c                    | j                          | j                  d      } |j                  di t         |j	                         dk(  sJ |j                         }t        |       |j                         D ]J  }|dv s|dk(  rt        ||         dk(  rJ |dk(  rt        ||         t        g d      k(  rBJ ||   JJ  y )Nr   rh   r   rG   rR   rJ   rK   rH   r   )
rT   rU   rV   r   r   peekprintrY   r=   rZ   )r.   r\   r6  r`   s       r   	test_peekr8  k  s    
LLN))+6JJNN#]#""" ??D	$Kyy{ %::cUltCy>Q&&&JtCy>S)Q%RRRR9$$$%r   c                     | j                          | j                  d      }| j                  d       t        j                  t
        d      5  |j                          d d d        y # 1 sw Y   y xY wr   )rT   rU   r   r   r   r   r6  r   s     r   3test_collection_peek_with_invalid_collection_throwsr:  }  sW    
LLN))&1J
V$	},K	L   r   c                     | j                          | j                  d      }| j                  d       t        j                  t
        d      5  |j                  dg       d d d        y # 1 sw Y   y xY w)NrC   r   r   )rO   )rT   rU   r   r   r   r   rX   r   s     r   4test_collection_query_with_invalid_collection_throwsr<    s`    
LLN))&1J
V$	},K	L /fX./ / /r  c                     | j                          | j                  d      }| j                  d       t        j                  t
        d      5  |j                  dgdg       d d d        y # 1 sw Y   y xY w)NrC   r   r   rD   rF   )rT   rU   r   r   r   r   updater   s     r   5test_collection_update_with_invalid_collection_throwsr?    sd    
LLN))&1J
V$	},K	L ;ug&:; ; ;s   A--A6rD   id2rM   onejt?	int_valuestring_valuefloat_valuerD  rh   rJ   rG   rK   c                 f   | j                          | j                  d      } |j                  di t         |j	                  ddg      }|d   d   d   dk(  sJ |d   d   d	   d
k(  sJ |d   d   d   dk(  sJ t        |d   d   d   t              sJ t        |d   d   d	   t              sJ y )Ntest_intrD   r@  r   rK   r   rD  rM   rF  rB  rh   r   )rT   rU   rV   metadata_recordsr   r   r   floatr.   r\   itemss      r   test_metadata_add_get_int_floatrN    s    
LLN))*5JJNN&%&NNu~N.Ea -222a /5888a -222eK(+K8#>>>eK(+M:EBBBr   c                 p   | j                          | j                  d      } |j                  d
i t         |j	                  g dgd      }|d   J |d   d   d   d   dk(  sJ |d   d   d   d   d	k(  sJ t        |d   d   d   d   t              sJ t        |d   d   d   d   t              sJ y )NrI  r   rM   r  rK   r   rD  rF  rB  r   )rT   rU   rV   rJ  rX   r   r   rK  rL  s      r   !test_metadata_add_query_int_floatrP    s    
LLN))*5JJNN&%&#)))*a * E )))a #K0A555a #M2e;;;eK(+A.{;SAAAeK(+A.}=uEEEr   c                     | j                          | j                  d      } |j                  d	i t         |j	                  ddi      }|d   d   d   dk(  sJ |d   d   d   dk(  sJ y )
NrI  rE  rA  whererK   r   rD  rM   r   rT   rU   rV   rJ  r   rL  s      r   test_metadata_get_where_stringrU    s|    
LLN))*5JJNN&%&NN.%!8N9Ea -222a 0E999r   c                     | j                          | j                  d      } |j                  d	i t         |j	                  ddi      }|d   d   d   dk(  sJ |d   d   d   dk(  sJ y )
NrI  rD  rM   rR  rK   r   rE  rA  r   rT  rL  s      r   test_metadata_get_where_intrW    s|    
LLN))*5JJNN&%&NN+q!1N2Ea -222a 0E999r   c                     | j                          | j                  d      } |j                  di t         |j	                  ddi      }|d   d   d   dk(  sJ |d   d   d	   d
k(  sJ |d   d   d   dk(  sJ y )NrI  rF  rB  rR  rK   r   rD  rM   rE  rA  r   rT  rL  s      r   test_metadata_get_where_floatrY    s    
LLN))*5JJNN&%&NN-!7N8Ea -222a 0E999a /5888r   c                 *   | j                          | j                  d      } |j                  di t         |j	                  dgddddg       |j                  dg      }|d	   d
   d   dk(  sJ |d	   d
   d   dk(  sJ |d	   d
   d   dk(  sJ y )NrI  rD   rh   twojt @rC  rG   rK   r   rK   r   rD  rE  rF  r   )rT   rU   rV   rJ  r>  r   rL  s      r   "test_metadata_update_get_int_floatr^    s    
LLN))*5JJNN&%&G!"E%PQ   NNwN'Ea -222a 0E999a /5888r   r)  nested5rg   c                     | j                          | j                  d      }t        j                  t        d      5   |j
                  di t         d d d        y # 1 sw Y   y xY w)Ntest_metadata_validationr  r   r   )rT   rU   r   r   
ValueErrorrV   bad_metadata_recordsr   s     r   test_metadata_validation_addre    sQ    
LLN))*DEJ	z	4 /
.-./ / /   AA&c                    | j                          | j                  d      } |j                  d	i t         t	        j
                  t        d      5  |j                  dgdddii       d d d        y # 1 sw Y   y xY w)
Nrb  r  r   rD   r)  r_  r`  r]  r   )rT   rU   rV   rJ  r   r   rc  r>  r   s     r   test_metadata_validation_updaterh    ss    
LLN))*DEJJNN&%&	z	4 Mug'Hc?1KLM M Ms   A55A>c                     | j                          | j                  d      }t        j                  t        d      5  |j                  dddii       d d d        y # 1 sw Y   y xY w)Ntest_where_validationrS  r   r)  r_  r`  rR  )rT   rU   r   r   rc  r   r   s     r   test_where_validation_getrk    sV    
LLN))*ABJ	z	1 9g#789 9 9rf  c                     | j                          | j                  d      }t        j                  t        d      5  |j                  g ddddii       d d d        y # 1 sw Y   y xY w)	Nrj  rS  r   r   r   r   r)  r_  r`  )r   rS  )rT   rU   r   r   rc  rX   r   s     r   test_where_validation_queryrn    s`    
LLN))*ABJ	z	1 W)GhPS_;UVW W Ws   A  A)r\  r[  )rD  rF  rE  c                     | j                          | j                  d      } |j                  di t         |j	                  dddii      }t        |d         dk(  sJ y )	Ntest_where_ltrD  $ltrh   rR  rK   rM   r   rT   rU   rV   operator_recordsr   r=   rL  s      r   rp  rp    s^    
LLN))/:JJNN&%&NN+qz!:N;Eu[!"a'''r   c                     | j                          | j                  d      } |j                  di t         |j	                  dddii      }t        |d         dk(  sJ y )	Ntest_where_lterD  $lte       @rR  rK   rh   r   rr  rL  s      r   ru  ru  %  s_    
LLN))*:;JJNN&%&NN+}!=N>Eu[!"a'''r   c                     | j                          | j                  d      } |j                  di t         |j	                  dddii      }t        |d         dk(  sJ y )	Nru  rF  $gtgffffffrR  rK   rh   r   rr  rL  s      r   test_where_gtrz  -  s_    
LLN))*:;JJNN&%&NN-%!?N@Eu[!"a'''r   c                     | j                          | j                  d      } |j                  di t         |j	                  dddii      }t        |d         dk(  sJ y )	Nru  rF  $gter\  rR  rK   rM   r   rr  rL  s      r   test_where_gter}  5  s_    
LLN))*:;JJNN&%&NN-&%!ANBEu[!"a'''r   c                     | j                          | j                  d      } |j                  di t         |j	                  dddii      }t        |d         dk(  sJ y )	Nru  rE  $ner[  rR  rK   rM   r   rr  rL  s      r   test_where_ne_stringr  =  s_    
LLN))*:;JJNN&%&NN.5%.!ANBEu[!"a'''r   c                    | j                          | j                  d      } |j                  d
i t         |j	                  dddii      }t        |d         dk(  sJ |j	                  ddd	ii      }t        |d         dk(  sJ y )Nru  rD  r  rM   rR  rK   rF  $eqr\  r   rr  rL  s      r   test_where_ne_eq_numberr  E  s    
LLN))*:;JJNN&%&NN+qz!:N;Eu[!"a'''NN-%!@NAEu[!"a'''r   c                    | j                          | j                  d      } |j                  di t         t	        j
                  t              5  |j                  dddii       d d d        t	        j
                  t              5  |j                  dddii       d d d        t	        j
                  t              5  |j                  dddd	i       d d d        t	        j
                  t              5  |j                  d
dddiii       d d d        t	        j
                  t              5  |j                  ddidddiid       d d d        t	        j
                  t              5  |j                  ddddiidddiigi       d d d        t	        j
                  t              5  |j                  ddddiigi       d d d        t	        j
                  t              5  |j                  dg i       d d d        t	        j
                  t              5  |j                  dddii       d d d        t	        j
                  t              5  |j                  ddddiiddigi       d d d        y # 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   fxY w# 1 sw Y   3xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   y xY w)Ntest_where_valid_operatorsrD  z$invalidrh   rR  rq  2rM   )rq  ry  $andry  )rD  $orr  r8   	$containsrC   firstsecondr   )rT   rU   rV   rs  r   r   rc  r   r   s     r   r  r  O  s   
LLN))*FGJJNN&%&	z	" =kJ?;<= 
z	" :kE3<89: 
z	" Bk1Q+?@AB 
z	" Bf{UAJ&?@AB 
z	" 
!&
K%3LM 	 	


 
z	" 
K%4{UAJ6OPQ 	 	


 
z	" Ce{UAJ&?%@ABC 
z	" *eR[)* 
z	" ;cK#89:; 
z	" 
;01 (+ 	 	

 
A= =: :B BB B
 


 

C C* *; ;
 
sx   J 	JJ:J'4J41K1K,K"K'K3 J
JJ$'J14J>KKK$'K03K<)r   r   r   r   )r   r   r   r   r   r   r  c                 ,   | j                          | j                  d      } |j                  di t         t	        j
                  t              5 } |j                  di t         d d d        dt        j                        v sJ y # 1 sw Y   #xY w)Ntest_dimensionality_validation	dimensionr   )
rT   rU   rV   r   r   r   r   bad_dimensionality_recordsr   r)  r.   r\   es      r   "test_dimensionality_validation_addr    sx    
LLN))*JKJJNN%_%	y	! 5Q
4345#agg,&&&5 5   B

Bc                 ,   | j                          | j                  d      } |j                  di t         t	        j
                  t              5 } |j                  di t         d d d        dt        j                        v sJ y # 1 sw Y   #xY w)N$test_dimensionality_validation_queryr  r   )rT   rU   rV   r   r   r   r   rX   bad_dimensionality_queryr   r)  r  s      r   r  r    sz    
LLN))*PQJJNN%_%	y	! 5Q
4345#agg,&&&5 5r  c                    | j                          | j                  d      } |j                  di t         t	        j
                  t        d      5  |j                  dddii       d d d        t	        j
                  t        d      5  |j                  g dd	di
       d d d        t	        j
                  t        d      5  |j                  d	g i       d d d        t	        j
                  t        d      5  |j                  d	ddii       d d d        t	        j
                  t        d      5  |j                  dddii       d d d        t	        j
                  t              5  |j                  dddii       d d d        t	        j
                  t              5  |j                  dddiddigi       d d d        t	        j
                  t              5  |j                  dd	digi       d d d        t	        j
                  t              5  |j                  dg i       d d d        t	        j
                  t              5  |j                  ddd	digid	digi       d d d        y # 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   axY w# 1 sw Y   6xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   y xY w)Nr  zwhere documentr   rq  z$nestedrh   where_documentrm  r  )r   r  textrE   $not_containsr  z$unsupporteddocr  r   )	rT   rU   rV   rs  r   r   rc  r   rX   r   s     r   #test_query_document_valid_operatorsr    s   
LLN))*FGJJNN&%&	z)9	: ?uy!n&=>? 
z)9	: V)[RSDTUV 
z)9	: 9{B&789 
z)9	: H{VW4E&FGH 
z)9	: L8I&JKL 
z	" Iv/F&GHI 
z	" 
!^U$;ne=T#UV 	 	


 
z	" GuU/C.D&EFG 
z	" 3ubk23 
z	" 
;"6!78;:NO 	 	

 
A? ?V V9 9H HL LI I
 

G G3 3
 
sx   I<J	J J#:J02J=*K
&KK#K/<J	JJ #J-0J:=K
KK #K,/K8zthis is doc1 and it's great!zdoc2 is also great!)rJ   rH   rG   rK   c                 \   | j                          | j                  d      } |j                  di t         |j	                  ddi      }t        |d         dk(  sJ |j	                  ddi      }t        |d         dk(  sJ |j	                  dd	i      }t        |d         d
k(  sJ y )Ntest_get_where_documentr  doc1r  rK   rM   greatrh   badr   r   )rT   rU   rV   contains_recordsr   r=   rL  s      r   r  r    s    
LLN))*CDJJNN&%&NN;*?N@Eu[!"a'''NN;*@NAEu[!"a'''NN;*>N?Eu[!"a'''r   c                    | j                          | j                  d      } |j                  di t         |j	                  g dddid      }t        |d   d         dk(  sJ |j	                  g d	dd
id      }t        |d   d         dk(  sJ t        j                  t              5 }|j	                  g d	ddid      }dt        |j                        v sJ 	 d d d        y # 1 sw Y   y xY w)Ntest_query_where_document)rM   r   r   r  r  rM   )r   r  rP   rK   r   rm  r  rh   r  
datapointsr   )rT   rU   rV   r  rX   r=   r   r   r   r   r)  )r.   r\   rM  r  s       r   r  r    s   
LLN))*EFJJNN&%&"K3HTU  E u[!!$%***"K3IUV  E u[!!$%***	y	! ,Q  &U7KWX ! 
 s177|+++	, , ,s   -2C))C2c                 h   | j                          | j                  d      } |j                  d	i t         |j	                  ddi       |j                         dk(  sJ |j	                  ddi       |j                         dk(  sJ |j	                  ddi       |j                         dk(  sJ y )
Ntest_delete_where_documentr  r  r  rM   r  r  r   r   )rT   rU   rV   r  r   r   r   s     r   r  r    s    
LLN))*FGJJNN&%&k6%:;"""k5%9:"""k7%;<"""r   )g?g      @r   )gffffff?gGz@r   )rD   r@  id3id4r  )rD  rE  rF  is)rD  rF  rE  r  ri   g/$@three   gjt@four)z this document is first and greatz!this document is second and greatz this document is third and greatz!this document is fourth and great)rJ   rG   rK   rH   c           	         | j                          | j                  d      } |j                  di t         |j	                  dddddiiddd	iigid
digi      }t        |d         dk(  sJ |j	                  dddddiidddiigiddddiidddiigigi      }t        |d         dk(  sJ |j	                  dddddiidddiigiddddiidddiigigi      }t        |d         dk(  sJ y )Ntest_logical_operatorsr  r  rD  r|  ri   rF  rq  gffffff?r  r  rR  rK   r  rE  r  r  r  rh   rM   r[  rA  r   rT   rU   rV   logical_operator_recordsr   r=   rL  s      r   test_where_logical_operatorsr  $  s   
LLN))*BCJJNN.-.NN+{3meS\5RSTu
  E u[!"a'''NN$uaj1'%)9: $uaj1'%9
  E$ u[!"a'''NN$uaj1'%8 $uaj1'%8
  E$ u[!"a'''r   c                    | j                          | j                  d      } |j                  di t         |j	                  dddiddigi      }t        |d         dk(  sJ |j	                  d	ddidd
igi      }t        |d         dk(  sJ |j	                  d	ddidd
igidddii      }t        |d         dk(  sJ y )Ntest_document_logical_operatorsr  r  r  r  r  rK   rM   r  r  rh   rD  r  )r  rS  r   r  rL  s      r   %test_where_document_logical_operatorsr  \  s   
LLN))*KLJJNN.-.NNg&e$
  E u[!"a'''NNg&h'
  E u[!"a'''NNg&h'
 %
  
E u[!"a'''r   rm  this document is firstthis document is secondc                 x   | j                          | j                  d      } |j                  di t         g d}|j	                  g d|d      }|d   J |d   d   d   d	k(  sJ |d
   d   d   d   dk(  sJ t        |d         t        |      k(  sJ g d}|j	                  g d|d      }|d
   J |d   d   d   d	k(  sJ t        |d         t        |      k(  sJ |j	                  g dg dgg d      }|d   J |d
   J |d   J |d   J |d   d   d   d	k(  sJ |d   d   d   dk(  sJ y )Ntest_query_includerK   rH   rL   rm  rM   r   rQ   rP   rJ   rG   r   rD   rK   rD  rR   )rJ   rH   rL   )rM   rh   r   rh   rH   rL   r@  r   )rT   rU   rV   r   rX   rZ   r.   r\   rQ   rM  s       r   r  r    s   
LLN))*>?JJNNW5G"  E
 &&&<?1&&&a #K0A555uZ !S\1116G"  E
 %%%<?1&&&uZ !S\111#[1  E
 %%%%%%&&&%%%<?1&&&<?1&&&r   c                    | j                          | j                  d      } |j                  di t         ddg}|j	                  |ddi      }|d   J |d   d	   d
k(  sJ |d   d	   d   dk(  sJ |d   d	   dk(  sJ t        |d         t        |      k(  sJ ddg}|j	                  |      }|d   J |d   d	   d
k(  sJ t        |d   d   d	   d      sJ t        |d         t        |      k(  sJ |j	                  g       }|d   J |d   J |d   J |d   d	   d
k(  sJ |d   g k(  sJ t        j                  t        d      5  |j	                  ddg      }d d d        t        j                  t        d      5  |j	                  d       }d d d        y # 1 sw Y   @xY w# 1 sw Y   y xY w)Ntest_get_includerK   rH   rD  rM   )rQ   rS  rJ   rG   r   rD   r  rR   r   r   rQ   r   	undefinedr   )
rT   rU   rV   r   r   rZ   r;   r   r   rc  r  s       r   r  r    s%   
LLN))*<=JJNNWK(GNN7;2BNCE&&&<?e###a -222a $<<<<uZ !S\111[)GNN7N+E%%%<?e###l+A.q13777uZ !S\111NN2N&E%%%%%%&&&<?e###"""	z	3 C['ABC 
z	3 -t,- -C C- -s   F(F4(F14F=c                     | j                          | j                  d      } |j                  di t         |j	                  g dg dd      }|d   d   d   dk(  sJ |d   d   d	   d
k(  sJ y )Ntest_query_orderr   r  rh   r  rH   r   r  rM   r  r   rT   rU   rV   r   rX   rL  s      r   r  r    s    
LLN))*<=JJNNW)7  E a #'@@@@a #'????r   c                 l   | j                          | j                  d      }t        j                  t              5 }|j                  g ddgi g       d d d        dt        j                        v sJ t        j                  t              5 }|j                  d       d d d        dt        |j                        v sJ t        j                  t              5 }|j                  ddg       d d d        dt        |j                        v sJ y # 1 sw Y   xY w# 1 sw Y   ~xY w# 1 sw Y   ;xY w)	Ntest_invalid_idrm  rM   rG  IDr   validr   )
rT   rU   r   r   rc  rV   r   r)  r   r   r  s      r   r  r    s   
LLN))*;<J	z	" Fa)!EF3qww< 
z	" a13qww< 
z	" ,awl+,3qww<F F
 
, ,s#   DDD*DD'*D3c                 r   d}| j                          | j                  d      } |j                  di t         |j	                  g dd      }|d   d   d   d	kD  sJ | j                          | j                  dd
ddd      } |j                  di t         |j	                  g dd      }|d   d   d   d|z
  kD  sJ |d   d   d   d|z   k  sJ | j                          | j                  dddi      } |j                  di t         |j	                  g dd      }|d   d   d   dk  sJ y )Ng-q=test_index_paramsr  )333333?gQ?g?rM   r  rL   r   r  r        )r  zhnsw:construction_efzhnsw:Mr   r  r  ipr   r  )r.   EPSr\   rM  s       r   r  r    s   
C
LLN))/B)CJJNNW)  E a #a''' LLN))  ("PQR * J JNNW)  E a #a#g---a #a#g--- LLN)) L$+? * J JNNW)  E a #b(((r   c                     | j                          t        j                  t              5  | j	                  dddi      } |j
                  di t         d d d        y # 1 sw Y   y xY w)Nr  r  foobarr  r   )rT   r   r   r   rU   rV   r   r   s     r   test_invalid_index_paramsr  (  s_    
LLN	+	, "--$h/G . 

 	
!!	" " "s   ,AA(c                    |j                  d      } | j                          | j                  dddi      }|j                  dd       |j                  d	      }|j	                  d      }|j
                  d   dk(  sJ g d
}|j                  dd|      }|j                         D ]H  }||v s|dk(  rt        ||         dk(  rJ |dk(  rt        ||         t        |      k(  r@J ||   HJ  y )Nr/   rC   r  r  r  rD   rE   rF   r3   rI   rM   rN   rG   rR   )
rS   rT   rU   rV   rW   r  rX   rY   r=   rZ   )r.   r[   r\   r]   r^   r_   r`   s          r   !test_persist_index_loading_paramsr  2  s   $$%89F
LLN))% * J NNuN0""#ABD$$J |,444DH			 
 
B
 wwy #8Or#w<1$$$Jr#w<3x=000c7?"?#r   c                    | j                          | j                  d      }t        j                  j	                  dd      j                  t        j                        j                         }|j                  |t        t        |            D cg c]  }d| 	 c}       |j                         t        |      k(  sJ y c c}w )Nr   i     zhttp://example.com/r   )rT   rU   rl   randomrandastypefloat32r   rV   rk   r=   r   )r.   r\   large_recordsis       r   test_add_larger  P  s    
LLN))+6J IINN4-44RZZ@GGIMNN 05c-6H0IJ1"1#&J  
 ]!3333 Ks   C
c                 t    | j                          | j                         }dd l}|j                  d|      sJ y )Nr   z\d+\.\d+\.\d+)rT   get_versionrer   )r.   versionr  s      r   test_get_versionr  a  s3    
LLN  "G 88$g...r   c                    | j                          | j                  d      } |j                  di t         t	        | j                               dk(  sJ | j                  d       t	        | j                               dk(  sJ y )Ntest_delete_collectionrM   r   r   )rT   rU   rV   r   r=   r   r   r   s     r   r  r  l  st    
LLN))*BCJJNNWv&&()Q...
56v&&()Q...r   c                      t               } t        d      D cg c]  }d }} | |      }t        |      dk(  sJ y c c}w )N@   zthis is a test)r   rk   r=   )r   rq   docsrJ   s       r   test_default_embeddingr  w  sD    13&+Bi00D0#D)Jz?b    1s   	<c                    t         j                  j                  dd      j                  t         j                        j                         }t         j                  j                  dd      j                  t         j                        j                         }t        t        |            D cg c]  }d| 	 }}t        t        |            D cg c]  }d| 	 }}| j                          | j                  d      }|j                  ||       | j                  d      }|j                  ||       t        | j                               dk(  sJ |j                         t        |      k(  sJ |j                         t        |      k(  sJ |j                  |d	   d
      }|j                  |d	   d
      }	t        |d         d	kD  sJ t        |	d         d	kD  sJ t        |d   d	         d	kD  sJ t        |	d   d	         d	kD  sJ |d   d	   d	   |d	   k(  sJ |	d   d	   d	   |d	   k(  sJ y c c}w c c}w )Nr   r  zhttp://example.com/1/zhttp://example.com/2/coll1r   coll2rh   r   rM   r  rG   )rl   r  r  r  r  r   rk   r=   rT   rU   rV   r   r   rX   )
r.   embeddings1embeddings2r  ids1ids2r  r  results1results2s
             r   test_multiple_collectionsr  ~  s   ))..S)00<CCEK))..S)00<CCEK16s;7G1HIA#A3'IDI16s;7G1HIA#A3'IDI
LLN$$W-E	II$I/$$W-E	II$I/v&&()Q...;;=C,,,,;;=C,,,,{{KNa{HH{{KNa{HH x!###x!###xq!"Q&&&xq!"Q&&&E?1a DG+++E?1a DG+++1 JIs   -H6H;c                    | j                          | j                  d      } |j                  di t         t        d   d   gg dgdgddigd} |j                  di | |j                  |d	   d
g d      }t        |d   d         d
k(  sJ |d   d   d   |d   d   k(  sJ |d   d   d   |d   d   k(  sJ |d   d   d   d   dk(  sJ t        |d	   d   d   |d	   d         sJ y )Ntest_update_queryrG   r   r   皙?333333?zupdated documentfoobar)rG   rJ   rH   rK   rJ   rM   r   r   rH   rK   r   )rT   rU   rV   r   r>  rX   r=   r@   )r.   r\   updated_recordsresultss       r   r  r    sR   
LLN))*=>JJNNW q!"&'()en%	O J(( (68  G
 wu~a !Q&&&5>!Q?5#9!#<<<<;"1%)Ea)HHHH;"1%e,555a #_\%B1%E  r   c                 d   | j                          | j                  d      } |j                  d
i t         g d}|j	                  g dgd|      }|j                         D ]K  }||v s|dk(  rt        ||   d         dk(  r!J |d	k(  rt        ||         t        |      k(  rCJ ||   KJ  y )Nr   rI   r   r  r   rG   r   rh   rR   r   )rT   rU   rV   r   rX   rY   r=   rZ   )r.   r\   r^   r  r`   s        r   <test_get_nearest_neighbors_where_n_results_more_than_elementr    s    
LLN))+6JJNNWDH)*  G
 ||~ ((?cUlws|A'1,,,Jws|$H5553<'''(r   c                 F   | j                          | j                  d      } |j                  d	i t         t	        j
                  t              5 }|j                  g dgdg d       d d d        dt        j                        v sJ |j                  t        k(  sJ t	        j
                  t              5 }|j                  g dgdg d       d d d        dt        |j                        v sJ |j                  t        k(  sJ y # 1 sw Y   xY w# 1 sw Y   DxY w)
Nr   r   rI   r   z<Number of requested results -1, cannot be negative, or zero.rA  r   r   )rT   rU   rV   r   r   r   	TypeErrorrX   r   r)  typerc  )r.   r\   excs      r   test_invalid_n_results_paramr    s   
LLN))+6JJNNW	y	! 
S-.I 	 	

 JS		N    88y   	z	" 
c-.I 	 	

 C		N"""88z!!!%
 

 
s   D:DDD )g@gQ	@g@rD   r@  r  rE  )r  r  zthis document is third)      @r  r   )r   g(\@g@r  
one_of_onez this document is even more firstzthis document is new and fourthc                    | j                          | j                  d      } |j                  di t         |j	                         dk(  sJ  |j
                  di t         |j	                         dk(  sJ |j                  g dt        d   d         }t        |d   d   t        d   d         sJ |d	   d   t        d	   d   k(  sJ |d
   d   t        d
   d   k(  sJ |j                  |d   dg d      }t        |d   d   d   t        d   d         sJ |d	   d   d   t        d	   d   k(  sJ |d
   d   d   t        d
   d   k(  sJ |j                  t        d   d          |j                  t        d   d   g dgddig       |j	                         dk(  sJ |j                  g ddg      }t        |d   d   g d      sJ |d	   d   ddik(  sJ |d
   d   J y )NrC   ri   r  r5  rG   r   )rQ   rG   rJ   rK   rH   rM   r   rh   r   )r   Gz?gGz@rE  za new string value)rG   rJ   rK   r  r   )rT   rU   rV   initial_recordsr   upsertnew_recordsr   r@   rX   r   )r.   r\   
get_resultquery_results       r   test_upsertr    sy   
LLN))&1JJNN%_%"""J$$"""8k%>PQR>S   J < #[%>q%A   k"1%[)A!)DDDDk"1%[)A!)DDDD###L18 $ L
 \"1%a(+l*CA*F   $Q'*k+.Fq.IIII$Q'*k+.Fq.IIII/%034E"1%%&"$89:  
 """8ug   J z,7:<MNNNk"1%.:N)OOOOk"1%---r   c                     | j                          | j                  d      }| j                  d       t        j                  t
        d      5   |j                  di t         d d d        y # 1 sw Y   y xY wr   )rT   rU   r   r   r   r   r  r
  r   s     r   5test_collection_upsert_with_invalid_collection_throwsr  1  s`    
LLN))&1J
V$	},K	L -
,O,- - -r   c                 z   | j                          | j                  d      }g dg dgddgd}t        j                  t              5 } |j
                  di | d d d        dt        j                        v sJ t        j                  t              5 }|j                  g dgd	
       d d d        dt        |j                        v sJ dgdgdggdgdgdgggddgd}t        j                  t              5 } |j                  di | d d d        dt        |j                        v sJ g dgg dggddgd}t        j                  t              5 } |j                  di | d d d        dt        |j                        v sJ y # 1 sw Y   3xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   HxY w)Ntest_invalid_embeddings)0r  r  )z1.2z2.243.2rD   r@  r   	embedding)z1.1z2.3r  rM   r  r   r   r   r   r   r   r   )rT   rU   r   r   rc  rV   r   r)  rX   r>  r  )r.   r\   invalid_recordsr  s       r   r  r  =  s   
LLN))*CDJ '(>?u~O 
z	" *a
))*#agg,&&& 
z	" 
a34 	 	


 #agg,&&& aS1##(>?u~O 
z	" -a
,O,-#agg,&&& ((+;*<=u~O 
z	" -a
,O,-#agg,&&&9* *

 
- -- -s0   FF;F%F1FF"%F.1F:c                 ,   | j                          | j                  d      } |j                  di t         t	        j
                  t              5 } |j                  di t         d d d        dt        j                        v sJ y # 1 sw Y   #xY w)N$test_dimensionality_update_exceptionr  r   )rT   rU   rV   r   r   r   r   r>  r  r   r)  r  s      r   $test_dimensionality_exception_updater  h  z    
LLN))*PQJJNN%_%	y	! 8Q
7678#agg,&&&8 8r  c                 ,   | j                          | j                  d      } |j                  di t         t	        j
                  t              5 } |j                  di t         d d d        dt        j                        v sJ y # 1 sw Y   #xY w)N$test_dimensionality_upsert_exceptionr  r   )rT   rU   rV   r   r   r   r   r  r  r   r)  r  s      r   $test_dimensionality_exception_upsertr  u  r  r  win32)reruns	conditionc                     t         j                  j                  d      rt        j                  d       | j                          y )NCHROMA_INTEGRATION_TEST_ONLY"Skipping test for integration test)r(   environr   r   r   r   )
client_ssls    r   test_ssl_self_signedr'    s,    	zz~~4589r   c                 0   t         j                  j                  d      rt        j                  d       | j                          | j                  j                  j                  }t        j                  t              5 }t        j                  d|       d d d        t        j                  t        j                         |j                   |j                   j"                        }| j%                          ddj'                  |      v sJ y # 1 sw Y   wxY w)Nr#  r$  T)sslportCERTIFICATE_VERIFY_FAILED )r(   r%  r   r   r   r   _server	_settingschroma_server_http_portr   rc  r%   
HttpClient	tracebackformat_exceptionr  r)  __traceback__r'   join)r&  _portr  stack_traces       r   'test_ssl_self_signed_without_ssl_verifyr7    s    	zz~~4589((@@E	z	" 2a512,,QWWqww 5 5K !!#&"''+*>>>>2 2s   >DDc                    | j                          | j                  d      }d}d}t        j                  j	                  dd      j                  t        j                        j                         }t        |      D cg c]  }|  }}|j                  ||       t        d|d      D cg c]  }|  c}t        j                  j	                  |      j                  t        j                        j                         }|j                  ||g       }|d   D 	
cg c]  }	|	D ]  }
|
  }}	}
t        fd	|D              sJ y c c}w c c}w c c}
}	w )
Ntest_query_id_filtering_smallr   r  r   r   r   r   rG   rP   rQ   rG   c              3   &   K   | ]  }|v  
 y wr   r   .0id	query_idss     r   	<genexpr>z8test_query_id_filtering_small_dataset.<locals>.<genexpr>       :2rY:   rT   rU   rl   r  r  r  r  r   rk   rV   rX   r>   )r.   r\   num_vectorsdimsmall_recordsr  rG   query_embeddingr  sublistitemall_returned_idsr?  s               @r   %test_query_id_filtering_small_datasetrK    s;   
LLN))*IJJK
CIINN3,33BJJ?FFHM -
.aaS6
.C
.NN   
 "'q+r!:;AA3;IiinnS)00<CCEO(	  G -4ENOwOtOOO:)9::::# / < Ps   ?
E-
E!Ec                    | j                          | j                  d      }d}d}t        j                  j	                  ||      j                  t        j                        j                         }t        |      D cg c]  }|  }}|j                  ||       t        d|d      D cg c]  }|  c}t        j                  j	                  |      j                  t        j                        j                         }|j                  ||g       }|d   D 	
cg c]  }	|	D ]  }
|
  }}	}
t        fd	|D              sJ t        d
      D cg c]L  }t        j                  j	                  |      j                  t        j                        j                         N }}|j                  |dg       }|d   D ]  }t        fd|D              rJ  y c c}w c c}w c c}
}	w c c}w )Ntest_query_id_filtering_medium  r  r   r   r   r:  rG   c              3   &   K   | ]  }|v  
 y wr   r   r<  s     r   r@  z9test_query_id_filtering_medium_dataset.<locals>.<genexpr>  rA  rB  ri   c              3   &   K   | ]  }|v  
 y wr   r   r<  s     r   r@  z9test_query_id_filtering_medium_dataset.<locals>.<genexpr>  s     8r2?8rB  rC  )r.   r\   rD  rE  medium_recordsr  rG   rG  r  rH  rI  rJ  rq   multi_query_embeddingsmulti_results
result_setr?  s                   @r   &test_query_id_filtering_medium_datasetrU    s   
LLN))*JKJK
CYY^^K5<<RZZHOOQN -
.aaS6
.C
.NN!  
 "'q+r!:;AA3;IiinnS)00<CCEO(	  G -4ENOwOtOOO:)9:::: BGq<=		s""2::.557  $$/	 % M $E* 9
8Z88889= / < Ps   ?
G -
G%!G*AG0c                 L   | j                          | j                  d      }d}d}t        j                  j	                  ||      j                  t        j                        j                         }t        |      D cg c]  }|  }}t        |      D cg c]  }d|i }}|j                  |||       t        dd      D cg c]  }|  }}|j                  |       t        dd	      D cg c]  }|  }	}|d
z   }
t        ||
      D cg c]  }|  }}t        j                  j	                  t        |	      t        |      z   |      j                  t        j                        j                         }t        t        |            D cg c]  }|dd	 }}|j                  ||	|z   |       t        dd      D cg c]  }|  c}t        dd      D cg c]  }|  c}z   t        |dz   |dz         D cg c]  }|  c}z   dg}t        j                  j	                  |      j                  t        j                        j                         }|j                  ||
|      }|d   D cg c]  }|D ]  }|  }}}t        fd|D              sJ t        |d         D ]1  \  }}t        |      D ]  \  }}||	v s||v s|d   |   |   d   rJ  3 |	d   }|j                  ||d|      }|d   d   d   d   sJ |d   }t!        j"                  t$              5 }|j                  ||d|       d d d        dt'        j(                        v sJ y c c}w c c}w c c}w c c}w c c}w c c}w c c}w c c}w c c}w c c}}w # 1 sw Y   VxY w)Ntest_query_id_filtering_e2er  r   indexrG  r      r   2   r  T)rX  upsertedr  #   -      rK   r:  rG   c              3   &   K   | ]  }|v  
 y wr   r   )r=  r>  valid_query_idss     r   r@  z.test_query_id_filtering_e2e.<locals>.<genexpr>  s     @r_$@rB  r[  r   rM   zError finding id)rT   rU   rl   r  r  r  r  r   rk   rV   r   r=   r  rX   r>   	enumerater   r   r   r   r)  )r.   r\   rE  rD  rJ   r  rG   rK   ids_to_deleteids_to_upsert_existingnew_num_vectorsids_to_upsert_newupsert_embeddingsupsert_metadatasr^   rG  r  rH  rI  rJ  result_indexid_list
item_indexitem_idupserted_id
deleted_idr,  r`  s                              @r   rW  rW    s   
LLN))*GHJ
CKS188DKKMJ -
.aaS6
.C
.',['9:!'1:I:NN   &+2r]3sV3M3-( /4Bm<f<<!B&O).{O)LMAA3MM 			s12S9J5KKSQ	

		  16c:K6L0M+,!&  $"%66"   q"&AA3& R=
)aaS6
)	* "+/;3CD
qcF
	
  }HiinnS)00<CCEO(!	  G -4ENOwOtOOO@/?@@@@!*75>!: Rg#,W#5 	RJ00G?P4P{+L9*EjQQQ	RR )+K(	  G ;"1%j111q!J	y	! 
U,	 	 	

 U[[!1111[ /: 4 =M 	'
)
 P*
 
sH   ?
M'M,
M13
M6
M;N 
N!
N

N<N/NN#c                     ddl m}   | g dg d        | g g        t        j                  t        d      5   | dd	d
g       ddd       t        j                  t        d      5   | ddgd       ddd       t        j                  t        d      5   | g dd	d
g       ddd       t        j                  t        d      5   | g dg d       ddd       t        j                  t        d      5   | g dg d       ddd       t        j                  t        d      5   | g dg d       ddd       t        j                  t        d      5   | g dg d       ddd        | g dg d        | g dg d        | g dg d       t        j                  t        d      5   | ddgd	dg       ddd       t        j                  t        d      5   | ddgd	d
g       ddd        | dgdg        | ddgddg       t        j                  t        d      5   | g d g d       ddd       t        j                  t        d      5   | g d!g d"       ddd       t        j                  t        d      5   | g d#g d$       ddd       y# 1 sw Y   sxY w# 1 sw Y   PxY w# 1 sw Y   +xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   HxY w# 1 sw Y   #xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   yxY w)%z.Test SparseVector validation in __post_init__.r   SparseVectorr   rh   r  r         ??indicesvaluesz*Expected SparseVector indices to be a listr   
not_a_listr   r  Nz)Expected SparseVector values to be a listrM   ,indices and values must have the same lengthr   rM   rh   z%SparseVector indices must be integers)r   not_intrh   r  )SparseVector indices must be non-negativer   r   rh   #SparseVector values must be numbers)r   
not_numberr  )              ?rw  rg   )rM         @ri   )r   rN  i'  *   gQ	@TF2indices must be sorted in strictly ascending order)r   rh   rM   )r   rM   rM   rh   )r   r  r  皙?)r  ri   rM   )rs  r  r   )chromadb.base_typesrp  r   r   rc  ro  s    r   test_validate_sparse_vectorr  2  s   0 ?; B' 
z)U	V >\3*=> 
z)T	U :aVL9: 
H
 ; 	YSz:; 
z)P	Q H.GH 
z)T	U AZ@A 
z)N	O IY/GHI 
z)P	Q F__EF 95 ;7 +OD 
z)N	O 9aVS$K89 
z)P	Q ;aYSz:; "tf- !Qu6 
N
 @ 	Y?@ 
N
 H 	\2FGH 
N
 @ 	Y?@ @C> >: :; ;H HA AI IF F9 9; ;@ @H H@ @s   J8*KKK>K,0K9"L>L0L <L-.L9 M8KKKK),K69LLL L*-L69MMc            	      :   ddl m}  ddlm}  |g dg d      } |g dg d      }d	|d
d}d|dd} | |        | |       d |g g       d} | |       t	        j
                  t        d      5  d |ddgdg      d}ddd       ddddgid}t	        j
                  t        d      5   | |       ddd       t	        j
                  t        d      5  d |g dg d      d}	ddd       t	        j
                  t        d      5  d  |ddgdd!g      d}
ddd       d" |ddgddg       |g d#g d$      d%d&} | |       d'g d(g d)dd*}t	        j
                  t        d      5   | |       ddd        |t        t        d+            t        d+      D cg c]  }t        |      d,z   c}      }d-|d.} | |       y# 1 sw Y   UxY w# 1 sw Y   -xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   xY wc c}w )/z<Test that sparse vectors are properly validated in metadata.r   )validate_metadataro  rq  rr  ru  rM   ri   r  )r  r  r  z
document 1rs  )r  sparse_embeddingscorez
document 2皙?zempty sparse)r  
sparse_vecry  r   invalidrM   r   )r  r  Nzmissing indicesrw  r  zLExpected metadata value to be a str, int, float, bool, SparseVector, or Noner|  znegative indexr}  r  r~  znon-numeric valuenot_a_numberzmultiple sparse vectors)rh   ri   r  )r  r  rs  r  )r  sparse_1sparse_2regular_fieldsome_configrz  r  rw  r  )r   sparse_vectorrN  gMbP?zlarge sparse)r  large_sparse_vec)
chromadb.api.typesr  r  rp  r   r   rc  listrk   rK  )r  rp  sparse_vector_1sparse_vector_2
metadata_1
metadata_2metadata_emptyinvalid_metadatainvalid_metadata_2invalid_metadata_3invalid_metadata_4metadata_multiplemetadata_nestedr  large_sparsemetadata_larges                   r   )test_sparse_vector_in_metadata_validationr    sW   40 #9_MO"9_MO +J +J
 j!j! "2b9N n% 
H
 
  ,aVSE J

 "%Sz2 
\
 . 	,-	. 
z)T	U 
$ ,"?!

 
z)N	O 
' ,aVS.DY Z

 * !Qc
C ?K	 '(  %./JO 
\
 + 	/*	+  U4[!*/+6Qa5 6L -,ONn%{
 
. .
 

 
(+ + 7sB   5G2	G&G3H ;	H.H
G#&G03G= H	Hc                     ddl m} m} ddlm} dt
        t        dg ddg did	} | |      }t        |d
   |      sJ |d
   j                  g dk(  sJ |d
   j                  g dk(  sJ  ||        |g dg d      }d|d	} | |      }|d
   |u sJ  ||       dt
        t        dg ddg did	}t        j                  t        d      5   | |       ddd       dt
        t        dg ddg did	}	t        j                  t        d      5   | |	       ddd       dt
        t        dg ddg did	}
t        j                  t        d      5   | |
       ddd       dt
        t        dddgdg did	}t        j                  t        d      5   | |       ddd       dddid} | |      }t        |d    t              sJ |d    d   dk(  sJ d!t
        t        dg dg id	} | |      }t        |d
   |      sJ |d
   j                  g k(  sJ |d
   j                  g k(  sJ t
        t        ddd"gdd#d$git
        t        ddd%gdd&d'gid(d)} | |      }t        |d*   |      sJ t        |d+   |      sJ |d   d(k(  sJ y# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   KxY w# 1 sw Y   xY w),zNTest that dict-format sparse vectors are normalized to SparseVector instances.r   )normalize_metadatar  ro  ztest documentrv  rq  rw  r  r  sparser  r  )rs        ?r  ru  unsorted)r  r   rh   )r  r  rw  r  r   N
duplicates)r   rh   rh   negative)r   r   rh   zindices must be non-negativemismatchrh   ry  regularr`   r)  )r  r   r   emptyrM   r  rw  ri   r        @r  )sparse1sparse2r  r  r  )r  r  r  r  rp  r
   r   r   rv  rw  r   r   rc  r   )r  r  rp  metadata_dict_format
normalizedsparse_instancemetadata_instance_formatnormalized2metadata_unsortedmetadata_duplicatesmetadata_negativemetadata_mismatchmetadata_regular_dictnormalized3r  normalized4r  normalized5s                     r   ,test_sparse_vector_dict_format_normalizationr    sy   H0  .yo
 $$89Jj*L999h''9444h&&/999 j! #9_MO!  %%=>Kx O333k" .yo
 
N
 . 	,-. .yo
 
N
 0 	./0 .zo
 
z)G	H .,-.
 .1vo
 
H
 . 	,-. '" %%:;Kk(+T222x '7222 5y"hPRSN %^4Kk(+\:::x ((B...x ''2---
 .1vsCj
 .1vsCj

  %%67Kk),l;;;k),l;;;y!R'''c. .0 0. .. .s0   	J	J	J)	J6JJ&)J36K c                     ddl m} m} ddlm}  | g dddt
        t        dddgd	d
dgidd |ddgddg      dddigg d      }t        |d   d   d   |      sJ t        |d   d   d   |      sJ d|d   d   vsJ  ||       |d   d   d   j                  ddgk(  sJ |d   d   d   j                  d
dgk(  sJ |d   d   d   j                  ddgk(  sJ |d   d   d   j                  ddgk(  sJ y)zITest that dict-format sparse vectors work in normalize_insert_record_set.r   )normalize_insert_record_setvalidate_insert_record_setro  r  doc2doc3Ntest1rv  rh   rw  r  rw  r  r  rM   ri   r  r  ru  r  test3)zdoc onezdoc twoz	doc three)rG   rJ   rK   rH   rK   r  )
r  r  r  r  rp  r
   r   r   rv  rw  )r  r  rp  
record_sets       r   ,test_sparse_vector_dict_format_in_record_setr  a  su    1 -$  61vsCj  &1vsCjI W
 6%J, j-a0:LIIIj-a0:LIII:k215555 z* k"1%h/77Aq6AAAk"1%h/663*DDDk"1%h/77Aq6AAAk"1%h/663*DDDr   c            
         ddl m}   | g dgg dgddgddgd	d
gggddiddiddiggg dgg dgd      }|j                         }t        |      dk(  sJ t        |d         dk(  sJ |d   d   d   dk(  sJ |d   d   d   dk(  sJ |d   d   d   ddgk(  sJ |d   d   d   ddik(  sJ |d   d   d   dk(  sJ |d   D ]   }d|v sJ d|v sJ d|v sJ d|v sJ d|v r J   | ddgg dgdd gg d!gd"dgdgdgggd#did#d$igd"gd%d&gg dgddgddggd      }|j                         }t        |      d$k(  sJ t        |d         d$k(  sJ t        |d         dk(  sJ |d   d   ddd#did%d'k(  sJ |d   d   dd d#d$id&d'k(  sJ |d   d   d(d)dgdd*k(  sJ |d   d   d+d,dgd-d*k(  sJ |d   d$   d.d/dgd0d*k(  sJ  | g g g g g g d      }|j                         }|g k(  sJ  | g dgg d1gd"ddid"ddiggg d2gg d3gd      }|j                         }t        |      dk(  sJ t        |d         dk(  sJ |d   d   dddidd4k(  sJ |d   d   d5d6d7k(  sJ |d   d$   d8ddid0d4k(  sJ  | dd5ggd"d"d"d"g gd      }|j                         }t        |      dk(  sJ t        |d         d$k(  sJ |d   d   ddik(  sJ |d   d   dd5ik(  sJ  | d9ggd:ggd9d;iggd<d=gggd>ggd?ggd@      }|dA   d9ggk(  sJ |j	                  dB      d:ggk(  sJ dC|v sJ t        |      dDk(  sJ |j                         }t        |d         dk(  sJ |d   d   d   d9k(  sJ t        dE       y")FzTTest the SearchResult.rows() method for converting column-major to row-major format.r   )SearchResultr  r  r  rw  r  r        @g      @r`   r8   r9   r   )rt  r  ffffff?)documentr  r  )rG   rH   rJ   rK   scoresselectrM   ri   r>  rD   r  r  r  r  r  rt  a1a2)b1b2b3doc_a1doc_a2)doc_b1doc_b2doc_b3Nxrh   rs  r  )r>  r  r  r  r  r  )r>  r  r  r  r  r  r  r  r  r  )Nr  N)rt  Nr  )r  r  r  )r>  r  r  r@  r  )r>  r  r  rC   ztest docTr   r  r	  r>   )rG   rH   rK   rJ   r  r  rG   rH   rK      z%All SearchResult.rows() tests passed!)r  r  rowsr=   r   r7  )r  resultr  rows       r   test_search_result_rowsr    s   / )*23 #Jc
S#J?@!3<%s|DE&'89	
	F ;;=Dt9>>tAw<1 71:du$$$71:j!V+++71:k"sCj00071:j!eS\11171:g#%%% Aw s{{S   c!!!S   #~~ 4L"45#X.0NOu%  8c1X.5Sz?3"G,{G.DE
	
F ;;=Dt9>>tAw<1tAw<1 71:!H	    71:!H	    71:U	    71:U	    71:U	    	
	F ;;=D2:: )*./8TC845'(89	
	F ;;=Dt9>>tAw<1 71:C8cJJJJ 71:6:::: 71:C8cJJJJ EN#d	
	F ;;=Dt9>>tAw<171:$&&&71:$&&& H:%,"D>*+ #J<.vhwi	
	F %=fXJ&&&::k"
|n444&   v;! ;;=DtAw<171:dv%%%	
12r   c                     ddl } ddlm}m}m}  | |ddgd       |dd	gd
d      g      }|j                         }dddiddddidddidddgddddigididddidddiddd	gd
dddigidigigi}||k(  sJ  | |ddgd       |dd	gd
d      gddgd      }|j                         }dddiddddidddidddgddddigididddidddiddd	gd
dddigidigigi}	||	k(  sJ  | |ddgd       |dd	gd
d       |d      g      }
|
j                         }d|v sJ d|d   d   v sJ |d   d   d   }t        |      dk(  sJ  | j                  t        d      5   | |ddgd       |dd	gd      gdg       }|j                          ddd        | j                  t        d!      5   | |ddgd       |dd	gd      gdd"g       }|j                          ddd        | j                  t        d#      5   |g       }|j                          ddd        | j                  t        d$      5   | |d      gd%&      }|j                          ddd        | j                  t        d$      5   | |d      gd&      }|j                          ddd        | |ddgd       |dd	gd
d      gd'dgdd(      }|j                         }dddidddd)idddidddgddddigididdd*idddiddd	gd
dddigidigigi}||k(  sJ  | |ddgd       |dd	gd      gd+      }|j                         }dddidddd,idddidddgddddigididdd,idddiddd	gddddigidigigi}||k(  sJ  | j                  t        d-      5   | |ddgd       |dd	gd      gd.d.gd/      }|j                          ddd       t        d0       y# 1 sw Y   ]xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   VxY w)1z9Test the Rrf (Reciprocal Rank Fusion) to_dict conversion.r   N)RrfKnnValr   r  T)rX   return_rankr  r  r  )rX   r`   r  $mul$valr   $sum$divr  <   $knn
#embedding   rX   r`   limitr  leftrightrw  r   )ranksweightskr  rM   ri   z/Number of weights .* must match number of ranksr   )r  r  z All weights must be non-negativeg      zRRF requires at least one rankzk must be positiver  )r  r  )r  r  	normalizer  g      ?g      ?)r  r  rs  zSum of weights must be positiver  )r  r  r  zAll RRF tests passed!)
r   &chromadb.execution.expression.operatorr  r  r  to_dictr=   r   rc  r7  )r   r  r  r  rrfr  expectedrrf_weightedresult_weightedexpected_weighted	rrf_threeresult_threetermsrrf_badrrf_negative	rrf_empty	rrf_neg_k
rrf_zero_krrf_normalizedresult_normalizedexpected_normalizedrrf_normalize_defaultsresult_defaultsexpected_defaultsrrf_zero_weightss                            r   test_rrf_to_dictr  ;	  s    DD sCjd3sCj&8dK	
C [[]F 	RL%+SM &%+RL(.693Z3?57;?	1*%&
)"&!& %+SM &%+RL(.693Z3E57;?	1*%&
)"&!'%'*
,H\ X sCjd3sCj&8dK
 c

L #**,O 	RL%+SM &%+SM(.693Z3?57;?	1*%&
)"&!& %+SM &%+SM(.693Z3E57;?	1*%&
)"&!'%'*
,\ //// sCjd3sCj&8dKH	
I $$&L \!!!\&)!,,,, #F+Eu:?? 
K
 
 3*$73*$7 E
 	
 
z)K	L 3*$73*$7 $K
 	 
z)I	J G	
 
z)=	> S
b)	
 
z)=	> #c(q)
 sCjd3sCj&8dK
 c

N '..0 	RL%+TN &%+SM(.693Z3?57;?	1*%&
)"&!& %+TN &%+SM(.693Z3E57;?	1*%&
)"&!'%'*
,\  3333 !sCjd3sCjd3
  -446O 	RL%+SM &%+RL(.693Z3?57;?	1*%&
)"&!& %+SM &%+RL(.693Z3?57;?	1*%&
)"&!'%'*
,\ //// 
z)J	K 	#3*$73*$7 #J
 	  "	# 

!"W
 
  
 
 |	# 	#sH   +2O8=3PP	"P"P,04P98PPPP),P69Qc                   :    e Zd ZdZd Zd Zd Zd Zd Zd Z	d Z
y	)
TestSearchDictSupportz%Test Search class dict input support.c                     ddl m} ddlm}  |ddi      }|j                  J t        |j                  |      sJ  |dddid	d
diigi      }|j                  J y)z-Test Search accepts dict for where parameter.r   SearchWherestatusactiverR  Nr  r  ry  rs  )"chromadb.execution.expression.planr  r  r  _wherer   )ro   r  r  searchs       r   test_search_with_dict_wherez1TestSearchDictSupport.test_search_with_dict_where
  sr    =@ x23}}(((&--/// v8(<wPS>U'VWX}}(((r   c                     ddl m} ddlm}  |ddddgii      }|j                  J t        |j                  |      sJ  |d
di      }|j                  J y	)z,Test Search accepts dict for rank parameter.r   r  Rankr  rX   r   r  rankNr  rs  )r  r  r  r!  _rankr   )ro   r  r!  r  s       r   test_search_with_dict_rankz0TestSearchDictSupport.test_search_with_dict_rank
  sf    =? fwc
&;<=||'''&,,--- fc]+||'''r   c                    ddl m}  |ddd      }|j                  j                  dk(  sJ |j                  j                  dk(  sJ  |d      }|j                  j                  dk(  sJ |j                  j                  dk(  sJ y)z5Test Search accepts dict and int for limit parameter.r   r  r   r  r  offsetr  N)r  r  _limitr  r(  ro   r  r  s      r   test_search_with_dict_limitz1TestSearchDictSupport.test_search_with_dict_limit
  s    = a89}}""b(((}}##q((( b!}}""b(((}}##q(((r   c                     ddl m}  |dddgi      }|j                  J  |ddg      }|j                  J  |dd	h      }|j                  J y)
z=Test Search accepts dict, list, and set for select parameter.r   r  rY   	#document#scorer  N	#metadatar  )r  r  _selectr+  s      r   test_search_with_dict_selectz2TestSearchDictSupport.test_search_with_dict_select
  sm    = h(?@A~~))) [9:~~))) \:;~~)))r   c                     ddl m} ddlm}  | |d      dk(  dddd	giid
dg      }|j                  J |j
                  J |j                  j                  d
k(  sJ |j                  J y)z2Test Search with mixed expression and dict inputs.r   r  )Keyr  r  r  rX   r   r  r   r.  rS  r#  r  r  N)	r  r  r  r5  r  r$  r*  r  r2  )ro   r  r5  r  s       r   test_search_mixed_inputsz.TestSearchDictSupport.test_search_mixed_inputs
  s    =>h-8+7S#J/0=	
 }}(((||'''}}""b(((~~)))r   c                     ddl m}  |       j                  ddi      j                  ddi      }|j                  J |j
                  J y)z)Test Search builder methods accept dicts.r   r  r  r  r  rs  N)r  r  rS  r#  r  r$  r+  s      r   &test_search_builder_methods_with_dictsz<TestSearchDictSupport.test_search_builder_methods_with_dicts
  sI    =8 45::FC=I}}(((||'''r   c                    ddl }ddlm}  |j                  t        d      5   |d       ddd        |j                  t        d      5   |d	
       ddd        |j                  t        d      5   |d       ddd        |j                  t        d      5   |d       ddd       y# 1 sw Y   xY w# 1 sw Y   jxY w# 1 sw Y   LxY w# 1 sw Y   yxY w)z(Test Search rejects invalid input types.r   Nr  zwhere must ber   r  rR  zrank must bers  r"  zlimit must be10r)  zselect must be{   r0  )r   r  r  r   r  )ro   r   r  s      r   test_search_invalid_inputsz0TestSearchDictSupport.test_search_invalid_inputs
  s    =V]]9O< 	$#	$ V]]9N; 		 V]]9O< 		 V]]9,<= 	#	 		$ 	$	 		 		 	s/   
B4
C 7
C!
C4B= C	CC!N)r   r   r   __doc__r  r%  r,  r3  r7  r9  r=  r   r   r   r  r  
  s(    /)()* * (r   r  c                   @    e Zd ZdZd Zd Zd Zd Zd Zd Z	d Z
d	 Zy
)TestWhereFromDictz"Test Where.from_dict() conversion.c                     ddl m}m} |j                  ddi      }t	        ||      sJ |j                  dddii      }t	        ||      sJ y)z Test simple equality conversion.r   )r  Eqr  r  r  N)r  r  rB  	from_dictr   )ro   r  rB  rS  s       r   test_simple_equalityz&TestWhereFromDict.test_simple_equality
  sT    D 8 45%$$$ E8+< =>%$$$r   c                    ddl m}m}m}m}m}m} |j                  dddii      }t        ||      sJ |j                  dddii      }t        ||      sJ |j                  dd	dii      }t        ||      sJ |j                  dd
dii      }t        ||      sJ |j                  dddii      }t        ||      sJ y)z%Test comparison operator conversions.r   )r  NeGtGteLtLter  r  inactiver  ry  rs  r|  rq  r  rv  N)	r  r  rF  rG  rH  rI  rJ  rC  r   )ro   r  rF  rG  rH  rI  rJ  rS  s           r   test_comparison_operatorsz+TestWhereFromDict.test_comparison_operators
  s    VV E:+> ?@%$$$ 5#, 78%$$$ 63- 89%%%% 5#, 78%$$$ 63- 89%%%%r   c                     ddl m}m}m} |j	                  ddddgii      }t        ||      sJ |j	                  dddd	gii      }t        ||      sJ y
)z%Test membership operator conversions.r   )r  InNinr  $inr  pendingz$nindeletedarchivedN)r  r  rN  rO  rC  r   )ro   r  rN  rO  rS  s        r   test_membership_operatorsz+TestWhereFromDict.test_membership_operators  sf    II EHi3H+I JK%$$$ FY
4K+L MN%%%%r   c                 8   ddl m}m}m}m}m} |j                  dddii      }t        ||      sJ |j                  dddii      }t        ||      sJ |j                  ddd	ii      }t        ||      sJ |j                  dd
dii      }t        ||      sJ y)z!Test string operator conversions.r   )r  ContainsNotContainsRegexNotRegexr  r  rE   r  spamz$regexz^test.*z
$not_regexz\d+N)r  r  rV  rW  rX  rY  rC  r   )ro   r  rV  rW  rX  rY  rS  s          r   test_string_operatorsz'TestWhereFromDict.test_string_operators  s    	
 	
 +w)? @A%*** /6)B CD%--- (I)> ?@%''' ,)? @A%***r   c                     ddl m}m}m} |j	                  dddidddiigi      }t        ||      sJ |j	                  d	ddidd
igi      }t        ||      sJ y)z"Test logical operator conversions.r   )r  AndOrr  r  r  r  ry  rs  r  rQ  N)r  r  r]  r^  rC  r   )ro   r  r]  r^  rS  s        r   r  z(TestWhereFromDict.test_logical_operators3  sy    II x*Wucl,CDE
 %%%% (H)=)?T(U VW%$$$r   c                 r    ddl m}m} |j                  ddddiddigidd	d
iigi      }t	        ||      sJ y)zTest nested logical operations.r   )r  r]  r  r  r  r  rQ  r  r|  rs  N)r  r  r]  rC  r   )ro   r  r]  rS  s       r   test_nested_logical_operatorsz/TestWhereFromDict.test_nested_logical_operatorsA  sU    Eh1Hi3HIJvsm,
 %%%%r   c                 ^    ddl m}m} |j                  ddddgii      }t	        ||      sJ y)zTest special key handling.r   )r  rN  z#idrP  rD   r@  N)r  r  rN  rC  r   )ro   r  rN  rS  s       r   test_special_keysz#TestWhereFromDict.test_special_keysO  s3    D (? @A%$$$r   c                    ddl }ddlm}  |j                  t        d      5  |j                  d       ddd        |j                  t        d      5  |j                  i        ddd        |j                  t        d      5  |j                  d	g i       ddd       y# 1 sw Y   pxY w# 1 sw Y   JxY w# 1 sw Y   yxY w)
zTest invalid Where dict inputs.r   Nr  Expected dictr   
not a dictcannot be emptyzrequires at least one conditionr  )r   r  r  r   r  rC  rc  )ro   r   r  s      r   test_invalid_where_dictsz*TestWhereFromDict.test_invalid_where_dictsW  s    @V]]9O< 	*OOL)	* V]]:->? 	 OOB	  V]]:-NO 	*OOVRL)	* 	*	* 	*	  	 	* 	*s#   B$B0B<$B-0B9<CN)r   r   r   r>  rD  rL  rT  r[  r  r`  rb  rg  r   r   r   r@  r@  
  s-    ,
%&0
&+4%&%*r   r@  c                   :    e Zd ZdZd Zd Zd Zd Zd Zd Z	d Z
y	)
TestRankFromDictz!Test Rank.from_dict() conversion.c                 x    ddl m}m} |j                  ddi      }t	        ||      sJ |j
                  dk(  sJ y)zTest Val conversion.r   )r!  r  r  rs  N)r  r!  r  rC  r   r)  )ro   r!  r  r#  s       r   test_val_conversionz$TestRankFromDict.test_val_conversioni  s9    D~~vsm,$$$$zzS   r   c                    ddl }ddlm}m} |j	                  ddddgii      }t        ||      sJ t        |j                  |j                        r1 |j                  |j                   |j                  ddg            sJ |j                  ddgk(  sJ |j                  dk(  sJ |j                  d	k(  sJ |j	                  dddgd
dddi      }|j                  d
k(  sJ |j                  dk(  sJ |j                  sJ y)zTest KNN conversion.r   N)r!  r  r  rX   r   r  r  r  r     Tr  )numpyr  r!  r  rC  r   rX   ndarrayallcloserm   r`   r  r  )ro   rl   r!  r  r#  s        r   test_knn_conversionz$TestRankFromDict.test_knn_conversionq  s   D ~~v#s'<=>$$$$djj"**-2;;tzz8288S#J+?@@@::#s+++xx<'''zzR ~~!3Z- #'		
 xx----zzS   r   c                 \   ddl m}m}m}m}m} |j                  dddiddigi      }t        ||      sJ |j                  dddiddid	i      }t        ||      sJ |j                  d
ddiddigi      }t        ||      sJ |j                  dddiddid	i      }t        ||      sJ y)z%Test arithmetic operator conversions.r   )r!  SumSubMulDivr  r  rs  r  z$subr  r  r  rw  r  N)r  r!  rs  rt  ru  rv  rC  r   )ro   r!  rs  rt  ru  rv  r#  s          r   test_arithmetic_operatorsz*TestRankFromDict.test_arithmetic_operators  s    SS ~~v}'EFG$$$$ ~~vQT'VWX$$$$ ~~v}'EFG$$$$ ~~vQT'VWX$$$$r   c                     ddl m}m}m}m} |j                  dddii      }t        ||      sJ |j                  dddii      }t        ||      sJ |j                  ddd	ii      }t        ||      sJ y
)zTest math function conversions.r   )r!  AbsExpLogz$absr  g      z$expr  z$logrw  N)r  r!  ry  rz  r{  rC  r   )ro   r!  ry  rz  r{  r#  s         r   test_math_functionsz$TestRankFromDict.test_math_functions  s    NN ~~v~67$$$$ ~~v}56$$$$ ~~v}56$$$$r   c                     ddl m}m}m} |j	                  dddiddigi      }t        ||      sJ |j	                  dddiddigi      }t        ||      sJ y)	zTest min/max conversions.r   )r!  MaxMinz$maxr  rs  r  z$minN)r  r!  r~  r  rC  r   )ro   r!  r~  r  r#  s        r   test_aggregation_functionsz+TestRankFromDict.test_aggregation_functions  sl    II ~~v}'EFG$$$$ ~~v}'EFG$$$$r   c                     ddl m}m} |j                  ddddddgiid	d
igidd	did	digigi      }t	        ||      sJ y)z%Test complex nested rank expressions.r   )r!  rs  r  r  r  rX   r   r  r  r  rs  N)r  r!  rs  rC  r   )ro   r!  rs  r#  s       r   test_complex_rank_expressionz-TestRankFromDict.test_complex_rank_expression  sf    D~~v#s'<=}MNvsmfc];<
 $$$$r   c                    ddl }ddlm}  |j                  t        d      5  |j                  d       ddd        |j                  t        d      5  |j                  i        ddd        |j                  t        d      5  |j                  d	d
dgid       ddd        |j                  t        d      5  |j                  ddi       ddd       y# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   VxY w# 1 sw Y   yxY w)zTest invalid Rank dict inputs.r   Nr   rd  r   re  rf  zexactly one operatorrs  rX   r   )r  r  zrequires a numberr  znot a number)r   r  r!  r   r  rC  rc  )ro   r   r!  s      r   test_invalid_rank_dictsz(TestRankFromDict.test_invalid_rank_dicts  s    ?V]]9O< 	)NN<(	) V]]:->? 	NN2	 V]]:-CD 	DNNC'C51ABC	D V]]9,?@ 	5NNFN34	5 	5	) 	)	 		D 	D	5 	5s/   CC(C4?D C%(C14C= D	N)r   r   r   r>  rk  rq  rw  r|  r  r  r  r   r   r   ri  ri  f  s)    +! @%(% 
%%5r   ri  c                   4    e Zd ZdZd Zd Zd Zd Zd Zd Z	y)	TestLimitFromDictz"Test Limit.from_dict() conversion.c                 z    ddl m} |j                  ddi      }|j                  dk(  sJ |j                  dk(  sJ y)zTest limit without offset.r   Limitr  r  Nr  r  rC  r  r(  ro   r  r  s      r   test_limit_onlyz!TestLimitFromDict.test_limit_only  s;    @".{{b   ||q   r   c                 t    ddl m} |j                  ddi      }|j                  dk(  sJ |j                  J y)zTest offset without limit.r   r  r(  r   N)r  r  rC  r(  r  r  s      r   test_offset_onlyz"TestLimitFromDict.test_offset_only  s9    @2/||r!!!{{"""r   c                 |    ddl m} |j                  ddd      }|j                  dk(  sJ |j                  dk(  sJ y)zTest both limit and offset.r   r  r  r   r'  Nr  r  s      r   test_limit_and_offsetz'TestLimitFromDict.test_limit_and_offset  s<    @" ;<{{b   ||r!!!r   c                    ddl }ddlm}  |j                  t        d      5  |j                  ddi       ddd        |j                  t        d      5  |j                  ddi       ddd        |j                  t        d      5  |j                  d	di       ddd       y# 1 sw Y   rxY w# 1 sw Y   JxY w# 1 sw Y   yxY w)
zTest Limit validation.r   Nr  zmust be positiver   r  r   zmust be non-negativer(  r   r  r  r   rc  rC  ro   r   r  s      r   test_validationz!TestLimitFromDict.test_validation  s    @ V]]:-?@ 	+OOWbM*	+ V]]:-?@ 	*OOWaL)	* V]]:-CD 	,OOXrN+	, 	,	+ 	+	* 	*	, 	,s#   B(B4C (B14B= C	c                    ddl }ddlm}  |j                  t        d      5  |j                  d       ddd        |j                  t        d      5  |j                  dd	i       ddd        |j                  t        d      5  |j                  d
di       ddd       y# 1 sw Y   rxY w# 1 sw Y   JxY w# 1 sw Y   yxY w)zTest type validation.r   Nr  rd  r   re  zmust be an integerr  20r(  g      %@)r   r  r  r   r  rC  r  s      r   test_invalid_typesz$TestLimitFromDict.test_invalid_types  s    @V]]9O< 	*OOL)	* V]]9,@A 	-OOWdO,	- V]]9,@A 	.OOXt,-	. 	.	* 	*	- 	-	. 	.s#   B&B2	B>&B/2B;>Cc                     ddl }ddlm}  |j                  t        d      5  |j                  ddd       ddd       y# 1 sw Y   yxY w)	"Test rejection of unexpected keys.r   Nr  Unexpected keysr   r   r`   )r  r  r  r  s      r   test_unexpected_keysz&TestLimitFromDict.test_unexpected_keys  s?    @V]]:->? 	=OObU;<	= 	= 	=   AA
N)
r   r   r   r>  r  r  r  r  r  r  r   r   r   r  r    s#    ,!#",".=r   r  c                   4    e Zd ZdZd Zd Zd Zd Zd Zd Z	y)	TestSelectFromDictz#Test Select.from_dict() conversion.c                    ddl m}m} |j                  dg di      }|j                  |j
                  v sJ |j                  |j
                  v sJ |j                  |j
                  v sJ |j                  |j
                  v sJ y)zTest special key conversion.r   Selectr5  rY   )r.  r  r1  r/  N)	r  r  r5  rC  DOCUMENTrY   	EMBEDDINGMETADATASCOREro   r  r5  r  s       r   rb  z$TestSelectFromDict.test_special_keys'  sx    F!!GH
 ||v{{***}}+++||v{{***yyFKK'''r   c                     ddl m}m} |j                  dg di      } |d      |j                  v sJ  |d      |j                  v sJ  |d      |j                  v sJ y)	z!Test regular metadata field keys.r   r  rY   )titleauthordater  r  r  N)r  r  r5  rC  rY   r  s       r   test_metadata_keysz%TestSelectFromDict.test_metadata_keys3  s]    F!!6+F"GH7|v{{***8}+++6{fkk)))r   c                     ddl m}m} |j                  dg di      }|j                  |j
                  v sJ  |d      |j
                  v sJ |j                  |j
                  v sJ y)z&Test mix of special and metadata keys.r   r  rY   )r.  r  r/  r  N)r  r  r5  rC  r  rY   r  r  s       r   test_mixed_keysz"TestSelectFromDict.test_mixed_keys<  s]    F!!6+K"LM||v{{***7|v{{***yyFKK'''r   c                 j    ddl m} |j                  dg i      }t        |j                        dk(  sJ y)zTest empty keys list.r   r  rY   N)r  r  rC  r=   rY   )ro   r  r  s      r   test_empty_keysz"TestSelectFromDict.test_empty_keysE  s0    A!!62,/6;;1$$$r   c                    ddl }ddlm}  |j                  t        d      5  |j                  d       ddd        |j                  t        d      5  |j                  dd	i       ddd        |j                  t        d
      5  |j                  ddgi       ddd       y# 1 sw Y   sxY w# 1 sw Y   KxY w# 1 sw Y   yxY w)zTest Select validation.r   Nr  rd  r   re  zmust be a list/tuple/setrY   z
not a listzmust be a stringr<  )r   r  r  r   r  rC  ro   r   r  s      r   r  z"TestSelectFromDict.test_validationL  s    AV]]9O< 	+\*	+ V]]9,FG 	5fl34	5 V]]9,>? 	.fse_-	. 	.	+ 	+	5 	5	. 	.s#   B'B3	B?'B03B<?Cc                     ddl }ddlm}  |j                  t        d      5  |j                  g dd       ddd       y# 1 sw Y   yxY w)r  r   Nr  r  r   r`   )rY   r  )r   r  r  r   rc  rC  r  s      r   r  z'TestSelectFromDict.test_unexpected_keysZ  sA    AV]]:->? 	=bU;<	= 	= 	=r  N)
r   r   r   r>  rb  r  r  r  r  r  r   r   r   r  r  $  s#    -
(*(%.=r   r  c                   .    e Zd ZdZd Zd Zd Zd Zd Zy)TestRoundTripConversionz9Test that to_dict() and from_dict() round-trip correctly.c                     ddl m}m}m}  | |d      dk(   |d      dkD  g      }|j	                         }|j                  |      }|j	                         |k(  sJ y)z!Test Where round-trip conversion.r   )r  r]  r5  r  r  r  rs  N)r  r  r]  r5  r  rC  )ro   r  r]  r5  original	dict_formrestoreds          r   test_where_round_tripz-TestRoundTripConversion.test_where_round_tripf  s\    JJH13w<#3EFG$$&	??9-!Y...r   c                    	 ddl 	ddlm}m}m}  |ddg      dz   |d      dz  z   }|j                         }|j                  |      }|j                         }	fd	 ||      sJ y)
z Test Rank round-trip conversion.r   N)r!  r  r  r   r  rX   r  rs  c                    t        | t              rt        |t              rd| v rd|v r| d   |d   }}d|v rd|v r} j                  |d   j                        } j                  |d   j                        } j                  ||      sy|D ]!  }|dk7  s	||   |j                  |      k7  s! y yt        | j                               t        |j                               k7  ry| D ]  } | |   ||         r y yt        | t              rFt        |t              r6t        |       t        |      k7  ryt        fdt        | |      D              S | |k(  S )Nr  rX   dtypeFTc              3   6   K   | ]  \  }} ||        y wr   r   )r=  r8   r9   compare_dictss      r   r@  zVTestRoundTripConversion.test_rank_round_trip.<locals>.compare_dicts.<locals>.<genexpr>  s     G41a=A.Gs   )r   r   rm   r  rp  r   rZ   rY   r  r=   r>   r?   )	d1d2knn1knn2q1q2r`   r  rl   s	          r   r  zCTestRoundTripConversion.test_rank_round_trip.<locals>.compare_dicts{  sI   "d#
2t(<R<FbL!#FRZ$D$7d?%RXXd7m2::F%RXXd7m2::F*r{{2r2#(#' -C"g~$s)txx}2L',-  $ rwwy>S^3  %C(C"S':$% B%*R*>r7c"g% G3r2;GGGRxr   )rn  r  r!  r  r  r  rC  )
ro   r!  r  r  r  r  r  restored_dictr  rl   s
           @@r   test_rank_round_tripz,TestRoundTripConversion.test_rank_round_tripo  sm    IIc3Z(3.SC?$$&	>>), ((*	 > ]I666r   c                     ddl m}  |dd      }|j                         }|j                  |      }|j                         |k(  sJ y)z!Test Limit round-trip conversion.r   r  r  r   r'  N)r  r  r  rC  )ro   r  r  r  r  s        r   test_limit_round_tripz-TestRoundTripConversion.test_limit_round_trip  sD    @r"-$$&	??9-!Y...r   c                     ddl m}m}  ||j                   |d      |j                  h      }|j                         }|j                  |      }t        |j                         d         t        |d         k(  sJ y)z"Test Select round-trip conversion.r   r  r  rY   rY   N)r  r  r5  r  r  r  rC  rZ   )ro   r  r5  r  r  r  s         r   test_select_round_tripz.TestRoundTripConversion.test_select_round_trip  sk    Fc'lCIIFG$$&	##I.8##%f-.#i6G2HHHHr   c                 P  
 ddl ddlm} ddlm}m}m}m}  | |d      dk(   |ddg	       |d
       ||j                  h            }|j                         } ||d   r|d   nd|d   r|d   nd|d   |d         }|j                         }	
fd
 
|	|      sJ y)z+Test Search round-trip through dict inputs.r   Nr  )r5  r  r  r  r  r  r   r  r  r   r)  r  r6  filterr#  r  r  c                 2   t        | t              rt        |t              rpd| v rd|v r| d   |d   }}t        |t              rt        |t              rd|v rd|v r|d   |d   }}d|v rd|v r 
j                  |d   
j                        } 
j                  |d   
j                        } 
j                  ||      sy|D ]!  }|dk7  s	||   |j                  |      k7  s! y | D ]!  }|dk7  s	| |   |j                  |      k7  s! y yt        | j                               t        |j                               k7  ry| D ]G  }t        | |   t              r%t        ||   t              r 	| |   ||         r9 y| |   ||   k7  sG y y| |k(  S )Nr#  r  rX   r  FT)r   r   rm   r  rp  r   rZ   rY   )r  r  rank1rank2r  r  r  r  r`   compare_search_dictsrl   s            r   r  zLTestRoundTripConversion.test_search_round_trip.<locals>.compare_search_dicts  s   "d#
2t(<R<FbL#%f:r&z5E!%.:eT3J!U?v).vf$D&$7d?%-RXXd7m2::%N%-RXXd7m2::%N'2r{{2r':+0+/ !5C'*g~$s)txxPS}:T/4!5 ,. !5C'*f}CBFF3K9O/4!5 (, rwwy>S^3  %C!"S'40Z35N3BsGRWE#(CBsG+$% Rxr   )
rn  r  r  r  r5  r  r  r  r  r  )ro   r  r5  r  r  r  original_searchsearch_dict
new_searchnew_dictr  rl   s             @@r   test_search_round_tripz.TestRoundTripConversion.test_search_round_trip  s    =RR h-8+C:&b/~.	
 &--/ +6x+@+h'd(3F(;V$g&x(	

 %%'"	 H $Hk:::r   N)	r   r   r   r>  r  r  r  r  r  r   r   r   r  r  c  s!    C/+7Z/IA;r   r  )gư>)r4   N)r(   r+   sysr   r1  r   r   typingr   r   rn  rl   r   r%   chromadb.server.fastapichromadb.api.fastapir   r  r   r   r	   r
   r   chromadb.configr   chromadb.errorsr   r   r   "chromadb.utils.embedding_functionsr   fixturer   r/   r3   boolr;   rK  r@   markparametrizera   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%  r'  r-  r1  r3  r8  r:  r<  r?  rJ  rN  rP  rU  rW  rY  r^  rd  re  rh  rk  rn  rs  rp  ru  rz  r}  r  r  r  r  r  bad_number_of_results_queryr  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  r  r  flakyplatform
startswithr'  r7  rK  rU  rW  r  r  r  r  r  r  r  r@  ri  r  r  r  r   r   r   <module>r     sD   	  
   (       (  % 
 H   7 7( 7 7&"$ "F F$ F ):(;<# =#0 ):(;<%# =%#P ):(;<+$ =+$\ ):(;<+ =+*>
,.< #$45#%<=#(#$ #$45#%<=#(/+#\?F'#+($=./b#,
9,!!%$/; #$455>uE	a 
CF::99  #$455>XsO,w	.BC /M9W #$455>uEuE ((((((,
b ()>?5>  -/DE 
 )*:; '')
Z #$4502GH5>uEuE	 (,,#" 		 (uERuERweTvUS	 .5(p$(T ./5>uE	a +,EF$'N-J@$ &$)N"#<4"//!,<:(("4 .0@A uE	a	!
  #$455>eL	a
 	+)+.\-%'V'' !s||'>'>w'GH I !s||'>'>w'GH? I?;6&9RT2nO@d[&||(~*EZj3Z@#H
d dN}* }*@x5 x5v@= @=F<= <=~M; M;r   