
    ei>/                    8   d dl mZmZmZmZmZmZmZmZm	Z	m
Z
mZmZmZmZmZmZ d dlmZ d dlmZmZmZ d dlmZ d dlZ G d deee            Z G d d	eee            Z G d
 d      ZddZddZddZ ddZ!ddZ"ddZ#ddZ$ddZ%ddZ&ddZ'ddZ(ddZ)y)    )SchemaSparseVectorIndexConfigSparseEmbeddingFunctionSparseVectorStringInvertedIndexConfigIntInvertedIndexConfigFloatInvertedIndexConfigBoolInvertedIndexConfigVectorIndexConfigHnswIndexConfigSpannIndexConfigFtsIndexConfigEmbeddingFunction
EmbeddingsCmekCmekProvider)Key)ListDictAny)ValidationErrorNc                       e Zd ZdZddefdZdee   dee   fdZe	defd       Z
deeef   fdZe	d	eeef   dd fd
       Zy)MockSparseEmbeddingFunctionz+Mock sparse embedding function for testing.namec                     || _         y N_name)selfr   s     g/home/obispo/Crisostomo_bridge/mision_env/lib/python3.12/site-packages/chromadb/test/api/test_schema.py__init__z$MockSparseEmbeddingFunction.__init__   s	    
    inputreturnc                 J    |D cg c]  }t        ddgddg       c}S c c}w )Nr            ?)indicesvalues)r   )r   r#   _s      r    __call__z$MockSparseEmbeddingFunction.__call__   s%    INOAaVS#J?OOOs    c                       y)Nmock_sparse r.   r"   r    r   z MockSparseEmbeddingFunction.name"   s    r"   c                     d| j                   iS )Nr   r   r   s    r    
get_configz&MockSparseEmbeddingFunction.get_config&   s    

##r"   configc                 8    t        | j                  dd            S )Nr   r-   )r   getr2   s    r    build_from_configz-MockSparseEmbeddingFunction.build_from_config)   s    *6::fm+LMMr"   N)r-   )__name__
__module____qualname____doc__strr!   r   r   r+   staticmethodr   r   r   r1   r6   r.   r"   r    r   r      s    5S Pd3i PD,> P #  $DcN $ N$sCx. N5R N Nr"   r   c                       e Zd ZdZddefdZdee   defdZe	defd       Z
deeef   fdZe	d	eeef   dd fd
       ZdefdZdee   fdZy)MockEmbeddingFunctionz$Mock embedding function for testing.
model_namec                     || _         y r   _model_name)r   r?   s     r    r!   zMockEmbeddingFunction.__init__1   s
    %r"   r#   r$   c                 n    dd l }|D cg c]!  }|j                  g d|j                        # c}S c c}w )Nr   )r'   g       @g      @)dtype)numpyarrayfloat32)r   r#   npr*   s       r    r+   zMockEmbeddingFunction.__call__4   s.     FKK

;KKKs   &2c                       y)Nmock_embeddingr.   r.   r"   r    r   zMockEmbeddingFunction.name:   s    r"   c                     d| j                   iS )Nr?   rA   r0   s    r    r1   z MockEmbeddingFunction.get_config>   s    d..//r"   r2   c                 8    t        | j                  dd            S )Nr?   
mock_model)r>   r4   r5   s    r    r6   z'MockEmbeddingFunction.build_from_configA   s    $VZZl%KLLr"   c                      y)Ncosiner.   r0   s    r    default_spacez#MockEmbeddingFunction.default_spaceE   s    r"   c                 
    g dS )N)rO   l2ipr.   r0   s    r    supported_spacesz&MockEmbeddingFunction.supported_spacesH   s    %%r"   N)rM   )r7   r8   r9   r:   r;   r!   r   r   r+   r<   r   r   r   r1   r6   rP   rT   r.   r"   r    r>   r>   .   s    .&3 &Ld3i LJ L  #    0DcN 0 M$sCx. M5L M Ms &$s) &r"   r>   c                   `   e Zd ZdZd-dZd-dZd-dZd-dZd-dZd-dZ	d-d	Z
d-d
Zd-dZd-dZd-dZd-dZd-dZd-dZd-dZd-dZd-dZd-dZd-dZd-dZd-dZd-dZd-dZd-dZd-dZd-dZd-dZd-dZd-dZ d-d Z!d-d!Z"d-d"Z#d-d#Z$d-d$Z%d-d%Z&d-d&Z'd-d'Z(d-d(Z)d-d)Z*d-d*Z+d-d+Z,d-d,Z-y).TestNewSchemaz$Test cases for the new Schema class.Nc                 ,	   t               }|j                  J |j                  j                  J |j                  j                  j                  J |j                  j                  j                  j                  du sJ |j                  j                  j
                  J |j                  j                  j
                  j                  du sJ |j                  j                  J |j                  j                  j                  J |j                  j                  j                  j                  du sJ |j                  j                  J |j                  j                  j                  J |j                  j                  j                  j                  du sJ |j                  j                  J |j                  j                  j                  J |j                  j                  j                  j                  du sJ |j                  j                  J |j                  j                  j                  J |j                  j                  j                  j                  du sJ |j                  j                  J |j                  j                  j                  J |j                  j                  j                  j                  du sJ |j                   J t#        |j                         dk(  sJ d|j                   v sJ |j                   d   j                  J |j                   d   j                  j                  J |j                   d   j                  j                  j                  du sJ |j                   d   j                  j
                  J |j                   d   j                  j
                  j                  du sJ d|j                   v sJ |j                   d   j                  J |j                   d   j                  j                  J |j                   d   j                  j                  j                  du sJ |j                   d   j                  j                  j$                  j&                  dk(  sJ y)z5Test that Schema() initializes with correct defaults.NFT   	#document
#embedding)r   defaultsstring	fts_indexenabledstring_inverted_index
float_listvector_indexsparse_vectorsparse_vector_index	int_valueint_inverted_indexfloat_valuefloat_inverted_indexbooleanbool_inverted_indexkeyslenr2   
source_keyr   schemas     r    "test_default_schema_initializationz0TestNewSchema.test_default_schema_initializationO   s    *** %%111%%//;;;%%//775@@@%%;;GGGOO""88@@DH	
H ))555))66BBBOO&&33;;uD	
D ,,888,,@@LLLOO))==EEN	
N ((444((;;GGGOO%%88@@DH	
H **666**??KKKOO''<<DDL	
L &&222&&::FFFOO##77??4G	
G {{&&&6;;1$$$ fkk))){{;'..:::{{;'..88DDD{{;'..88@@DHHH{{;'..DDPPP{{;'..DDLLPUUUU v{{***{{<(33???{{<(33@@LLL{{<(33@@HHDPPPKK%00==DDOO	
r"   c                    t               }t               }|j                  |d      }||u sJ d|j                  v sJ |j                  d   j                  J |j                  d   j                  j
                  	 J |j                  d   j                  j
                  j                  du sJ |j                  d   j                  j
                  j                  |k(  sJ |j                  d   j                  J |j                  d   j                  J |j                  d   j                  J |j                  d   j                  J |j                  d   j                  J |j                  j                  J |j                  j                  j
                  J |j                  j                  j
                  j                  du sJ y)zJTest creating a sparse vector index on a specific key with default config.custom_sparse_keyr2   keyNTF)r   r   create_indexrj   rb   rc   r^   r2   r\   r`   rd   rf   rh   r[   r   rn   r2   results       r    &test_create_sparse_vector_index_on_keyz4TestNewSchema.test_create_sparse_vector_index_on_key   s    )*$$F8K$L  #fkk111 {{./==IIIKK+,::NN	
 KK+,::NNVV	
 KK+,::NNUU	

 {{./66>>>{{./::BBB{{./99AAA{{./;;CCC{{./77??? ,,888,,@@LLLOO))==EEN	
Nr"   c                    t               }t        d      }t        |d      }|j                  |d      }||u sJ d|j                  v sJ |j                  d   j
                  J |j                  d   j
                  j                  	 J |j                  d   j
                  j                  }|j                  du sJ |j                  j                  |k(  sJ |j                  j                  dk(  sJ |j                  j                  j                         d	k(  sJ |j                  j                  j                         d
dik(  sJ |j                  j
                  J |j                  j
                  j                  J |j                  j
                  j                  j                  du sJ |j                  j
                  j                  j                  j                  	 J y)zTTest creating a sparse vector index with custom config including embedding function.custom_sparse_efr   custom_document_fieldembedding_functionrl   sparse_embeddingsrr   NTr-   r   F)r   r   r   rt   rj   rb   rc   r^   r2   r}   rl   r   r1   r[   )r   rn   embedding_funcr2   rv   sparse_indexs         r    2test_create_sparse_vector_index_with_custom_configz@TestNewSchema.test_create_sparse_vector_index_with_custom_config   s    5:LM(-:Q

 $$F8K$L  #fkk111{{./==IIIKK+,::NN	

 {{

-++ 	 ##t+++ ""55GGG""--1HHHH ""55::<MMM""55@@B&G
 
 	
 

 ,,888,,@@LLLOO))==EEN	
N OO))==DDWW	
r"   c                    t               }t               }|j                  |d      }||u sJ d|j                  v sJ |j                  d   j                  J |j                  d   j                  j
                  J |j                  d   j                  j
                  j                  du sJ d|j                  v sJ |j                  d   j                  J |j                  d   j                  j
                  J |j                  d   j                  j
                  j                  du sJ d|j                  v sJ |j                  d   j                  J |j                  j                  J |j                  j                  j
                  J |j                  j                  j
                  j                  du sJ y)7Test disabling string inverted index on a specific key.custom_text_keyrr   NFrY   rZ   T)r   r   delete_indexrj   r\   r_   r^   r[   ru   s       r    test_delete_index_on_keyz&TestNewSchema.test_delete_index_on_key   s    +,$$F8I$J  !FKK/// {{,-44@@@{{,-44JJVVVKK)*11GGOOSXX	
X fkk))){{;'..:::{{;'..DDPPPKK$++AAIIUR	
R v{{***KK%,,4	
4 %%111%%;;GGGOO""88@@DH	
Hr"   c                    t               }t        dt                     }t               }|j	                  |d      j                  |d      j                  |d      }||u sJ d|j                  v sJ d|j                  v sJ d|j                  v sJ |j                  d   j                  J |j                  d   j                  j                  J |j                  d   j                  j                  j                  du sJ |j                  d   j                  j                  j                  j                  dk(  sJ |j                  d   j                  J |j                  d   j                  J |j                  d   j                  J |j                  d   j                  J |j                  d   j                   J |j                  d   j                  J |j                  d   j                  j"                  J |j                  d   j                  j"                  j                  d	u sJ |j                  d   j                  J |j                  d   j                  J |j                  d   j                  J |j                  d   j                  J |j                  d   j                   J |j                  d   j                  J |j                  d   j                  j"                  J |j                  d   j                  j"                  j                  d	u sJ |j                  d   j                  J |j                  d   j                  J |j                  d   j                  J |j                  d   j                  J |j                  d   j                   J |j$                  j                  J |j$                  j                  j                  J |j$                  j                  j                  j                  d	u sJ |j$                  j                  J |j$                  j                  j"                  J |j$                  j                  j"                  j                  du sJ d
|j                  v sJ d|j                  v sJ |j                  d
   j                  J |j                  d
   j                  j&                  J |j                  d
   j                  j&                  j                  du sJ |j                  d   j                  J |j                  d   j                  j(                  J |j                  d   j                  j(                  j                  du sJ y)zDTest chaining create_index() and delete_index() operations together.raw_textrl   r}   embeddings_keyrr   
text_key_1
text_key_2NTFrY   rZ   )r   r   r   r   rt   r   rj   rb   rc   r^   r2   rl   r\   r`   rd   rf   rh   r_   r[   r]   ra   )r   rn   sparse_configstring_configrv   s        r    )test_chained_create_and_delete_operationsz7TestNewSchema.test_chained_create_and_delete_operations"  sj    0!6Q6S
 23 }:JK\L\A\L\A 	   6;;...v{{***v{{*** {{+,::FFFKK()77KKW	
W KK()77KKSS	
 KK m//zz	
 {{+,33;;;{{+,77???{{+,66>>>{{+,88@@@{{+,44<<< {{<(//;;;{{<(//EEQQQ{{<(//EEMMQVVVV {{<(66>>>{{<(33;;;{{<(22:::{{<(44<<<{{<(00888 {{<(//;;;{{<(//EEQQQ{{<(//EEMMQVVVV {{<(66>>>{{<(33;;;{{<(22:::{{<(44<<<{{<(00888 ,,888,,@@LLLOO))==EEN	
N %%111%%;;GGGOO""88@@DH	
H fkk)))v{{***{{;'..:::{{;'..88DDDKK$++55==E	
E{{<(33???{{<(33@@LLLKK%00==EEM	
Mr"   c                    t               }t        dd      }|j                  |      }||u sJ |j                  j                  J |j                  j                  j
                  J |j                  j                  j
                  j                  du sJ |j                  j                  j
                  j                  j                  dk(  sJ |j                  j                  j
                  j                  j                  dk(  sJ |j                  d   j                  J |j                  d   j                  j
                  J |j                  d   j                  j
                  j                  du sJ |j                  d   j                  j
                  j                  j                  dk(  sJ |j                  d   j                  j
                  j                  j                  d	k(  sJ t        d
      }t        j                  t        d      5  |j                  |d       ddd       t        j                  t        d      5  |j                  |d	       ddd       t        d      }t        j                  t        d      5  |j                  |d       ddd       y# 1 sw Y   xY w# 1 sw Y   YxY w# 1 sw Y   yxY w)z5Test vector index configuration and key restrictions.rO   custom_source)spacerl   r5   NFrZ   TrY   rR   )r   /Vector index cannot be enabled on specific keysmatch
my_vectorsrr   .Cannot create index on special key '#document'rS   /Cannot create index on special key '#embedding')r   r   rt   r[   r`   ra   r^   r2   r   rl   rj   pytestraises
ValueError)r   rn   vector_configrv   vector_config2vector_config3s         r    )test_vector_index_config_and_restrictionsz7TestNewSchema.test_vector_index_config_and_restrictions  s   )_U $$M$: ))555))66BBBOO&&33;;uD	
D))66==CCxOOOOO&&33::EEX	
X
 {{<(33???{{<(33@@LLLKK%00==EEM	
M KK%00==DDJJhV	
V KK%00==DDOO	

 +6]]O
 	I ~<H	I ]]N
 	H ~;G	H +6]]O
 	I ~<H	I 	I	I 	I	H 	H	I 	Is$   J1J=K	1J:=K	Kc                 F   t               }t        d      }t        ddd      }t        |d|d	      }|j	                  |
      }||u sJ |j
                  j                  J |j
                  j                  j                  }|J |j                  du sJ |j                  j                  |u sJ |j                  j                  j                         dk(  sJ |j                  j                  j                         ddik(  sJ |j                  j                  dk(  sJ |j                  j                  J |j                  j                  j                  dk(  sJ |j                  j                  j                   dk(  sJ |j                  j                  j"                  dk(  sJ |j                  j$                  dk(  sJ |j&                  d   j                  J |j&                  d   j                  j                  }|J |j                  du sJ |j                  j                  |u sJ |j                  j                  dk(  sJ |j                  j                  J |j                  j                  j                  dk(  sJ |j                  j$                  dk(  sJ y)zATest setting embedding function and HNSW config for vector index.custom_model_v2r?          d   ef_constructionmax_neighbors	ef_searchrR   r{   )r}   r   hnswrl   r5   NFrJ   r?   rZ   TrY   )r   r>   r   r   rt   r[   r`   ra   r^   r2   r}   r   r1   r   r   r   r   r   rl   rj   )r   rn   mock_efhnsw_configr   rv   defaults_vectorembedding_vectors           r    2test_vector_index_with_embedding_function_and_hnswz@TestNewSchema.test_vector_index_with_embedding_function_and_hnsw  s    (3DE%rS

 *&.	
 $$M$: ))555 //44AA***&&%///%%88GCCC%%88==?CSSSS%%88CCE+J
 
 	
 
 %%++t333%%**666%%**::cAAA%%**88B>>>%%**44;;;%%004KKKK {{<(33???!;;|4??LL+++''4///&&99WDDD&&,,444&&++777&&++;;sBBB##..+=	
=r"   c                 v   t               }t               }|j                  |      }||u sJ |j                  j                  J |j                  j                  j
                  J |j                  j                  j
                  j                  du sJ |j                  j                  j
                  j                  |k(  sJ |j                  d   j                  J |j                  d   j                  j
                  J |j                  d   j                  j
                  j                  du sJ |j                  d   j                  j
                  j                  |k(  sJ t               }t        j                  t        d      5  |j                  |d	       ddd       t        j                  t        d
      5  |j                  |d	       ddd       t        j                  t        d      5  |j                  |d	       ddd       y# 1 sw Y   xxY w# 1 sw Y   MxY w# 1 sw Y   yxY w)z2Test FTS index configuration and key restrictions.r5   NFrY   T,FTS index cannot be enabled on specific keysr   custom_text_fieldrr   r   rZ   r   )r   r   rt   r[   r\   r]   r^   r2   rj   r   r   r   )r   rn   
fts_configrv   fts_config2s        r    &test_fts_index_config_and_restrictionsz4TestNewSchema.test_fts_index_config_and_restrictions  s   #%
 $$J$7 %%111%%//;;;OO"",,44=	
=%%//66*DDD {{;'..:::{{;'..88DDDKK$++55==E	
E{{;'..88??:MMM %&]]L
 	M {8KL	M ]]O
 	F {E	F ]]N
 	E {D	E 	E	M 	M	F 	F	E 	Es$   HH#:H/H #H,/H8c                 X   t               }t               }t        j                  t        d      5  |j                  |d       ddd       t        j                  t        d      5  |j                  |d       ddd       t               }t        j                  t        d      5  |j                  |d       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   xY w# 1 sw Y   YxY w# 1 sw Y   yxY w)zHTest that #embedding and #document keys are blocked for all index types.r   r   rY   rr   Nr   rZ   )r   r   r   r   r   rt   r   )r   rn   r   r   s       r    -test_special_keys_blocked_for_all_index_typesz;TestNewSchema.test_special_keys_blocked_for_all_index_types  s'    23]]N
 	G }+F	G ]]O
 	H },G	H 01]]N
 	G }+F	G ]]O
 	H },G	H 	H'	G 	G	H 	H	G 	G	H 	Hs/   C<'D(DD <DDD D)c                     t               }t        j                  t        d      5  |j	                  d       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)z8Test that enabling all indexes for a key is not allowed.z.Cannot enable all index types for key 'my_key'r   my_keyrs   Nz/Cannot disable all index types for key 'my_key')r   r   r   r   rt   r   rm   s     r    &test_cannot_enable_all_indexes_for_keyz4TestNewSchema.test_cannot_enable_all_indexes_for_key;  s     ]]N
 	. H-	. ]]O
 	. H-	. 	.	. 	.	. 	.s   A8B8BBc                 T   t               }t               }t        j                  t        d      5  |j                  |       ddd       t        j                  t        d      5  |j                  |d       ddd       t               }t        j                  t        d      5  |j                  |       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   xY w# 1 sw Y   YxY w# 1 sw Y   yxY w)	z9Test that deleting vector and FTS indexes is not allowed.0Deleting vector index is not currently supportedr   r5   Nr   rr   -Deleting FTS index is not currently supportedmy_text)r   r   r   r   r   r   r   r   rn   r   r   s       r    &test_cannot_delete_vector_or_fts_indexz4TestNewSchema.test_cannot_delete_vector_or_fts_indexK  s    *+]]P
 	6 }5	6 ]]P
 	H },G	H $%
]]M
 	3 z2	3 ]]M
 	B zyA	B 	B'	6 	6	H 	H	3 	3	B 	Bs/   C:&D'DD:DDDD'c                    t               }|j                  j                  J |j                  j                  j                  J |j                  j                  j                  j                  du sJ t               }|j                  |      }||u sJ |j                  j                  j                  J |j                  j                  j                  j                  du sJ |j                  j                  j                  j                  |k(  sJ |j                  d   j                  J |j                  d   j                  j                  J |j                  d   j                  j                  j                  du sJ y)z.Test disabling string inverted index globally.NTr5   FrY   )	r   r[   r\   r_   r^   r   r   r2   rj   r   rn   r   rv   s       r    +test_disable_string_inverted_index_globallyz9TestNewSchema.test_disable_string_inverted_index_globallyi  s[    %%111%%;;GGG%%;;CCtKKK 23$$M$: %%;;GGG%%;;CCuLLL%%;;BBmSSS {{;'..:::{{;'..DDPPPKK$++AAIIUR	
Rr"   c                    t               }t               }|j                  |d      }||u sJ d|j                  v sJ |j                  d   j                  J |j                  d   j                  j
                  J |j                  d   j                  j
                  j                  du sJ |j                  d   j                  j
                  j                  |k(  sJ |j                  d   j                  J |j                  d   j                  J |j                  d   j                  J |j                  j                  J |j                  j                  j
                  J |j                  j                  j
                  j                  du sJ |j                  d   j                  J |j                  d   j                  j
                  J |j                  d   j                  j
                  j                  du sJ |j                  d   j                  J |j                  d   j                  j                  J |j                  d   j                  j                  j                  du sJ y)r   my_text_fieldrr   NFTrY   rZ   )r   r   r   rj   r\   r_   r^   r2   r`   rb   rd   r[   ra   r   s       r    )test_disable_string_inverted_index_on_keyz7TestNewSchema.test_disable_string_inverted_index_on_key  s\    23$$M$O &++---{{?+22>>>{{?+22HHTTTKK(//EEMMQVV	
V KK(//EELL	

 {{?+66>>>{{?+99AAA{{?+55=== %%111%%;;GGG%%;;CCtKKK {{;'..:::{{;'..DDPPP{{;'..DDLLPUUUU{{<(33???{{<(33@@LLL{{<(33@@HHDPPPr"   c                    t               }|j                  j                  J |j                  j                  j                  J |j                  j                  j                  j                  du sJ t               }|j                  |      }||u sJ |j                  j                  j                  j                  du sJ |j                  j                  j                  j                  |k(  sJ t               }|j                  |d      }||u sJ d|j                  v sJ |j                  d   j                  J |j                  d   j                  j                  J |j                  d   j                  j                  j                  du sJ |j                  d   j                  j                  j                  |k(  sJ |j                  d   j                  J |j                  d   j                  J |j                  d   j                  J |j                  d   j                  J |j                  d   j                  J |j                  d   j                  J |j                  d   j                  J y)	zATest disabling int inverted index globally and on a specific key.NTr5   F	age_fieldrr   rY   rZ   )r   r[   rd   re   r^   r   r   r2   rj   r\   r`   rb   rf   rh   )r   rn   
int_configrv   int_config2s        r    test_disable_int_inverted_indexz-TestNewSchema.test_disable_int_inverted_index  si    ((444((;;GGG((;;CCtKKK ,-
$$J$7 ((;;CCuLLL((;;BBjPPP -.$$K[$I fkk))){{;'11==={{;'11DDPPP{{;'11DDLLPUUUUKK$..AAHHKW	
W {{;'..666{{;'22:::{{;'55==={{;'33;;;{{;'//777 {{;'..:::{{<(33???r"   c                    t               }|j                         }d|v sJ d|v sJ t        |d   t              sJ t        |d   t              sJ |d   }d|v sJ d|d   v sJ |d   d   d   du sJ |d   d   d   i k(  sJ d|d   v sJ |d   d   d   d	u sJ |d   d   d   i k(  sJ d
|v sJ d|d
   v sJ |d
   d   d   du sJ |d
   d   d   }d|v sJ |d   dk(  sJ d|v sJ |d   d   dk(  sJ |d   d   dk(  sJ |d   d   i k(  sJ d|v sJ d|d   v sJ |d   d   d   du sJ |d   d   d   }d|v sJ |d   ddik(  sJ d|v sJ d|d   v sJ |d   d   d   d	u sJ |d   d   d   i k(  sJ d|v sJ d|d   v sJ |d   d   d   d	u sJ |d   d   d   i k(  sJ d|v sJ d|d   v sJ |d   d   d   d	u sJ |d   d   d   i k(  sJ |d   }d|v sJ d|d   v sJ d|d   d   v sJ |d   d   d   d   d	u sJ |d   d   d   d   i k(  sJ d|d   d   v sJ |d   d   d   d   du sJ |d   d   d   d   i k(  sJ d|v sJ d
|d   v sJ d|d   d
   v sJ |d   d
   d   d   d	u sJ |d   d
   d   d   }d|v sJ |d   dk(  sJ d|v sJ |d   dk(  sJ d|v sJ |d   d   dk(  sJ |d   d   dk(  sJ |d   d   i k(  sJ t        j                  |      }|j
                  j                  J |j
                  j                  j                  J |j
                  j                  j                  j                  du sJ |j
                  j                  j                  j                  |j
                  j                  j                  j                  k(  sJ |j
                  j                  j                  J |j
                  j                  j                  j                  d	u sJ |j
                  j                  j                  j                  |j
                  j                  j                  j                  k(  sJ |j
                  j                  J |j
                  j                  j                  J |j
                  j                  j                  j                  du sJ |j
                  j                  j                  j                  |j
                  j                  j                  j                  k(  sJ |j
                  j                  j                  j                  j                  dk(  sJ |j
                  j                  j                  j                  j                  	 J |j
                  j                  j                  j                  j                  j                         dk(  sJ |j
                  j                  j                  j                  j                  j                         dk(  sJ |j
                  j                   J |j
                  j                   j"                  J |j
                  j                   j"                  j                  du sJ |j
                  j                   j"                  j                  |j
                  j                   j"                  j                  k(  sJ |j
                  j$                  J |j
                  j$                  j&                  J |j
                  j$                  j&                  j                  d	u sJ |j
                  j$                  j&                  j                  |j
                  j$                  j&                  j                  k(  sJ |j
                  j(                  J |j
                  j(                  j*                  J |j
                  j(                  j*                  j                  d	u sJ |j
                  j(                  j*                  j                  |j
                  j(                  j*                  j                  k(  sJ |j
                  j,                  J |j
                  j,                  j.                  J |j
                  j,                  j.                  j                  d	u sJ |j
                  j,                  j.                  j                  |j
                  j,                  j.                  j                  k(  sJ d|j0                  v sJ |j0                  d   j                  J |j0                  d   j                  j                  J |j0                  d   j                  j                  j                  d	u sJ |j0                  d   j                  j                  j                  |j0                  d   j                  j                  j                  k(  sJ |j0                  d   j                  j                  J |j0                  d   j                  j                  j                  du sJ |j0                  d   j                  j                  j                  |j0                  d   j                  j                  j                  k(  sJ d|j0                  v sJ |j0                  d   j                  J |j0                  d   j                  j                  J |j0                  d   j                  j                  j                  d	u sJ |j0                  d   j                  j                  j                  |j0                  d   j                  j                  j                  k(  sJ |j0                  d   j                  j                  j                  j2                  dk(  sJ |j0                  d   j                  j                  j                  j2                  dk(  sJ |j0                  d   j                  j                  j                  j                  dk(  sJ |j0                  d   j                  j                  j                  j                  	 J |j0                  d   j                  j                  j                  j                  j                         dk(  sJ |j0                  d   j                  j                  j                  j                  j                         dk(  sJ y) z;Test serialization and deserialization of a default Schema.r[   rj   r\   r]   r^   Fr2   r_   Tr`   ra   r   rR   r}   typeknownr   defaultrb   rc   unknownintre   floatrg   boolri   rY   rZ   rl   N)r   serialize_to_json
isinstancedictdeserialize_from_jsonr[   r\   r]   r^   r_   r`   ra   r2   r   r}   r   rb   rc   rd   re   rf   rg   rh   ri   rj   rl   )	r   original	json_datar[   r   sparse_vector_configrj   embedding_vector_configdeserializeds	            r    )test_serialize_deserialize_default_schemaz7TestNewSchema.test_serialize_deserialize_default_schema  s    8 ..0	 Y&&&""")J/666)F+T222 Z( 8###hx0000!+.y9UBBB!+.x8B>>>&(8*<<<<!"9:9EMMM!"9:8DJJJ x''',!7777%n5i@EIII .~>xH-'''W%---#}444126:gEEE126:iGGG128<BBB (***$(AAAA()>?	JeSSS'89NO 
 $';;;;#$89fi=PPPP    #x66634Y?4GGG34X>"DDD ("""%'):::: !78CtKKK !78BbHHH !!!$(8888 56yATIII 56x@BFFF   d"""4,,,,d;/9999K *;7	BdJJJK *;7ARGGG&${*;H*EEEEK *+BCINRWWWWK *+BCHMQSSSS t###tL1111l!3L!AAAAL!,/?	JdRRR"&|"4\"B>"R#
 1111&w/47776666&|4CCC#'>>>>&';<VDOOO&';<VD	QQQ&';<XF"LLL 33I> $$++777$$++55AAA$$++55==FFF!!((22::  ''1199:	
: $$++AAMMM$$++AAIITQQQ!!((>>FF  ''==EEF	
F
 $$//;;;$$//<<HHH$$//<<DDMMM!!,,99AA  ++88@@A	
A $$//<<CCIITQQQ !!,,99@@SS	
 !!,,99@@SSXXZ	
 ((55<<OOTTV	

 $$22>>>$$22FFRRR$$22FFNNRWWWW!!//CCKK  ..BBJJK	
K
 $$..:::$$..AAMMM$$..AAIITQQQ!!++>>FF  **==EEF	
F
 $$00<<<$$00EEQQQ$$00EEMMQUUUU!!--BBJJ  ,,AAIIJ	
J
 $$,,888$$,,@@LLL$$,,@@HHDPPP!!))==EE  ((<<DDE	
E
 l/////  -44@@@  -44>>JJJ  -44>>FF$NNNk*11;;CC}}[)00::BBC	
C   -44JJVVVk*11GGOOSXX	
X k*11GGOO}}[)00FFNNO	
O
 |00000  .99EEE  .99FFRRR  .99FFNNRVVVVl+66CCKK}}\*55BBJJK	
K
 l+66CCJJUU	
 MM,'22??FFQQ	

 l+66CCJJPPTXX	
X jff-?-?	
 jff-?-?	
 MMjff-?-?	
r"   c                    t               }t        dd      }|j                  |       |j                         }|d   }d|v sJ d|d   v sJ |d   d   }|d   d	u sJ |d
   d   dk(  sJ |d
   d   d   dk(  sJ |d   }d|v sJ |d   d   d   }|d   du sJ |d
   d   dk(  sJ |d
   d   dk(  sJ |d
   d   d   dk(  sJ t        j                  |      }|j
                  j                  J |j
                  j                  j                  J |j
                  j                  j                  j                  d	u sJ |j
                  j                  j                  j                  j                  dk(  sJ |j
                  j                  j                  j                  j                  	 J d|j                  v sJ |j                  d   j                  J |j                  d   j                  j                  J |j                  d   j                  j                  j                  du sJ |j                  d   j                  j                  j                  j                  dk(  sJ |j                  d   j                  j                  j                  j                  dk(  sJ |j                  d   j                  j                  j                  j                  	 J y)z^Test serialization/deserialization of Schema with vector config where embedding_function=None.rO   N)r   r}   r5   r[   r`   ra   r^   Fr2   r   r}   r   legacyrj   rZ   Trl   rY   )r   r   rt   r   r   r[   r`   ra   r^   r2   r   r}   rj   rl   )	r   r   r   r   r[   vector_jsonrj   embedding_vector_jsonr   s	            r    3test_serialize_deserialize_with_vector_config_no_efzATestNewSchema.test_serialize_deserialize_with_vector_config_no_ef  sU    8)#
 	]3 ..0	 Z(x''',!7777|,^<9%...8$W-9998$%9:6BhNNN  t### $\ 2< @ P$Y/4777!(+G4@	
@$X.|<KKK$X./CDVLPXXXX 33I> $$//;;;$$//<<HHH$$//<<DDMMM!!,,99@@FF(R	
R !!,,99@@SS	

 |00000  .99EEE  .99FFRRR  .99FFNNRVVVVl+66CCJJPP	
 l+66CCJJUU	

 jff-?-?	
r"   c                 ,   ddl m} t        |d<   	 t               }t        d      }t	        ddd	      }t        |d
|      }|j                  |       |j                         }|d   }d|v sJ d|d   v sJ |d   d   }|d   du sJ |d   d   d
k(  sJ |d   d   d   dk(  sJ |d   d   d   dk(  sJ |d   d   d   d   dk(  sJ d|d   v sJ |d   d   d   dk(  sJ |d   d   d   dk(  sJ |d   d   d   dk(  sJ |d   }	d|	v sJ |	d   d   d   }
|
d   du sJ |
d   d   d
k(  sJ |
d   d    d!k(  sJ |
d   d   d   dk(  sJ |
d   d   d   dk(  sJ |
d   d   d   d   dk(  sJ d|
d   v sJ |
d   d   d   dk(  sJ |
d   d   d   dk(  sJ |
d   d   d   dk(  sJ t        j                  |      }|j                  j                  J |j                  j                  j                  J |j                  j                  j                  j                  du sJ |j                  j                  j                  j                  j                  d
k(  sJ |j                  j                  j                  j                  j                  	 J |j                  j                  j                  j                  j                  j!                         dk(  sJ |j                  j                  j                  j                  j                  j#                         }|d   dk(  sJ |j                  j                  j                  j                  j$                  J |j                  j                  j                  j                  j$                  j&                  dk(  sJ |j                  j                  j                  j                  j$                  j(                  dk(  sJ |j                  j                  j                  j                  j$                  j*                  dk(  sJ d|j,                  v sJ |j,                  d   j                  J |j,                  d   j                  j                  J |j,                  d   j                  j                  j                  du sJ |j,                  d   j                  j                  j                  j                  d
k(  sJ |j,                  d   j                  j                  j                  j.                  d!k(  sJ |j,                  d   j                  j                  j                  j                  	 J |j,                  d   j                  j                  j                  j                  j!                         dk(  sJ |j,                  d   j                  j                  j                  j                  j#                         }|d   dk(  sJ |j,                  d   j                  j                  j                  j$                  	 J |j,                  d   j                  j                  j                  j$                  j&                  dk(  sJ |j,                  d   j                  j                  j                  j$                  j(                  dk(  sJ |j,                  d   j                  j                  j                  j$                  j*                  dk(  sJ 	 d|v r|d= y"y"# d|v r|d= w w xY w)#zLTest serialization/deserialization of Schema with custom embedding function.r   known_embedding_functionsrJ   custom_model_v3r      0      r   rS   r}   r   r   r5   r[   r`   ra   r^   Fr2   r   r}   r   r   r   r?   r   r   r   r   rj   rZ   Trl   rY   N)"chromadb.utils.embedding_functionsr   r>   r   r   r   rt   r   r   r[   r`   ra   r^   r2   r   r}   r   r1   r   r   r   r   rj   rl   )r   r   r   	custom_efr   r   r   r[   r   rj   r   r   	ef_configef_config_embeddings                 r    )test_serialize_deserialize_with_custom_efz7TestNewSchema.test_serialize_deserialize_with_custom_ef
  s    	Q6K!"23k	@xH-9JKI) #2K .#, M
 !!!7 !224I !,H8+++!Xl%;;;;"<0@Ky)U222x(1T999x()=>vF'QQQH%&:;FCGWWW H%&:;HElS$%% [2222x(01BCsJJJx(0ARGGGx(0=DDD V$D4'''$($6|$D^$T!%i0D88 &h/8D@@ &h/=LL &h/0DEfMQXXX &h/0DEfM#$$ &h/0DEhO  %%% 28<<<<(26:;LMQTTTT(26:?KrQQQ(26:;G3NNN "77	BL  ((33???((33@@LLL((33@@HHEQQQ%%00==DDJJdRR %%00==DDWW %%00==DDWW\\^#$$ %--88EELL__jjlI\*.????((33@@GGLLXXX%%00==DDIIYY %%00==DDIIWW %%00==DDIISS
  <#4#4444$$\2==III$$\2==JJVVV!!,/::GGOOSWWW !!,/::GGNNTT !! *\\&& !! *\\&&1C1C !! *\\&&1C1CDDF#$$ #/"3"3#jff-?-?

   '|48IIII !!,/::GGNNSS !! *\\&&oo !! *\\&&mm !! *\\&&ii  #<<-.>? =#<<-.>? =s   [.\	 	
\c                    ddl m} t        |d<   	 t               }t        d      }t	        dddd	
      }t        |d|      }|j                  |       |j                         }|d   }d|v sJ d|d   v sJ |d   d   }|d   du sJ |d   d   dk(  sJ |d   d   d   dk(  sJ |d   d   d   dk(  sJ |d   d   d   d   dk(  sJ d|d   v sJ |d   d   d   dk(  sJ |d   d   d   dk(  sJ |d   d   d   dk(  sJ |d   d   d   d	k(  sJ |d   j                  d      J |d!   }	d"|	v sJ |	d"   d   d   }
|
d   d#u sJ |
d   d   dk(  sJ |
d   d$   d%k(  sJ |
d   d   d   dk(  sJ |
d   d   d   dk(  sJ |
d   d   d   d   dk(  sJ d|
d   v sJ |
d   d   d   dk(  sJ |
d   d   d   dk(  sJ |
d   d   d   dk(  sJ |
d   d   d   d	k(  sJ |
d   j                  d      J t        j                  |      }|j                  j                  J |j                  j                  j                  J |j                  j                  j                  j                  du sJ |j                  j                  j                  j                  j                  dk(  sJ |j                  j                  j                  j                  j                   	 J |j                  j                  j                  j                  j                   j#                         dk(  sJ |j                  j                  j                  j                  j                   j%                         }|d   dk(  sJ |j                  j                  j                  j                  j&                  J |j                  j                  j                  j                  j&                  j(                  dk(  sJ |j                  j                  j                  j                  j&                  j*                  dk(  sJ |j                  j                  j                  j                  j&                  j,                  dk(  sJ |j                  j                  j                  j                  j&                  j.                  d	k(  sJ |j                  j                  j                  j                  j0                  J d"|j2                  v sJ |j2                  d"   j                  J |j2                  d"   j                  j                  J |j2                  d"   j                  j                  j                  d#u sJ |j2                  d"   j                  j                  j                  j                  dk(  sJ |j2                  d"   j                  j                  j                  j4                  d%k(  sJ |j2                  d"   j                  j                  j                  j                   	 J |j2                  d"   j                  j                  j                  j                   j#                         dk(  sJ |j2                  d"   j                  j                  j                  j                   j%                         }|d   dk(  sJ |j2                  d"   j                  j                  j                  j&                  	 J |j2                  d"   j                  j                  j                  j&                  j(                  dk(  sJ |j2                  d"   j                  j                  j                  j&                  j*                  dk(  sJ |j2                  d"   j                  j                  j                  j&                  j,                  dk(  sJ |j2                  d"   j                  j                  j                  j&                  j.                  d	k(  sJ |j2                  d"   j                  j                  j                  j0                  	 J 	 d|v r|d= y y # d|v r|d= w w xY w)&zETest serialization/deserialization of Schema with SPANN index config.r   r   rJ   spann_modelr   r   2   r      )search_nprobewrite_nprober   r   rO   )r}   r   spannr5   r[   r`   ra   r^   Fr2   r   r}   r   r   r   r?   r   r   r   r   r   r   Nrj   rZ   Trl   rY   )r   r   r>   r   r   r   rt   r   r4   r   r[   r`   ra   r^   r2   r   r}   r   r1   r   r   r   r   r   r   rj   rl   )r   r   r   r   spann_configr   r   r[   r   rj   r   r   r   r   s                 r    ,test_serialize_deserialize_with_spann_configz:TestNewSchema.test_serialize_deserialize_with_spann_config  sd	    	Q6K!"23B	@xH-GI+!CSVL .#,HLM !!!7 !224I !,H8+++!Xl%;;;;"<0@Ky)U222x(1X===x()=>vF'QQQH%&:;FCGWWW H%&:;HElS !! k(3333x(1/BcIIIx(1.ARGGGx(12CDKKKx(1+>#EEEx(,,V4<<< V$D4'''$($6|$D^$T!%i0D88 &h/8HDD &h/=LL &h/0DEfMQXXX &h/0DEfM#$$ &h/0DEhO  !!! 3H====(27;OLPSSSS(27;NKrQQQ(27;<MNRUUUU(27;KHCOOO(266v>FFF "77	BL  ((33???((33@@LLL((33@@HHEQQQ%%00==DDJJhVV %%00==DDWW %%00==DDWW\\^#$$ %--88EELL__jjlI\*m;;; %%00==DDJJVV %%00==DDJJXX %%00==DDJJWW %%00==DDJJZZ %%00==DDJJTT  ((33@@GGLLTTT  <#4#4444$$\2==III$$\2==JJVVV!!,/::GGOOSWWW !!,/::GGNNTT !! *\\&& !! *\\&&1C1C !! *\\&&1C1CDDF#$$ #/"3"3#jff-?-?

   '|4EEE !!,/::GGNNTT !! *\\&&}} !! *\\&&|| !! *\\&& !! *\\&&yy !!,/::GGNNSS
  #<<-.>? =#<<-.>? =s   `4a 
ac                 x   ddl m} t        |d<   t        |d<   	 t	               }t        d      }t        dd	      }t        |d
|      }|j                  |       t        d      }t        |d      }|j                  |d       t               }|j                  |d       t               }	|j                  |	d       t               }
|j                  |
d       |j                         }|d   }|d   }|d   d   d   d   d
k(  sJ |d   d   d   d   d   dk(  sJ |d   d   d   d   d   dk(  sJ d|v sJ d|v sJ d|v sJ d|v sJ d|v sJ d |v sJ |d   }d!|v sJ |d!   d"   d#   d$u sJ |d!   d"   d   d%   dk(  sJ |d!   d"   d   d&   d'   d(k(  sJ |d!   d"   d   d&   d)   dk(  sJ |d!   d"   d   d&   d   d)   dk(  sJ d*|vsJ d|vsJ d+|vsJ d,|vsJ d-|vsJ |d   }d*|v sJ |d*   d.   d#   d/u sJ |d*   d.   d   i k(  sJ d0|d*   vsJ d!|vsJ d|vsJ d+|vsJ d,|vsJ d-|vsJ |d   }d+|v sJ |d+   d1   d#   d/u sJ |d+   d1   d   i k(  sJ d*|vsJ d!|vsJ d|vsJ d,|vsJ d-|vsJ |d   }d,|v sJ |d,   d2   d#   d/u sJ |d,   d2   d   i k(  sJ d*|vsJ d!|vsJ d|vsJ d+|vsJ d-|vsJ |d    }d|v sJ |d   d   d#   d$u sJ |d   d   d   d   d
k(  sJ |d   d   d   d%   dk(  sJ |d   d   d   d&   d'   d(k(  sJ |d   d   d   d&   d)   dk(  sJ |d   d   d   d&   d   d3   dk(  sJ |d   d   d   d   d   dk(  sJ |d   d   d   d   d   dk(  sJ |d   d   d   j                  d4      	 J d*|vsJ d!|vsJ d+|vsJ d,|vsJ d-|vsJ |d   }d*|v sJ |d*   d0   d#   d$u sJ |d*   d0   d   i k(  sJ |d*   d.   d#   d/u sJ |d*   d.   d   i k(  sJ d!|vsJ d|vsJ d+|vsJ d,|vsJ d-|vsJ t	        j                  |      }|j                   j"                  J |j                   j"                  j$                  J |j                   j"                  j$                  j&                  j(                  d
k(  sJ |j                   j"                  j$                  j&                  j*                  J |j                   j"                  j$                  j&                  j*                  j,                  dk(  sJ |j                   j"                  j$                  j&                  j*                  j.                  dk(  sJ |j                   j"                  j$                  j&                  j0                  	 J |j                   j"                  j$                  j&                  j0                  j3                         dk(  sJ d|j4                  v sJ |j4                  d   j6                  J |j4                  d   j6                  j8                  	 J |j4                  d   j6                  j8                  j:                  d$u sJ |j4                  d   j6                  j8                  j&                  j<                  dk(  sJ |j4                  d   j>                  J |j4                  d   j"                  J |j4                  d   j@                  J d|j4                  v sJ |j4                  d   j>                  J |j4                  d   j>                  jB                  J |j4                  d   j>                  jB                  j:                  d/u sJ |j4                  d   j6                  J |j4                  d   j"                  J d|j4                  v sJ |j4                  d   j@                  J |j4                  d   j@                  jD                  J |j4                  d   j@                  jD                  j:                  d/u sJ |j4                  d   j>                  J |j4                  d   j"                  J d|j4                  v sJ |j4                  d   jF                  J |j4                  d   jF                  jH                  J |j4                  d   jF                  jH                  j:                  d/u sJ |j4                  d   j>                  J |j4                  d   j6                  J d |j4                  v sJ |j4                  d    j"                  J |j4                  d    j"                  j$                  J |j4                  d    j"                  j$                  j&                  j(                  d
k(  sJ |j4                  d    j"                  j$                  j&                  j<                  dk(  sJ |j4                  d    j"                  j$                  j&                  j*                  	 J |j4                  d    j"                  j$                  j&                  j*                  j,                  dk(  sJ |j                   j>                  J |j                   j>                  jB                  J |j                   j>                  jB                  j:                  d$u sJ |j                   j@                  J |j                   j@                  jD                  J |j                   j@                  jD                  j:                  d$u sJ |j                   j6                  J |j                   j6                  j8                  J |j                   j6                  j8                  j:                  d/u sJ 	 d|v r|d= d|v r|d= y5y5# d|v r|d= d|v r|d= w w xY w)6zLTest serialization/deserialization with multiple mixed schema modifications.r   r   rJ   r-   mixed_test_modelr   i,  @   r   r   rS   r   r5   sparse_modelrz   
text_fieldr|   embeddings_fieldrr   tagscountpricer[   rj   r`   ra   r2   r   r   r   r   rY   rZ   rb   rc   r^   Trl   r}   r   r   r   r\   r   r   r   r_   Fr]   re   rg   r?   r   N)%r   r   r>   r   r   r   r   rt   r   r   r   r   r	   r   r4   r   r[   r`   ra   r2   r   r   r   r   r}   r   rj   rb   rc   r^   rl   r\   rd   r_   re   rf   rg   )r   r   r   r   r   r   	sparse_efr   r   r   float_configr   r[   rj   embeddings_field_json	tags_json
count_json
price_jsonembedding_jsondocument_jsonr   s                        r    6test_serialize_deserialize_complex_mixed_modificationszDTestNewSchema.test_serialize_deserialize_complex_mixed_modifications  s    	Q6K!"233N!-0O	=xH .9KLI)#RPK-#,D{M !!!7 4HI3#,M !!<N!O 67M!!F!C 01J!!!A 45L!!7!C !224I !,HV$D L).9(CGLPTTTT&~6x@H%  &~6x@H#  &---T>!>d?"?d?"?$&&&4''' %));$<!"&;;;;%o67LMiX &o67LMhW      &o67LMhW(  &o67LMhW( !!! &o67LMhW("$ """ #8888'<<<< 5555"7777!6666 VIy(((X&'>?	JeSSSX&'>?IROOOi&9999")333y000	))))+++*** gJJ&&&e$%9:9ENNNe$%9:8DJJJ:---"*444z111*,,,+++ gJj(((g&'=>yIURRRg&'=>xHBNNN:---"*444z111
***+++ ",/N>111!,/?	JdRRR|,^<XFwOSWWW |,^<XF|T |,^<XF(  |,^<XF( $$$ |,^<XF((* &&& |,^<XFvN%  |,^<XFvN#  |,^<XFJJ7S >111".888....000/// !-M},,, *;7	BdJJJ *;7ARGGG *+BCINRWWWW *+BCHMQSSSS"-777}444----///... "77	BL  ((33???((33@@LLL((33@@GGMMQUUUU((33@@GGLLXXX%%00==DDIIYY %%00==DDIIWW %%00==DDWW %%00==DDWW\\^#$$
 &):)::::$$%78FFRRR!!"45CCWW !!&- 3 3GG !!&- 3 3FF::    $$%78??GGG$$%78CCKKK$$%78BBJJJ \.....$$V,33???$$V,33IIUUU!!&)00FFNNRWWW  $$V,::BBB$$V,77??? l/////$$W-77CCC$$W-77JJVVV!!'*44GGOOSXXX  $$W-44<<<$$W-88@@@ l/////$$W-99EEE!!'*66KKWW !!'*66KKSS  $$W-44<<<$$W-;;CCC  <#4#4444$$\2==III$$\2==JJVVV!!,/::GGNNTT !! *\\&& !!,/::GGNNSS !! *\\&&oo  ((//;;;((//EEQQQ%%,,BBJJdRR((22>>>((22EEQQQ%%//BBJJdRR((66BBB((66JJVVV%%33GGOOSXXX  #<<-.>? 99-m< :  #<<-.>? 99-m< :s   m=n( (n9c                    t               }t        dt                     }|j                  |d       t	               }|j                  |d       d|j
                  v sJ |j
                  d   }|j                  J |j                  j                  J |j                  j                  j                  du sJ |j                  J |j                  j                  J |j                  j                  j                  du sJ |j
                  d   j                  J |j
                  d   j                  J |j
                  d   j                  J |j
                  d   j                  J |j                         }|d   d   }d|v sJ d	|v sJ |d   d
   d   du sJ |d	   d   d   du sJ t        j                   |      }d|j
                  v sJ |j
                  d   }|j                  J |j                  j                  J |j                  j                  j                  du sJ |j                  J |j                  j                  J |j                  j                  j                  du sJ y)z;Test that multiple index types can coexist on the same key.sourcer   multi_fieldrr   NTrj   rb   r\   rc   r^   r_   )r   r   r   rt   r   rj   rb   rc   r^   r\   r_   r`   rd   rf   rh   r   r   )	r   rn   r   r   r  r   multi_field_jsonr   des_multi_fields	            r    %test_multiple_index_types_on_same_keyz3TestNewSchema.test_multiple_index_types_on_same_key  s    04O4Q
 	=mD 23=mD +++kk-0((444((<<HHH((<<DDLLL!!---!!77CCC!!77??4GGG {{=)44<<<{{=)33;;;{{=)55==={{=)11999 ,,.	$V,];"2222++++_-.CDYOSWW	
W)*AB9MQUUUU 33I> 1 1111&++M:,,888,,@@LLL,,@@HHDPPP%%111%%;;GGG%%;;CCtKKKr"   c                    t               }t               }|j                  |d       d|j                  v sJ |j                  d   }|j                  J |j                  j
                  J |j                  j
                  j                  du sJ |j                  |d       d|j                  v sJ |j                  d   }|j                  J |j                  j
                  J |j                  j
                  j                  du sJ |j                         }d|d   v sJ |d   d   }d|v sJ |d   d   d	   du sJ t        j                  |      }d|j                  v sJ |j                  d   }|j                  J |j                  j
                  J |j                  j
                  j                  du sJ y)
zcTest that disabling an index reverts to default behavior (key may still exist with disabled state).
temp_fieldrr   NTFrj   r\   r_   r^   )
r   r   rt   rj   r\   r_   r^   r   r   r   )	r   rn   r   temp_field_initialr  r   temp_field_jsonr   des_temp_fields	            r    $test_override_then_revert_to_defaultz2TestNewSchema.test_override_then_revert_to_default  s    23=lC v{{***#[[6!((444!((>>JJJ!((>>FF$NNN 	=lC v{{***[[.
  ,,,  66BBB  66>>%GGG ,,.	y0000#F+L9?***x()@A)LPUUUU 33I>|00000%**<8$$000$$::FFF$$::BBeKKKr"   c                    t               }t               }t        j                  t        d      5  |j                  |d       ddd       t               }t        j                  t        d      5  |j                  |d       ddd       t        j                  t        d      5  |j                          ddd       t        j                  t        d	      5  |j                  d
       ddd       t        j                  t        d      5  |j                  d
       ddd       t        j                  t        d      5  |j                  |       ddd       t        j                  t        d      5  |j                  |       ddd       t        j                  t        d      5  |j                  |d       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   xY w# 1 sw Y   kxY w# 1 sw Y   Bx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   yxY w)z6Test that invalid operations raise appropriate errors.r   r   rZ   rr   Nr   rY   z&Cannot enable all index types globallyz-Cannot enable all index types for key 'mykey'mykeyr   z.Cannot disable all index types for key 'mykey'r   r5   r   r   custom_fieldr   )r   r   r   r   r   rt   r   r   r   s       r    &test_error_handling_invalid_operationsz4TestNewSchema.test_error_handling_invalid_operations>  sJ    *+]]O
 	H },G	H $%
]]N
 	D z{C	D ]]:-UV 	"!	" ]]M
 	- G,	- ]]N
 	- G,	- ]]P
 	6 }5	6 ]]M
 	3 z2	3 ]]O
 	J }.I	J ]]L
 	G z~F	G 	G_	H 	H	D 	D	" 	"	- 	-	- 	-	6 	6	3 	3	J 	J	G 	Gsk   H1H(H"H/H<I	>I4I!+I-HH"H,/H9<I	II!I*-I6c                 (   t               }|j                         }t        |d         dk(  sJ d|d   v sJ d|d   v sJ t        j                  |      }|j                  }|j
                  J |j
                  j                  J |j
                  j                  j                  du sJ |j
                  j                  J |j
                  j                  j                  du sJ |j                  J |j                  j                  J |j                  j                  j                  du sJ |j                  J |j                  j                  J |j                  j                  j                  du sJ |j                  J |j                  j                  J |j                  j                  j                  du sJ |j                  J |j                  j                   J |j                  j                   j                  du sJ |j"                  J |j"                  j$                  J |j"                  j$                  j                  du sJ t        |j&                        dk(  sJ d|j&                  v sJ d|j&                  v sJ y)z;Test serialization/deserialization of an unmodified schema.rj   rX   rY   rZ   NTF)r   r   rk   r   r[   r\   r_   r^   r]   r`   ra   rb   rc   rd   re   rf   rg   rh   ri   rj   )r   r   r   r   r[   s        r    test_empty_schema_serializationz-TestNewSchema.test_empty_schema_serializationx  s    8 ..0	 9V$%***i////y0000 33I>  ((***44@@@44<<DDD((444((00E999""...""//;;;""//775@@@%%111%%99EEE%%99AAUJJJ!!---!!44@@@!!44<<DDD##///##88DDD##88@@DHHH+++33???33;;tCCC <$$%***l/////|00000r"   c                 "	   ddl m} t        |d<   	 t               }t        d      }t	        dd      }t        |d	|
      }|j                  |       |j                  t        dt                     d       |j                  t               d       |j                         }t        j                  |      }|j                         }t        j                  |      }	|	j                         }
t        j                  |
      }||	|fD ]  }|j                  j                  J |j                  j                  j                  J |j                  j                  j                  j                   j"                  d	k(  sJ |j                  j                  j                  j                   j$                  J |j                  j                  j                  j                   j$                  j&                  dk(  sJ |j                  j                  j                  j                   j$                  j(                  dk(  sJ |j                  j                  j                  j                   j*                  	 J |j                  j                  j                  j                   j*                  j-                         dk(  sJ d|j.                  v sJ |j.                  d   }|j0                  J |j0                  j2                  J |j0                  j2                  j4                  du sJ |j0                  j2                  j                   j6                  dk(  sJ d|j.                  v sJ |j.                  d   }|j8                  J |j8                  j:                  J |j8                  j:                  j4                  du rJ  t=        |j.                        t=        |	j.                        cxk(  rt=        |j.                        k(  sJ  J t?        |j.                  j/                               t?        |	j.                  j/                               cxk(  r't?        |j.                  j/                               k(  sJ  J 	 d|v r|d= yy# d|v r|d= w w xY w)zRTest that multiple serialization/deserialization cycles preserve schema integrity.r   r   rJ   roundtrip_modelr   r   (   r   rO   r   r5   textr   
embeddingsrr   r  NTF) r   r   r>   r   r   r   rt   r   r   r   r   r   r   r[   r`   ra   r2   r   r   r   r   r}   r   rj   rb   rc   r^   rl   r\   r_   rk   set)r   r   r   r   r   r   json1schema1json2schema2json3schema3rn   embeddings_overridetags_overrides                  r    .test_multiple_serialize_deserialize_roundtripsz<TestNewSchema.test_multiple_serialize_deserialize_roundtrips  sR    	Q6K!"23U	@xH-9JKI)#RPK-#,H;M !!!7!!.%:U:W !	 "  !!)B)D&!Q ..0E2259G --/E2259G --/E2259G #GW5 )S11===11>>JJJ11>>EEKKxWWW11>>EEJJVVVOO..;;BBGGWW OO..;;BBGGUU OO..;;BBUU   OO..;;BBUUZZ\'((
 $v{{222&,kk,&?#*88DDD*88LLXXX'55IIQQ (55IIPP[[
 ,,, &F 3$++777$++AAMMM$++AAIIURRRS)SX w||$GLL(9NS=NNNNNNGLL%%'(w||((*+,w||((*+,,,  #<<-.>? =#<<-.>? =s   N(R ;B?R 
Rc                 @   t               }t        d      D ]u  }d| }|dk(  r*|j                  t        d| t	                     |       7|dz  dk(  r|j                  t               |       [|j                  t               |       w t        |j                        d	k(  sJ d
|j                  v sJ |j                  d
   }|j                  J |j                  j                  J |j                  j                  j                  du sJ d|j                  v sJ |j                  d   }|j                  J |j                  j                  J |j                  j                  j                  du sJ d|j                  v sJ |j                  d   }|j                  J |j                  j                   J |j                  j                   j                  du sJ |j#                         }t        |d         d	k(  sJ t        j$                  |      }t        |j                        d	k(  sJ d
|j                  v sJ |j                  d
   }	|	j                  J |	j                  j                  J |	j                  j                  j                  du sJ |	j                  j                  j&                  j(                  dk(  sJ d|j                  v sJ |j                  d   }
|
j                  J |
j                  j                  J |
j                  j                  j                  du sJ d|j                  v sJ |j                  d   }|j                  J |j                  j                   J |j                  j                   j                  du sJ y)z2Test schema with many key overrides (stress test).r   field_r   source_r   rr   rX   r&   4   field_0NTfield_1Ffield_2rj   source_0field_49field_48)r   rangert   r   r   r   r   r   rk   rj   rb   rc   r^   r\   r_   rd   re   r   r   r2   rl   )r   rn   ikey_namer2  r3  r4  r   r   des_field_0des_field_49des_field_48s               r    test_many_keys_stressz#TestNewSchema.test_many_keys_stress  s    r 	SAs|HAv##2%,QC=+F+H ! $  Q!##+D+FH#U ##+A+C#R!	S& 6;;2%%% FKK'''++i($$000$$88DDD$$88@@DHHHFKK'''++i(~~)))~~33???~~33;;uDDDFKK'''++i(  ,,,  33???  33;;uDDD ,,.	9V$%+++ 33I><$$%+++ L-----"''	2((444((<<HHH((<<DDLLL%%99@@KK	
 \.....#((4""...""88DDD""88@@EIII\.....#((4%%111%%88DDD%%88@@EIIIr"   c                 N   t               }|j                  t        dt                     d      j	                  t               d      j	                  t               d      j	                  t               d      }||u sJ d|j                  v sJ |j                  d   }|j                  J |j                  j                  J |j                  j                  j                  d	u sJ d|j                  v sJ |j                  d   }|j                  J |j                  j                  J |j                  j                  j                  d
u sJ d|j                  v sJ |j                  d   }|j                  J |j                  j                  J |j                  j                  j                  d
u sJ d|j                  v sJ |j                  d   }|j                  J |j                  j                  J |j                  j                  j                  d
u sJ y)z@Test chaining multiple create_index and delete_index operations.r"  r   field1rr   field2field3field4NTF)r   rt   r   r   r   r   r   rj   rb   rc   r^   r\   r_   rd   re   )r   rn   rv   r@  rA  rB  rC  s          r    test_chained_operationsz%TestNewSchema.test_chained_operationsN  s    .%:U:W 	    \!:!<(\K\!:!<(\K\!7!9x\H 	  6;;&&&X&##///##77CCC##77??4GGG6;;&&&X&}}(((}}22>>>}}22::eCCC6;;&&&X&}}(((}}22>>>}}22::eCCC6;;&&&X&+++22>>>22::eCCCr"   c                 :   t               }|j                  j                  J |j                  j                  j                  J |j                  j                  j                  j                  du sJ |j                  j
                  J |j                  j
                  j                  J |j                  j
                  j                  j                  du sJ t               }|j                  |       |j                  j                  j                  J |j                  j                  j                  j                  du sJ t               }|j                  |       |j                  j
                  j                  J |j                  j
                  j                  j                  du sJ |j                  t               d       d|j                  v sJ |j                  d   j                  j                  j                  du sJ |j                  t               d       d|j                  v sJ |j                  d   j
                  j                  j                  du sJ |j                         }|d   d	   d
   d   du sJ |d   d   d   d   du sJ |d   d   d	   d
   d   du sJ |d   d   d   d   d   du sJ t        j                  |      }|j                  j                  j                  j                  du sJ |j                  j
                  j                  j                  du sJ |j                  d   j                  j                  j                  du sJ |j                  d   j
                  j                  j                  du sJ y)z8Test enabling/disabling float and bool inverted indexes.NTr5   Fr  rr   	is_activer[   r   rg   r^   r   ri   rj   )r   r[   rf   rg   r^   rh   ri   r	   r   r
   rt   rj   r   r   )r   rn   r  bool_configr   r   s         r    $test_float_and_bool_inverted_indexesz2TestNewSchema.test_float_and_bool_inverted_indexes{  s    **666**??KKK**??GG4OOO&&222&&::FFF&&::BBdJJJ 01<0**??KKK**??GG5PPP ./;/&&::FFF&&::BBeKKK 	#;#=7K&++%%%{{7#//DDLLPTTTT 	#:#<+Nfkk))){{;'//CCKKuTTT ,,.	j!'*+AB9MQVV	
V$V,-BCINRWWWWfg&w/0FG	R	
 fk*623HI)T	

 33I>$$00EEMMQVVVV$$,,@@HHEQQQg&22GGOOSWW	
W k*22FFNNRWW	
Wr"   c                 F   ddl m} t        |d<   	 t               }t        d      }t	        |      }|j                  |       |j                         }|d   d	   d
   }|d   d   dk(  sJ |d   d   d	   d
   }|d   d   dk(  sJ t        j                  |      }|j                  j                  J |j                  j                  j                  J |j                  j                  j                  j                  j                  dk(  sJ |j                  d   j                  J |j                  d   j                  j                  J |j                  d   j                  j                  j                  j                  dk(  sJ 	 d|v r|d= yy# d|v r|d= w w xY w)zVTest that space is correctly inferred from embedding function when not explicitly set.r   r   rJ   space_inference_testr   r}   r5   r[   r`   ra   r2   r   rO   rj   rZ   Nr   r   r>   r   r   rt   r   r   r[   r`   ra   r2   r   rj   	r   r   rn   r   r   r   r   r   r   s	            r    ,test_space_inference_from_embedding_functionz:TestNewSchema.test_space_inference_from_embedding_function  s    	Q6K!"23+	@XF .9OPI-#,M }5 002I (
3LA.QO"8,W5AAA  )0>|L  $H-g6(BBB "77	BL((33???((33@@LLL%%00==DDJJhVV  $$\2==III$$\2==JJVVV!!,/::GGNNTT  #<<-.>? =#<<-.>? =s   E;F 
F c                 H   ddl m} t        |d<   	 t               }t        d      }t	        |d      }|j                  |       |j                         }|d	   d
   d   }|d   d   dk(  sJ |d   d   d
   d   }|d   d   dk(  sJ t        j                  |      }|j                  j                  J |j                  j                  j                  J |j                  j                  j                  j                  j                  dk(  sJ |j                  d   j                  J |j                  d   j                  j                  J |j                  d   j                  j                  j                  j                  dk(  sJ 	 d|v r|d= yy# d|v r|d= w w xY w)zTTest that explicit space parameter overrides the embedding function's default space.r   r   rJ   override_testr   rR   r}   r   r5   r[   r`   ra   r2   r   rj   rZ   NrL  rM  s	            r    8test_explicit_space_overrides_embedding_function_defaultzFTestNewSchema.test_explicit_space_overrides_embedding_function_default  s    	Q6K!"23-	@XF
 .II-#,M }5 002I (
3LA.QO)'2d::  )0>|L  !*73t;; "77	BL((33???((33@@LLL((33@@GGMMQUUUU$$\2==III$$\2==JJVVV!!,/::GGNNTT  #<<-.>? =#<<-.>? =s   E<F 
F!c                    t               }t        dd      }|j                  |       |j                         }|d   d   d   }|d   d	   dk(  sJ |d   d
   d   dk(  sJ |d   d   d   d   }|d   d	   dk(  sJ |d   d
   d   dk(  sJ t        j                  |      }|j
                  j                  J |j
                  j                  j                  J |j
                  j                  j                  j                  j                  dk(  sJ |j
                  j                  j                  j                  j                  	 J y)zITest space handling when no embedding function is provided (legacy mode).NrS   rQ  r5   r[   r`   ra   r2   r   r}   r   r   rj   rZ   )r   r   rt   r   r   r[   r`   ra   r2   r   r}   )r   rn   r   r   r   r   r   s          r    /test_space_inference_with_no_embedding_functionz=TestNewSchema.test_space_inference_with_no_embedding_function!	  s    *#
 	=1 ,,.	 $J/=nMx(1T999x()=>vF(RRR$V,\:<HX)'2d:::)*>?G8SSS 33I>$$//;;;$$//<<HHH$$//<<CCIITQQQ!!,,99@@SS	
r"   c                    ddl m} t        |d<   	 t               }t        d      }t	        |      }|j                  |       |j                         }|d   d	   d
   d   d   dk(  sJ t        j                  |      }|j                         }|d   d	   d
   d   d   dk(  sJ t        j                  |      }|j                         }	|	d   d	   d
   d   d   dk(  sJ ||fD ]u  }
|
j                  j                  J |
j                  j                  j                  J |
j                  j                  j                  j                  j                  dk(  ruJ  	 d|v r|d= yy# d|v r|d= w w xY w)zQTest that inferred space remains stable across multiple serialization roundtrips.r   r   rJ   roundtrip_space_testr   rK  r5   r[   r`   ra   r2   r   rO   N)r   r   r>   r   r   rt   r   r   r[   r`   ra   r2   r   )r   r   r   r   r   r%  r&  r'  r(  r)  rn   s              r    (test_space_inference_multiple_roundtripsz6TestNewSchema.test_space_inference_multiple_roundtripsB	  s    	Q6K!"23)	@xH-9OPI-#,M !!!7 ..0Ej!,/?I'R 2259G --/Ej!,/?I'R 2259G --/Ej!,/?I'R
 #G, X11===11>>JJJ11>>EEKKxWWWX  #<<-.>? =#<<-.>? =s   EE' E' '
E1c                    t               }t        dt                     }|j                  |d       d|j                  v sJ |j                  d   }|j
                  J |j
                  j                  J |j
                  j                  j                  du sJ |j
                  j                  j                  j                  dk(  sJ t               }|j                  |d       d|j                  v sJ |j                  d   }|j                  J |j                  j                  J |j                  j                  j                  du sJ |j
                  j                  j                  j                  dk(  sJ y)zZTest that each key override has its own independent config (no inheritance from defaults).default_sourcer   r@  rr   NTrA  )r   r   r   rt   rj   rb   rc   r^   r2   rl   r   r\   r_   )r   rn   r   r@  r   rA  s         r    "test_keys_have_independent_configsz0TestNewSchema.test_keys_have_independent_configst	  s~    0':<
 	=h? 6;;&&&X&##///##77CCC##77??4GGG  44;;FF 	
 
 23=h? 6;;&&&X&}}(((}}22>>>}}22::dBBB   44;;FF 	
 r"   c                    ddl m} t        |d<   	 t               }t        d      }t	        dd      }t        |d	|
      }|j                  |       |j                  d   j                  j                  j                  j                  }|J |j                  dk(  sJ |j                  dk(  sJ t        d      }t	        dd      }t        |d|
      }	|j                  |	       |j                  j                  J |j                  j                  j                  J |j                  j                  j                  j                  j                  dk(  sJ |j                  j                  j                  j                  j                  J |j                  j                  j                  j                  j                  j                  dk(  sJ |j                  j                  j                  j                  j                  j                  dk(  sJ |j                  d   j                  J |j                  d   j                  j                  J |j                  d   j                  j                  j                  j                  }
|
J |
j                  dk(  sJ |
j                  dk(  sJ |j                  d   j                  j                  j                  j                  dk(  sJ 	 d|v r|d= yy# d|v r|d= w w xY w)zPTest that changes to global defaults don't affect already-created key overrides.r   r   rJ   initial_modelr   r      r   rO   r   r5   rZ   Nupdated_modelr   r   rR   )r   r   r>   r   r   r   rt   rj   r`   ra   r2   r   r   r   r[   r   )r   r   rn   ef1hnsw1vector_config1initial_embedding_hnswef2hnsw2r   updated_embedding_hnsws              r    :test_global_default_changes_dont_affect_existing_overrideszHTestNewSchema.test_global_default_changes_dont_affect_existing_overrides	  s    	Q6K!"238	@XF (?CC#CrJE.#&hUN ~6 &,[[&jffTT # *555)99S@@@)772=== (?CC#CrJE.#&dN ~6 ??--999??--::FFF??--::AAGG4OOO??--::AAFFRRR**77>>CCSS **77>>CCQQUWWW ;;|,77CCC;;|,77DDPPP%+[[&jffTT # *555)99S@@@)772===L)44AAHHNNRVVV
  #<<-.>? =#<<-.>? =s   K!K< <
Lc                    t               }|j                  t        dt                     d       |j                  t	               d       |j                  t	               d       |j
                  d   j                  j                  j                  j                  dk(  sJ |j
                  d   j                  j                  j                  du sJ |j
                  d   j                  j                  j                  du sJ |j                  t	               d       |j
                  d   j                  j                  j                  du sJ |j
                  d   }|j                  J |j                  j                  J |j                  j                  j                  du sJ |j                  j                  j                  j                  dk(  sJ |j
                  d   }|j                  J |j                  j                  J |j                  j                  j                  du sJ |j                         }t        j                  |      }|j
                  d   j                  j                  j                  j                  dk(  sJ |j
                  d   j                  j                  j                  du sJ |j
                  d   j                  j                  j                  du sJ y	)
zBTest that modifying one key's overrides doesn't affect other keys.source_ar   key_arr   key_bkey_cTFN)r   rt   r   r   r   rj   rb   rc   r2   rl   r\   r_   r^   r   r   r   )r   rn   ri  rk  r   r   s         r    +test_key_specific_overrides_are_independentz9TestNewSchema.test_key_specific_overrides_are_independent	  s    	*%:U:W 	 	 	
 	#<#>GL#<#>GL KK ..BBIITT	
 {{7#**@@HHDPPP{{7#**@@HHDPPP 	#<#>GL {{7#**@@HHEQQQ G$""...""66BBB""66>>$FFF""66==HHJVVVG$||'''||11===||1199TAAA ,,.	33I> m//zz	

   )00FFNNRWWWW  )00FFNNRVVVVr"   c                    t               }|j                  j                  J |j                  j                  j                  J |j                  j                  j                  j                  du sJ |j                  t                      |j                  j                  j                  j                  du sJ |j                  t               d       |j                  t               d       |j                  j                  j                  j                  du sJ |j                  d   }|j                  J |j                  j                  J |j                  j                  j                  du sJ |j                  d   }|j                  J |j                  j                  J |j                  j                  j                  du sJ |j                         }d|d   v sJ d|d   v sJ d	|d   v sJ d
|d   v sJ d|d   vsJ y)zDTest disabling an index globally, then enabling it on specific keys.NTr5   Fimportant_fieldrr   searchable_fieldrj   rY   rZ   other_field)
r   r[   r\   r_   r^   r   r   rt   rj   r   )r   rn   	important
searchabler   s        r    +test_global_default_disable_then_key_enablez9TestNewSchema.test_global_default_disable_then_key_enable
  s    %%111%%;;GGG%%;;CCtKKK 	#<#>?%%;;CCuLLL 	#<#>DUV#<#>DVW %%;;CCuLLL KK 12	+++55AAA55==EEE[[!34
  ,,,  66BBB  66>>$FFF ,,.	 !If$5555!Yv%6666i////y0000 If$5555r"   c                 .   t               }|j                  t        dt                     d       |j                  d   }|j
                  J |j
                  j                  J |j
                  j                  j                  du sJ |j                  J |j                  J |j                  J |j                  J |j                  J |j                         }|d   d   }d|v sJ d	|vsJ d
|vsJ d|vsJ d|vsJ d|vsJ t        j                  |      }|j                  d   }|j
                  J |j
                  j                  J |j
                  j                  j                  du sJ |j                  J |j                  J y)z\Test that when you override one aspect of a value type, other indexes still follow defaults.	my_sourcer   multi_index_fieldrr   NTrj   rb   r\   r   r   r   r`   )r   rt   r   r   rj   rb   rc   r^   r\   rd   rf   rh   r`   r   r   )r   rn   fieldr   
field_jsonr   	des_fields          r    )test_partial_override_fills_from_defaultsz7TestNewSchema.test_partial_override_fills_from_defaults>
  s    	*&;V;X $	 	 	
 /0 ""...""66BBB""66>>$FFF ||###&&&  (((}}$$$''' ,,.	v&':;
 *,,,z)))J&&&j(((Z''':--- 33I> %%&9:	 &&222&&::FFF&&::BBdJJJ '''""***r"   c                    t        j                  d      }|j                  t        j                  k(  sJ |j
                  dk(  sJ |j                         du sJ t        j                  d      }|j                         du sJ y)z(Test basic CMEK creation and validation.zRprojects/test-project/locations/us-central1/keyRings/test-ring/cryptoKeys/test-keyTzinvalid-formatFN)r   gcpproviderr   GCPresourcevalidate_pattern)r   cmekinvalid_cmeks      r    test_cmek_basic_creationz&TestNewSchema.test_cmek_basic_creationr
  s     xx`
 }} 0 0000MMcd	
d
 $$&$... xx 01,,.%777r"   c                    t        j                  d      }|j                         }|ddik(  sJ d|v sJ |d   dk(  sJ t        j                  |      }|j                  t
        j                  k(  sJ |j                  |j                  k(  sJ y)z,Test CMEK serialization and deserialization..projects/p/locations/l/keyRings/r/cryptoKeys/kr|  N)r   r|  to_dict	from_dictr}  r   r~  r  )r   r  	cmek_dictrestoreds       r    test_cmek_serializationz%TestNewSchema.test_cmek_serialization
  s    xxHI LLN	U$TUUUU	!!!#SSSS >>),  L$4$4444  DMM111r"   c                    t               }|j                  J t        j                  d      }|j	                  |      }||u sJ |j                  J |j                  j
                  t        j                  k(  sJ |j                  j                  dk(  sJ |j	                  d       |j                  J t        j                  d      }t               j	                  |      }|j                  J |j                  j                  dk(  sJ y)z:Test CMEK integration with Schema using set_cmek() method.Nz7projects/test/locations/us/keyRings/ring/cryptoKeys/keyr  )	r   r  r   r|  set_cmekr}  r   r~  r  )r   rn   r  rv   cmek2r(  s         r    test_cmek_in_schemaz!TestNewSchema.test_cmek_in_schema
  s    {{""" xxQR&  {{&&&{{##|'7'7777KK  HI	
I
 	{{""" IJ(##E*||'''||$$(XXXXr"   c                    t        j                  d      }t               j                  |      }|j	                         }d|v sJ |d   ddik(  sJ d|d   v sJ |d   d   dk(  sJ t        j
                  |      }|j                  J |j                  j                  t        j                  k(  sJ |j                  j                  |j                  k(  sJ y)z$Test Schema serialization with CMEK.r  r  r|  N)r   r|  r   r  r   r   r  r}  r   r~  r  )r   r  rn   r   r   s        r    test_cmek_schema_serializationz,TestNewSchema.test_cmek_schema_serialization
  s    xxHI""4( ,,.	 """ C%
 
 	
 
 	&))))fe$(XX	
X 33I>  ,,,  ))\-=-====  ))T]]:::r"   c                     t               }|j                         }d|vsJ t        j                  |      }|j                  J y)z@Test Schema serialization without CMEK (backward compatibility).r  N)r   r   r   r  )r   rn   r   r   s       r    +test_cmek_schema_without_cmek_serializationz9TestNewSchema.test_cmek_schema_without_cmek_serialization
  sM     ,,.	 Y&&& 33I>  (((r"   c                 ,   t        j                  t        d      5  t        j                  i i i d       ddd       t        j                  t        d      5  t        j                  i i ddid       ddd       y# 1 sw Y   IxY w# 1 sw Y   yxY w)zCTest that invalid CMEK data raises a warning and sets cmek to None.z,Unsupported or missing CMEK provider in datar   )r[   rj   r  Ninvalid_providerzsome-resource)r   r   r   r   r   r0   s    r    !test_cmek_invalid_deserializationz/TestNewSchema.test_cmek_invalid_deserialization
  s    ]]:-[\ 	((R8	
 ]]:-[\ 	(( "/A	 		 	
	 	s   A>B
>B
Br$   N).r7   r8   r9   r:   ro   rw   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r  r  r  r  r  r-  r>  rD  rH  rN  rR  rT  rW  rZ  rf  rl  rs  rz  r  r  r  r  r  r  r.   r"   r    rV   rV   L   s   .G
R+
Z1
f(
Tc
J5In1
f*EXH<. B<
4%QN)@Vj
XG
Rr@hI@VW=r
2Lh&LP8Gt*1X\@|JJX+DZ:
x2@h4@l
B0@d%
N?@B4Wl*6X2+h8&2Y@;0)r"   rV   c                      t               } t               }t        j                  t        d      5  | j                  |       ddd       y# 1 sw Y   yxY w)zITest that sparse vector index cannot be created globally (without a key).z5Sparse vector index must be created on a specific keyr   r5   N)r   r   r   r   r   rt   rn   r   s     r    -test_sparse_vector_cannot_be_created_globallyr  
  sM    XF+-M 
Q
 2 	=12 2 2s   AAc                     t               } t               }| j                  |d       | j                  d   j                  J | j                  d   j                  j
                  J | j                  d   j                  j
                  j                  du sJ t        j                  t        d      5  | j                  |d       ddd       y# 1 sw Y   yxY w)zITest that sparse vector index cannot be deleted (temporarily disallowed).r   rr   NTz7Deleting sparse vector index is not currently supportedr   )r   r   rt   rj   rb   rc   r^   r   r   r   r   r  s     r    $test_sparse_vector_cannot_be_deletedr  
  s    XF+-M }(;;;x ..:::;;x ..BBNNN;;x ..BBJJdRRR 
S
 @ 	=h?@ @ @s   4CCc                     t               } t               }| j                  |d       d| j                  v sJ | j                  d   j                  J | j                  d   j                  j
                  J | j                  d   j                  j
                  j                  du sJ t               }| j                  |t        d             d| j                  v sJ | j                  d   j                  J | j                  d   j                  j                  J | j                  d   j                  j                  j                  du sJ t               }|j                  |d       t               }|j                  |t        d             |j                  d   j                  J |j                  d   j                  j
                  J |j                  d   j                  J |j                  d   j                  j
                  J |j                  d   j                  j
                  j                  |j                  d   j                  j
                  j                  k(  sJ y)zLTest that create_index accepts both str and Key types for the key parameter.test_field_strrr   NTtest_field_key
same_field)r   r   rt   rj   r\   r_   r^   r   r   rd   re   rn   r   r   r(  r*  s        r    "test_create_index_accepts_key_typer  	  sF   XF ./M
}2BC v{{***;;'(//;;;;;'(//EEQQQ;;'(//EEMMQUUUU ()J
zs3C/DE v{{***;;'(22>>>;;'(22EEQQQ;;'(22EEMMQUUUU hG<@hG3|3DE <<%,,888<<%,,BBNNN<<%,,888<<%,,BBNNN\"))??GG<<%,,BBJJ	K	Kr"   c                     t               } t               }t               }| j                  |d       d| j                  v sJ | j                  d   j
                  J | j                  d   j
                  j                  J | j                  d   j
                  j                  j                  du sJ | j                  |t        d             d| j                  v sJ | j                  d   j                  J | j                  d   j                  j                  J | j                  d   j                  j                  j                  du sJ t               }|j                  |d       t               }|j                  |t        d             |j                  d   j
                  J |j                  d   j
                  j                  J |j                  d   j
                  J |j                  d   j
                  j                  J |j                  d   j
                  j                  j                  |j                  d   j
                  j                  j                  k(  sJ y)zLTest that delete_index accepts both str and Key types for the key parameter.r  rr   NFr  r  )r   r   r   r   rj   r\   r_   r^   r   rd   re   r  s        r    "test_delete_index_accepts_key_typer  3  sH   XF ./M')J }2BC v{{***;;'(//;;;;;'(//EEQQQ;;'(//EEMMQVVVV zs3C/DE v{{***;;'(22>>>;;'(22EEQQQ;;'(22EEMMQVVVV hG<@hG3|3DE <<%,,888<<%,,BBNNN<<%,,888<<%,,BBNNN\"))??GG<<%,,BBJJ	K	Kr"   c                     t               } t               }t        j                  t        d      5  | j                  |t        j                         ddd       t        j                  t        d      5  | j                  |t        j                         ddd       t        j                  t        d      5  | j                  |d       ddd       t        j                  t        d      5  | j                  |d       ddd       t        j                  t        d      5  | j                  |t        d	             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   nxY w# 1 sw Y   yxY w)
zPTest that create_index rejects special keys like Key.DOCUMENT and Key.EMBEDDING.r   r   rr   Nr   rY   key cannot begin with '#'#custom_key#custom)	r   r   r   r   r   rt   r   DOCUMENT	EMBEDDINGrn   r   s     r    &test_create_index_rejects_special_keysr  _  sc   XF-/M 
J
 D 	=cllCD 
K
 E 	=cmmDE 
J
 C 	=kBC 
z)D	E E=mDE 
z)D	E F=c)nEF F-D DE EC CE EF F;   "E5"E:E&1E2(E>EE#&E/2E;>Fc                     t               } t               }t        j                  t        d      5  | j                  |t        j                         ddd       t        j                  t        d      5  | j                  |t        j                         ddd       t        j                  t        d      5  | j                  |d       ddd       t        j                  t        d      5  | j                  |d       ddd       t        j                  t        d      5  | j                  |t        d	             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   nxY w# 1 sw Y   yxY w)
zPTest that delete_index rejects special keys like Key.DOCUMENT and Key.EMBEDDING.z.Cannot delete index on special key '#document'r   rr   Nz/Cannot delete index on special key '#embedding'rZ   r  r  r  )	r   r   r   r   r   r   r   r  r  r  s     r    &test_delete_index_rejects_special_keysr    sc   XF-/M 
J
 D 	=cllCD 
K
 E 	=cmmDE 
K
 D 	=lCD 
z)D	E E=mDE 
z)D	E F=c)nEF F-D DE ED DE EF Fr  c                     t        d      } | j                  dk(  sJ t        | j                  t              sJ t        t	        d            }|j                  dk(  sJ t        |j                  t              sJ t        t        j
                        }|j                  dk(  sJ t        |j                  t              sJ t        d      }t        t	        d            }|j                  |j                  k(  sJ t        d      }|j                  J t        t	        d            }|j                         }|d   dk(  sJ t        |d   t              sJ y)zFTest that VectorIndexConfig.source_key accepts both str and Key types.my_fieldrl   rY   testNserialize_testrl   )r   rl   r   r;   r   r  
model_dumpconfig1config2config3config4config5config6config7config_dicts           r    4test_vector_index_config_source_key_accepts_key_typer    sG     :6G+++g((#...  3z?;G+++g((#...  3<<8G,,,g((#...  62G3v;7G!3!3333  40G%%%  3/?+@AG$$&K|$(8888k,/555r"   c                     t        d      } | j                  dk(  sJ t        | j                  t              sJ t        t	        d            }|j                  dk(  sJ t        |j                  t              sJ t        t        j
                        }|j                  dk(  sJ t        |j                  t              sJ t        d      }t        t	        d            }|j                  |j                  k(  sJ t        d      }|j                  J t        t	        d            }|j                         }|d   dk(  sJ t        |d   t              sJ y)zLTest that SparseVectorIndexConfig.source_key accepts both str and Key types.r  r  rY   r  Nr  rl   )r   rl   r   r;   r   r  r  r  s           r    ;test_sparse_vector_index_config_source_key_accepts_key_typer    sG    &<G+++g((#... &ZAG+++g((#... &>G,,,g((#... &8G%V=G!3!3333 &6G%%% &5E1FGG$$&K|$(8888k,/555r"   c                     t        j                  t        d      5  t        d       ddd       t        j                  t        d      5  t        ddg       ddd       t        j                  t        d      5  t	        d       ddd       t        j                  t        d      5  t	        ddi       ddd       y# 1 sw Y   xY w# 1 sw Y   xxY w# 1 sw Y   TxY w# 1 sw Y   yxY w)z@Test that config validators reject invalid types for source_key.zsource_key must be str or Keyr   {   r  Nnotvalid)r   r   r   r   r   r.   r"   r    ,test_config_source_key_rejects_invalid_typesr    s     
z)H	I *S)* 
z)H	I 7eW%567 
z)H	I 03/0 
z)H	I =E7+;<= =* *7 70 0= =s/   CC>C.C*CCC'*C3c                     t        d      } | j                  dk(  sJ t        t        j                        }|j                  dk(  sJ t	        j
                  t        d      5  t        d       ddd       t	        j
                  t        d      5  t        t        j                         ddd       t	        j
                  t        d      5  t        d       ddd       t	        j
                  t        d      5  t        d       ddd       t	        j
                  t        d      5  t        d	       ddd       t        d
      }|j                  d
k(  sJ t        d      }|j                  dk(  sJ t        t        j                        }|j                  dk(  sJ t	        j
                  t        d      5  t        d       ddd       t	        j
                  t        d      5  t        t        j                         ddd       t	        j
                  t        d      5  t        d       ddd       t        d
      }|j                  d
k(  sJ 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   xY w# 1 sw Y   xY w# 1 sw Y   xY w)zGTest that source_key only allows #document, rejects other special keys.rY   r  z source_key cannot begin with '#'r   rZ   Nz	#metadataz#scorer  r  )	r   rl   r   r  r   r   r   r  r   )r  r  r  r  r  r  s         r    -test_config_source_key_validates_special_keysr    sC     ;7G,,,  3<<8G,,, 
z)K	L 3\23 
z)K	L 4S]]34 
z)K	L 2[12 
z)K	L /X./ 
z)K	L 0Y/0  :6G+++ &=G,,, &>G,,, 
z)K	L 9<89 
z)K	L :3==9: 
z)K	L 8;78 &<G+++]3 34 42 2/ /0 0"9 9: :8 8s`   $I(I5JJ2JJ)7J55K(I25I?JJJ&)J25J>K
c                  (   t               } t        j                  t              5 }| j	                  dt        d             ddd       t        j                        }d|j                         v sJ d|j                         v sJ y# 1 sw Y   GxY w)zkTest that SparseVectorIndexConfig raises ValueError when source_key is provided without embedding_function.invalid_sparser   r  )rs   r2   Nrl   r}   )	r   r   r   r   rt   r   r;   valuelower)rn   exc_info	error_msgs      r    5test_sparse_vector_config_requires_ef_with_source_keyr  /  s    XF 
z	" 
h *' 	 	

 HNN#I9??,,,,9??#4444
 
s   BBc            	      <   t        j                  t        t        f      5 } t	        t
        j                         ddd       t         j                        }d|j                         v sJ d|j                         v sd|j                         v sJ t        j                  t        t        f      5 } t        d       ddd       t        | j                        }d|j                         v sJ d|j                         v sd|j                         v sJ t        j                  t        t        f      5 } t        d       ddd       t        | j                        }d|j                         v sJ d|j                         v sd|j                         v sJ t        j                  t        t        f      5 } t        d       ddd       t        | j                        }d|j                         v sJ d|j                         v sd|j                         v sJ t        j                  t        t        f      5 } t        d	       ddd       t        | j                        }d|j                         v sJ d|j                         v sd|j                         v sJ t        j                  t        t        f      5 } t        d
       ddd       t        | j                        }d|j                         v sJ d|j                         v sd|j                         v sJ t        j                  t        t        f      5 } t        d       ddd       t        | j                        }d|j                         v sJ d|j                         v sd|j                         v sJ t        j                  t        t        f      5 } t!        d	       ddd       t        | j                        }d|j                         v sJ t!        ddddddd      }|j"                  dk(  sJ |j$                  dk(  sJ t        j                  t        t        f      5 } t'        d	       ddd       t        | j                        }d|j                         v sJ t'        d,i dddddddddddddddd d!dd"d#d$dd%d&d'dd(d)d*d)d+d}|j(                  dk(  sJ |j*                  dk(  sJ |j"                  d)k(  sJ t-        |d      rJ t-        |d      rJ t-        |d!      rJ y# 1 sw Y   xY w# 1 sw Y   zx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   xY w# 1 sw Y   2xY w)-z;Test that all config classes reject invalid/unknown fields.r   Nrs   extra	permittedr  )invalid_fieldr  r  gGz?Tr   r]     i  g333333?)r   r   r   num_threads
batch_sizesync_thresholdresize_factorr   r   search_rng_factorr'   search_rng_epsilong      $@nreplica_count   r   r   write_rng_factorwrite_rng_epsilong      @split_thresholdr   num_samples_kmeansinitial_lambdag      Y@reassign_neighbor_countmerge_threshold   num_centers_to_merge_tor   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   hasattr)r  r  r2   r   s       r    )test_config_classes_reject_invalid_fieldsr  C  sr    

O4	5 2CLL12 HNN#IIOO%%%%ioo'';)//:K+KKK 

O4	5 0/0 HNN#Iioo////ioo'';)//:K+KKK 

O4	5 -V,- HNN#Iioo////ioo'';)//:K+KKK 

O4	5 8!78 HNN#Iioo////ioo'';)//:K+KKK 

O4	5 2S12 HNN#Iioo////ioo'';)//:K+KKK 

O4	5 5 t45 HNN#Iioo////ioo'';)//:K+KKK 

O4	5 4d34 HNN#Iioo////ioo'';)//:K+KKK 

O4	5 +c*+ HNN#Iioo//// F !!S(((2%%% 

O4	5 ,s+, HNN#Iioo//// $    	
        !#  !"    !L& %%+++$$***''3...|%8999|%5666|%9::::]2 20 0- -8 82 25 54 4+ +*, ,sk   T)1T64U7U:U=U* U7VV)T36U UUU'*U47VVVr  )*chromadb.api.typesr   r   r   r   r   r   r	   r
   r   r   r   r   r   r   r   r   &chromadb.execution.expression.operatorr   typingr   r   r   pydanticr   r   r;   r   r>   rV   r  r  r  r  r  r  r  r  r  r  r  r  r.   r"   r    <module>r     s       $ 7 " " $ N"9$s)"D N*&-d3i8 &<]* ]*~T	2@$'T)XF@F@6B6B=":,z5(q;r"   