codekingpro/portable-devtools
114k
1from chromadb.api.types import (
2 Schema,
3 SparseVectorIndexConfig,
4 SparseEmbeddingFunction,
5 SparseVector,
6 StringInvertedIndexConfig,
7 IntInvertedIndexConfig,
8 FloatInvertedIndexConfig,
9 BoolInvertedIndexConfig,
10 VectorIndexConfig,
11 HnswIndexConfig,
12 SpannIndexConfig,
13 FtsIndexConfig,
14 EmbeddingFunction,
15 Embeddings,
16 Cmek,
17 CmekProvider,
18)
19from chromadb.execution.expression.operator import Key
20from typing import List, Dict, Any
21from pydantic import ValidationError
22import pytest
23
24
25class MockSparseEmbeddingFunction(SparseEmbeddingFunction[List[str]]):
26 """Mock sparse embedding function for testing."""
27
28 def __init__(self, name: str = "mock_sparse"):
29 self._name = name
30
31 def __call__(self, input: List[str]) -> List[SparseVector]:
32 return [SparseVector(indices=[0, 1], values=[1.0, 1.0]) for _ in input]
33
34 @staticmethod
35 def name() -> str:
36 return "mock_sparse"
37
38 def get_config(self) -> Dict[str, Any]:
39 return {"name": self._name}
40
41 @staticmethod
42 def build_from_config(config: Dict[str, Any]) -> "MockSparseEmbeddingFunction":
43 return MockSparseEmbeddingFunction(config.get("name", "mock_sparse"))
44
45
46class MockEmbeddingFunction(EmbeddingFunction[List[str]]):
47 """Mock embedding function for testing."""
48
49 def __init__(self, model_name: str = "mock_model"):
50 self._model_name = model_name
51
52 def __call__(self, input: List[str]) -> Embeddings:
53 import numpy as np
54
55 # Return mock embeddings (3-dimensional)
56 return [np.array([1.0, 2.0, 3.0], dtype=np.float32) for _ in input]
57
58 @staticmethod
59 def name() -> str:
60 return "mock_embedding"
61
62 def get_config(self) -> Dict[str, Any]:
63 return {"model_name": self._model_name}
64
65 @staticmethod
66 def build_from_config(config: Dict[str, Any]) -> "MockEmbeddingFunction":
67 return MockEmbeddingFunction(config.get("model_name", "mock_model"))
68
69 def default_space(self) -> str: # type: ignore
70 return "cosine"
71
72 def supported_spaces(self) -> List[str]: # type: ignore
73 return ["cosine", "l2", "ip"]
74
75
76class TestNewSchema:
77 """Test cases for the new Schema class."""
78
79 def test_default_schema_initialization(self) -> None:
80 """Test that Schema() initializes with correct defaults."""
81 schema = Schema()
82
83 # Verify defaults are populated
84 assert schema.defaults is not None
85
86 # Verify string value type defaults
87 assert schema.defaults.string is not None
88 assert schema.defaults.string.fts_index is not None
89 assert schema.defaults.string.fts_index.enabled is False # Disabled by default
90 assert schema.defaults.string.string_inverted_index is not None
91 assert (
92 schema.defaults.string.string_inverted_index.enabled is True
93 ) # Enabled by default
94
95 # Verify float_list value type defaults
96 assert schema.defaults.float_list is not None
97 assert schema.defaults.float_list.vector_index is not None
98 assert (
99 schema.defaults.float_list.vector_index.enabled is False
100 ) # Disabled by default
101
102 # Verify sparse_vector value type defaults
103 assert schema.defaults.sparse_vector is not None
104 assert schema.defaults.sparse_vector.sparse_vector_index is not None
105 assert (
106 schema.defaults.sparse_vector.sparse_vector_index.enabled is False
107 ) # Disabled by default
108
109 # Verify int_value type defaults
110 assert schema.defaults.int_value is not None
111 assert schema.defaults.int_value.int_inverted_index is not None
112 assert (
113 schema.defaults.int_value.int_inverted_index.enabled is True
114 ) # Enabled by default
115
116 # Verify float_value type defaults
117 assert schema.defaults.float_value is not None
118 assert schema.defaults.float_value.float_inverted_index is not None
119 assert (
120 schema.defaults.float_value.float_inverted_index.enabled is True
121 ) # Enabled by default
122
123 # Verify boolean type defaults
124 assert schema.defaults.boolean is not None
125 assert schema.defaults.boolean.bool_inverted_index is not None
126 assert (
127 schema.defaults.boolean.bool_inverted_index.enabled is True
128 ) # Enabled by default
129
130 # Verify keys are populated
131 assert schema.keys is not None
132 assert len(schema.keys) == 2 # Should have #document and #embedding
133
134 # Verify #document key override (FTS enabled, string inverted disabled)
135 assert "#document" in schema.keys
136 assert schema.keys["#document"].string is not None
137 assert schema.keys["#document"].string.fts_index is not None
138 assert schema.keys["#document"].string.fts_index.enabled is True
139 assert schema.keys["#document"].string.string_inverted_index is not None
140 assert schema.keys["#document"].string.string_inverted_index.enabled is False
141
142 # Verify #embedding key override (vector index enabled)
143 assert "#embedding" in schema.keys
144 assert schema.keys["#embedding"].float_list is not None
145 assert schema.keys["#embedding"].float_list.vector_index is not None
146 assert schema.keys["#embedding"].float_list.vector_index.enabled is True
147 assert (
148 schema.keys["#embedding"].float_list.vector_index.config.source_key
149 == "#document"
150 )
151
152 def test_create_sparse_vector_index_on_key(self) -> None:
153 """Test creating a sparse vector index on a specific key with default config."""
154 schema = Schema()
155
156 # Create sparse vector index on a custom key with default config
157 config = SparseVectorIndexConfig()
158 result = schema.create_index(config=config, key="custom_sparse_key")
159
160 # Should return self for chaining
161 assert result is schema
162
163 # Verify the key override was created
164 assert "custom_sparse_key" in schema.keys
165
166 # Verify sparse_vector type was set for this key
167 assert schema.keys["custom_sparse_key"].sparse_vector is not None
168 assert (
169 schema.keys["custom_sparse_key"].sparse_vector.sparse_vector_index
170 is not None
171 )
172
173 # Verify it's enabled and has the correct config
174 assert (
175 schema.keys["custom_sparse_key"].sparse_vector.sparse_vector_index.enabled
176 is True
177 )
178 assert (
179 schema.keys["custom_sparse_key"].sparse_vector.sparse_vector_index.config
180 == config
181 )
182
183 # Verify other value types for this key are None (not initialized)
184 assert schema.keys["custom_sparse_key"].string is None
185 assert schema.keys["custom_sparse_key"].float_list is None
186 assert schema.keys["custom_sparse_key"].int_value is None
187 assert schema.keys["custom_sparse_key"].float_value is None
188 assert schema.keys["custom_sparse_key"].boolean is None
189
190 # Verify defaults were not affected
191 assert schema.defaults.sparse_vector is not None
192 assert schema.defaults.sparse_vector.sparse_vector_index is not None
193 assert (
194 schema.defaults.sparse_vector.sparse_vector_index.enabled is False
195 ) # Still disabled by default
196
197 def test_create_sparse_vector_index_with_custom_config(self) -> None:
198 """Test creating a sparse vector index with custom config including embedding function."""
199 schema = Schema()
200
201 # Create custom sparse vector config with embedding function and source key
202 embedding_func = MockSparseEmbeddingFunction(name="custom_sparse_ef")
203 config = SparseVectorIndexConfig(
204 embedding_function=embedding_func, source_key="custom_document_field"
205 )
206
207 # Create sparse vector index on a custom key
208 result = schema.create_index(config=config, key="sparse_embeddings")
209
210 # Should return self for chaining
211 assert result is schema
212
213 # Verify the key override was created
214 assert "sparse_embeddings" in schema.keys
215 assert schema.keys["sparse_embeddings"].sparse_vector is not None
216 assert (
217 schema.keys["sparse_embeddings"].sparse_vector.sparse_vector_index
218 is not None
219 )
220
221 # Verify it's enabled
222 sparse_index = schema.keys[
223 "sparse_embeddings"
224 ].sparse_vector.sparse_vector_index
225 assert sparse_index.enabled is True
226
227 # Verify the config has our custom settings
228 assert sparse_index.config.embedding_function == embedding_func
229 assert sparse_index.config.source_key == "custom_document_field"
230
231 # Verify the embedding function is the same instance
232 assert sparse_index.config.embedding_function.name() == "mock_sparse"
233 assert sparse_index.config.embedding_function.get_config() == {
234 "name": "custom_sparse_ef"
235 }
236
237 # Verify global defaults were not overridden
238 assert schema.defaults.sparse_vector is not None
239 assert schema.defaults.sparse_vector.sparse_vector_index is not None
240 assert (
241 schema.defaults.sparse_vector.sparse_vector_index.enabled is False
242 ) # Still disabled by default
243 assert (
244 schema.defaults.sparse_vector.sparse_vector_index.config.embedding_function
245 is None
246 ) # No custom embedding function
247
248 def test_delete_index_on_key(self) -> None:
249 """Test disabling string inverted index on a specific key."""
250 schema = Schema()
251
252 # Create a config and disable it on a specific key
253 config = StringInvertedIndexConfig()
254 result = schema.delete_index(config=config, key="custom_text_key")
255
256 # Should return self for chaining
257 assert result is schema
258
259 # Verify the key override was created
260 assert "custom_text_key" in schema.keys
261
262 # Verify string inverted index is disabled for this key
263 assert schema.keys["custom_text_key"].string is not None
264 assert schema.keys["custom_text_key"].string.string_inverted_index is not None
265 assert (
266 schema.keys["custom_text_key"].string.string_inverted_index.enabled is False
267 )
268
269 # Verify other keys are not affected - check #document key
270 assert "#document" in schema.keys
271 assert schema.keys["#document"].string is not None
272 assert schema.keys["#document"].string.string_inverted_index is not None
273 assert (
274 schema.keys["#document"].string.string_inverted_index.enabled is False
275 ) # Was disabled by default in #document
276
277 # Verify other keys are not affected - check #embedding key (shouldn't have string config)
278 assert "#embedding" in schema.keys
279 assert (
280 schema.keys["#embedding"].string is None
281 ) # #embedding doesn't have string configs
282
283 # Verify global defaults are not affected
284 assert schema.defaults.string is not None
285 assert schema.defaults.string.string_inverted_index is not None
286 assert (
287 schema.defaults.string.string_inverted_index.enabled is True
288 ) # Global default is still enabled
289
290 def test_chained_create_and_delete_operations(self) -> None:
291 """Test chaining create_index() and delete_index() operations together."""
292 schema = Schema()
293
294 # Chain multiple operations:
295 # 1. Create sparse vector index on "embeddings_key"
296 # 2. Disable string inverted index on "text_key_1"
297 # 3. Disable string inverted index on "text_key_2"
298 sparse_config = SparseVectorIndexConfig(
299 source_key="raw_text", embedding_function=MockSparseEmbeddingFunction()
300 )
301 string_config = StringInvertedIndexConfig()
302
303 result = (
304 schema.create_index(config=sparse_config, key="embeddings_key")
305 .delete_index(config=string_config, key="text_key_1")
306 .delete_index(config=string_config, key="text_key_2")
307 )
308
309 # Should return self for chaining
310 assert result is schema
311
312 # Verify all three key overrides were created
313 assert "embeddings_key" in schema.keys
314 assert "text_key_1" in schema.keys
315 assert "text_key_2" in schema.keys
316
317 # Verify sparse vector index on "embeddings_key" is enabled
318 assert schema.keys["embeddings_key"].sparse_vector is not None
319 assert (
320 schema.keys["embeddings_key"].sparse_vector.sparse_vector_index is not None
321 )
322 assert (
323 schema.keys["embeddings_key"].sparse_vector.sparse_vector_index.enabled
324 is True
325 )
326 assert (
327 schema.keys[
328 "embeddings_key"
329 ].sparse_vector.sparse_vector_index.config.source_key
330 == "raw_text"
331 )
332
333 # Verify only sparse_vector is set for embeddings_key (other types are None)
334 assert schema.keys["embeddings_key"].string is None
335 assert schema.keys["embeddings_key"].float_list is None
336 assert schema.keys["embeddings_key"].int_value is None
337 assert schema.keys["embeddings_key"].float_value is None
338 assert schema.keys["embeddings_key"].boolean is None
339
340 # Verify string inverted index on "text_key_1" is disabled
341 assert schema.keys["text_key_1"].string is not None
342 assert schema.keys["text_key_1"].string.string_inverted_index is not None
343 assert schema.keys["text_key_1"].string.string_inverted_index.enabled is False
344
345 # Verify only string is set for text_key_1 (other types are None)
346 assert schema.keys["text_key_1"].sparse_vector is None
347 assert schema.keys["text_key_1"].float_list is None
348 assert schema.keys["text_key_1"].int_value is None
349 assert schema.keys["text_key_1"].float_value is None
350 assert schema.keys["text_key_1"].boolean is None
351
352 # Verify string inverted index on "text_key_2" is disabled
353 assert schema.keys["text_key_2"].string is not None
354 assert schema.keys["text_key_2"].string.string_inverted_index is not None
355 assert schema.keys["text_key_2"].string.string_inverted_index.enabled is False
356
357 # Verify only string is set for text_key_2 (other types are None)
358 assert schema.keys["text_key_2"].sparse_vector is None
359 assert schema.keys["text_key_2"].float_list is None
360 assert schema.keys["text_key_2"].int_value is None
361 assert schema.keys["text_key_2"].float_value is None
362 assert schema.keys["text_key_2"].boolean is None
363
364 # Verify global defaults are not affected
365 assert schema.defaults.sparse_vector is not None
366 assert schema.defaults.sparse_vector.sparse_vector_index is not None
367 assert (
368 schema.defaults.sparse_vector.sparse_vector_index.enabled is False
369 ) # Still disabled globally
370
371 assert schema.defaults.string is not None
372 assert schema.defaults.string.string_inverted_index is not None
373 assert (
374 schema.defaults.string.string_inverted_index.enabled is True
375 ) # Still enabled globally
376
377 # Verify pre-existing key overrides (#document, #embedding) are not affected
378 assert "#document" in schema.keys
379 assert "#embedding" in schema.keys
380 assert schema.keys["#document"].string is not None
381 assert schema.keys["#document"].string.fts_index is not None
382 assert (
383 schema.keys["#document"].string.fts_index.enabled is True
384 ) # Still enabled
385 assert schema.keys["#embedding"].float_list is not None
386 assert schema.keys["#embedding"].float_list.vector_index is not None
387 assert (
388 schema.keys["#embedding"].float_list.vector_index.enabled is True
389 ) # Still enabled
390
391 def test_vector_index_config_and_restrictions(self) -> None:
392 """Test vector index configuration and key restrictions."""
393 schema = Schema()
394 vector_config = VectorIndexConfig(space="cosine", source_key="custom_source")
395
396 # Test 1: CAN set vector config globally - applies to defaults and #embedding
397 result = schema.create_index(config=vector_config)
398 assert result is schema # Should return self for chaining
399
400 # Verify the vector config was applied to defaults (enabled state preserved as False)
401 assert schema.defaults.float_list is not None
402 assert schema.defaults.float_list.vector_index is not None
403 assert (
404 schema.defaults.float_list.vector_index.enabled is False
405 ) # Still disabled in defaults
406 assert schema.defaults.float_list.vector_index.config.space == "cosine"
407 assert (
408 schema.defaults.float_list.vector_index.config.source_key == "custom_source"
409 )
410
411 # Verify the vector config was also applied to #embedding (enabled state preserved as True)
412 # Note: source_key should NOT be overridden on #embedding - it should stay as "#document"
413 assert schema.keys["#embedding"].float_list is not None
414 assert schema.keys["#embedding"].float_list.vector_index is not None
415 assert (
416 schema.keys["#embedding"].float_list.vector_index.enabled is True
417 ) # Still enabled on #embedding
418 assert (
419 schema.keys["#embedding"].float_list.vector_index.config.space == "cosine"
420 )
421 assert (
422 schema.keys["#embedding"].float_list.vector_index.config.source_key
423 == "#document"
424 ) # Preserved, NOT overridden
425
426 # Test 2: Cannot create vector index on custom key
427 vector_config2 = VectorIndexConfig(space="l2")
428 with pytest.raises(
429 ValueError, match="Vector index cannot be enabled on specific keys"
430 ):
431 schema.create_index(config=vector_config2, key="my_vectors")
432
433 # Test 3: Cannot create vector index on #document key (special key blocked globally)
434 with pytest.raises(
435 ValueError, match="Cannot create index on special key '#document'"
436 ):
437 schema.create_index(config=vector_config2, key="#document")
438
439 # Test 4: Cannot create vector index on #embedding key (special key blocked globally)
440 vector_config3 = VectorIndexConfig(space="ip")
441 with pytest.raises(
442 ValueError, match="Cannot create index on special key '#embedding'"
443 ):
444 schema.create_index(config=vector_config3, key="#embedding")
445
446 def test_vector_index_with_embedding_function_and_hnsw(self) -> None:
447 """Test setting embedding function and HNSW config for vector index."""
448 schema = Schema()
449
450 # Create a custom embedding function and HNSW config
451 mock_ef = MockEmbeddingFunction(model_name="custom_model_v2")
452 hnsw_config = HnswIndexConfig(
453 ef_construction=200, max_neighbors=32, ef_search=100
454 )
455
456 # Set vector config with embedding function, space, and HNSW config
457 vector_config = VectorIndexConfig(
458 embedding_function=mock_ef,
459 space="l2", # Override default space from EF
460 hnsw=hnsw_config,
461 source_key="custom_document_field",
462 )
463
464 result = schema.create_index(config=vector_config)
465 assert result is schema
466
467 # Verify defaults: should have EF, space, HNSW, and source_key
468 assert schema.defaults.float_list is not None
469 defaults_vector = schema.defaults.float_list.vector_index
470 assert defaults_vector is not None
471 assert defaults_vector.enabled is False
472 assert defaults_vector.config.embedding_function is mock_ef
473 assert defaults_vector.config.embedding_function.name() == "mock_embedding"
474 assert defaults_vector.config.embedding_function.get_config() == {
475 "model_name": "custom_model_v2"
476 }
477 assert defaults_vector.config.space == "l2"
478 assert defaults_vector.config.hnsw is not None
479 assert defaults_vector.config.hnsw.ef_construction == 200
480 assert defaults_vector.config.hnsw.max_neighbors == 32
481 assert defaults_vector.config.hnsw.ef_search == 100
482 assert defaults_vector.config.source_key == "custom_document_field"
483
484 # Verify #embedding: should have EF, space, HNSW, but source_key is preserved as "#document"
485 assert schema.keys["#embedding"].float_list is not None
486 embedding_vector = schema.keys["#embedding"].float_list.vector_index
487 assert embedding_vector is not None
488 assert embedding_vector.enabled is True
489 assert embedding_vector.config.embedding_function is mock_ef
490 assert embedding_vector.config.space == "l2"
491 assert embedding_vector.config.hnsw is not None
492 assert embedding_vector.config.hnsw.ef_construction == 200
493 assert (
494 embedding_vector.config.source_key == "#document"
495 ) # Preserved, NOT overridden by user config
496
497 def test_fts_index_config_and_restrictions(self) -> None:
498 """Test FTS index configuration and key restrictions."""
499 schema = Schema()
500 fts_config = FtsIndexConfig()
501
502 # Test 1: MUST specify key="#document" for FTS — global (no key) is not allowed
503 with pytest.raises(
504 ValueError, match="FTS index can only be enabled on #document key"
505 ):
506 schema.create_index(config=fts_config)
507
508 # Enable FTS explicitly on #document
509 result = schema.create_index(config=fts_config, key="#document")
510 assert result is schema # Should return self for chaining
511
512 # Verify FTS is enabled on #document
513 assert schema.keys["#document"].string is not None
514 assert schema.keys["#document"].string.fts_index is not None
515 assert schema.keys["#document"].string.fts_index.enabled is True
516 assert schema.keys["#document"].string.fts_index.config == fts_config
517
518 # Test 2: Cannot create FTS index on custom key
519 fts_config2 = FtsIndexConfig()
520 with pytest.raises(
521 ValueError, match="FTS index can only be enabled on #document key"
522 ):
523 schema.create_index(config=fts_config2, key="custom_text_field")
524
525 # Test 3: Cannot create FTS index on #embedding key (special key blocked)
526 with pytest.raises(
527 ValueError, match="Cannot create index on special key '#embedding'"
528 ):
529 schema.create_index(config=fts_config2, key="#embedding")
530
531 # Test 4: Cannot create non-FTS index on #document key
532 with pytest.raises(
533 ValueError, match="Cannot create index on special key '#document'"
534 ):
535 schema.create_index(config=StringInvertedIndexConfig(), key="#document")
536
537 def test_special_keys_blocked_for_all_index_types(self) -> None:
538 """Test that #embedding and #document keys are blocked for all index types."""
539 schema = Schema()
540
541 # Test with StringInvertedIndexConfig on #document
542 string_config = StringInvertedIndexConfig()
543 with pytest.raises(
544 ValueError, match="Cannot create index on special key '#document'"
545 ):
546 schema.create_index(config=string_config, key="#document")
547
548 # Test with StringInvertedIndexConfig on #embedding
549 with pytest.raises(
550 ValueError, match="Cannot create index on special key '#embedding'"
551 ):
552 schema.create_index(config=string_config, key="#embedding")
553
554 # Test with SparseVectorIndexConfig on #document
555 sparse_config = SparseVectorIndexConfig()
556 with pytest.raises(
557 ValueError, match="Cannot create index on special key '#document'"
558 ):
559 schema.create_index(config=sparse_config, key="#document")
560
561 # Test with SparseVectorIndexConfig on #embedding
562 with pytest.raises(
563 ValueError, match="Cannot create index on special key '#embedding'"
564 ):
565 schema.create_index(config=sparse_config, key="#embedding")
566
567 def test_cannot_enable_all_indexes_for_key(self) -> None:
568 """Test that enabling all indexes for a key is not allowed."""
569 schema = Schema()
570
571 # Try to enable all indexes for a custom key (config=None, key="my_key")
572 with pytest.raises(
573 ValueError, match="Cannot enable all index types for key 'my_key'"
574 ):
575 schema.create_index(key="my_key")
576
577 # Try to disable all indexes for a custom key (config=None, key="my_key")
578 with pytest.raises(
579 ValueError, match="Cannot disable all index types for key 'my_key'"
580 ):
581 schema.delete_index(key="my_key")
582
583 def test_cannot_delete_vector_or_fts_index(self) -> None:
584 """Test that deleting vector index is not allowed and FTS delete is restricted."""
585 schema = Schema()
586
587 # Vector delete - fully disallowed
588 vector_config = VectorIndexConfig()
589 with pytest.raises(
590 ValueError, match="Deleting vector index is not currently supported"
591 ):
592 schema.delete_index(config=vector_config)
593
594 # Try to delete vector index on a custom key
595 with pytest.raises(
596 ValueError, match="Deleting vector index is not currently supported"
597 ):
598 schema.delete_index(config=vector_config, key="my_vectors")
599
600 # FTS delete: only allowed on #document, other cases must error
601 fts_config = FtsIndexConfig()
602 with pytest.raises(
603 ValueError, match="Deleting FTS index is only supported on #document key"
604 ):
605 schema.delete_index(config=fts_config)
606
607 # Try to delete FTS index on a custom key
608 with pytest.raises(
609 ValueError, match="Deleting FTS index is only supported on #document key"
610 ):
611 schema.delete_index(config=fts_config, key="my_text")
612
613 # Positive case: deleting FTS index on #document should succeed and disable FTS
614 schema.delete_index(config=fts_config, key="#document")
615 assert schema.keys["#document"].string is not None
616 assert schema.keys["#document"].string.fts_index is not None
617 assert schema.keys["#document"].string.fts_index.enabled is False
618 assert schema.keys["#document"].string.fts_index.config == fts_config
619
620 def test_disable_string_inverted_index_globally(self) -> None:
621 """Test disabling string inverted index globally."""
622 schema = Schema()
623
624 # Verify string inverted index is enabled by default in global defaults
625 assert schema.defaults.string is not None
626 assert schema.defaults.string.string_inverted_index is not None
627 assert schema.defaults.string.string_inverted_index.enabled is True
628
629 # Disable string inverted index globally
630 string_config = StringInvertedIndexConfig()
631 result = schema.delete_index(config=string_config)
632 assert result is schema # Should return self for chaining
633
634 # Verify it's now disabled in defaults
635 assert schema.defaults.string.string_inverted_index is not None
636 assert schema.defaults.string.string_inverted_index.enabled is False
637 assert schema.defaults.string.string_inverted_index.config == string_config
638
639 # Verify key overrides are not affected (e.g., #document still has its config)
640 assert schema.keys["#document"].string is not None
641 assert schema.keys["#document"].string.string_inverted_index is not None
642 assert (
643 schema.keys["#document"].string.string_inverted_index.enabled is False
644 ) # #document has it disabled
645
646 def test_disable_string_inverted_index_on_key(self) -> None:
647 """Test disabling string inverted index on a specific key."""
648 schema = Schema()
649
650 # Disable string inverted index on a custom key
651 string_config = StringInvertedIndexConfig()
652 result = schema.delete_index(config=string_config, key="my_text_field")
653 assert result is schema
654
655 # Verify it's disabled on the custom key
656 assert "my_text_field" in schema.keys
657 assert schema.keys["my_text_field"].string is not None
658 assert schema.keys["my_text_field"].string.string_inverted_index is not None
659 assert (
660 schema.keys["my_text_field"].string.string_inverted_index.enabled is False
661 )
662 assert (
663 schema.keys["my_text_field"].string.string_inverted_index.config
664 == string_config
665 )
666
667 # Verify other value types on this key are None (sparse override)
668 assert schema.keys["my_text_field"].float_list is None
669 assert schema.keys["my_text_field"].sparse_vector is None
670 assert schema.keys["my_text_field"].int_value is None
671
672 # Verify global defaults are not affected
673 assert schema.defaults.string is not None
674 assert schema.defaults.string.string_inverted_index is not None
675 assert schema.defaults.string.string_inverted_index.enabled is True
676
677 # Verify other key overrides are not affected
678 assert schema.keys["#document"].string is not None
679 assert schema.keys["#document"].string.string_inverted_index is not None
680 assert schema.keys["#document"].string.string_inverted_index.enabled is False
681 assert schema.keys["#embedding"].float_list is not None
682 assert schema.keys["#embedding"].float_list.vector_index is not None
683 assert schema.keys["#embedding"].float_list.vector_index.enabled is True
684
685 def test_disable_int_inverted_index(self) -> None:
686 """Test disabling int inverted index globally and on a specific key."""
687 schema = Schema()
688
689 # Verify int inverted index is enabled by default
690 assert schema.defaults.int_value is not None
691 assert schema.defaults.int_value.int_inverted_index is not None
692 assert schema.defaults.int_value.int_inverted_index.enabled is True
693
694 # Test 1: Disable int inverted index globally
695 int_config = IntInvertedIndexConfig()
696 result = schema.delete_index(config=int_config)
697 assert result is schema
698
699 # Verify it's now disabled in defaults
700 assert schema.defaults.int_value.int_inverted_index.enabled is False
701 assert schema.defaults.int_value.int_inverted_index.config == int_config
702
703 # Test 2: Disable int inverted index on a specific key
704 int_config2 = IntInvertedIndexConfig()
705 result = schema.delete_index(config=int_config2, key="age_field")
706 assert result is schema
707
708 # Verify it's disabled on the custom key
709 assert "age_field" in schema.keys
710 assert schema.keys["age_field"].int_value is not None
711 assert schema.keys["age_field"].int_value.int_inverted_index is not None
712 assert schema.keys["age_field"].int_value.int_inverted_index.enabled is False
713 assert (
714 schema.keys["age_field"].int_value.int_inverted_index.config == int_config2
715 )
716
717 # Verify sparse override (only int_value is set)
718 assert schema.keys["age_field"].string is None
719 assert schema.keys["age_field"].float_list is None
720 assert schema.keys["age_field"].sparse_vector is None
721 assert schema.keys["age_field"].float_value is None
722 assert schema.keys["age_field"].boolean is None
723
724 # Verify other keys are not affected
725 assert schema.keys["#document"].string is not None
726 assert schema.keys["#embedding"].float_list is not None
727
728 def test_serialize_deserialize_default_schema(self) -> None:
729 """Test serialization and deserialization of a default Schema."""
730 # Create a default schema
731 original = Schema()
732
733 # Serialize to JSON
734 json_data = original.serialize_to_json()
735
736 # Verify the top-level structure
737 assert "defaults" in json_data
738 assert "keys" in json_data
739 assert isinstance(json_data["defaults"], dict)
740 assert isinstance(json_data["keys"], dict)
741
742 # Verify defaults structure in detail
743 defaults = json_data["defaults"]
744
745 # Check string
746 assert "string" in defaults
747 assert "fts_index" in defaults["string"]
748 assert defaults["string"]["fts_index"]["enabled"] is False
749 assert defaults["string"]["fts_index"]["config"] == {}
750 assert "string_inverted_index" in defaults["string"]
751 assert defaults["string"]["string_inverted_index"]["enabled"] is True
752 assert defaults["string"]["string_inverted_index"]["config"] == {}
753
754 # Check float_list
755 assert "float_list" in defaults
756 assert "vector_index" in defaults["float_list"]
757 assert defaults["float_list"]["vector_index"]["enabled"] is False
758 vector_config = defaults["float_list"]["vector_index"]["config"]
759 assert "space" in vector_config
760 assert vector_config["space"] == "l2" # Default space
761 assert "embedding_function" in vector_config
762 assert vector_config["embedding_function"]["type"] == "known"
763 assert vector_config["embedding_function"]["name"] == "default"
764 assert vector_config["embedding_function"]["config"] == {}
765
766 # Check sparse_vector
767 assert "sparse_vector" in defaults
768 assert "sparse_vector_index" in defaults["sparse_vector"]
769 assert defaults["sparse_vector"]["sparse_vector_index"]["enabled"] is False
770 sparse_vector_config = defaults["sparse_vector"]["sparse_vector_index"][
771 "config"
772 ]
773 # SparseVectorIndexConfig has embedding_function field with unknown default
774 assert "embedding_function" in sparse_vector_config
775 assert sparse_vector_config["embedding_function"] == {"type": "unknown"}
776
777 # Check int
778 assert "int" in defaults
779 assert "int_inverted_index" in defaults["int"]
780 assert defaults["int"]["int_inverted_index"]["enabled"] is True
781 assert defaults["int"]["int_inverted_index"]["config"] == {}
782
783 # Check float
784 assert "float" in defaults
785 assert "float_inverted_index" in defaults["float"]
786 assert defaults["float"]["float_inverted_index"]["enabled"] is True
787 assert defaults["float"]["float_inverted_index"]["config"] == {}
788
789 # Check bool
790 assert "bool" in defaults
791 assert "bool_inverted_index" in defaults["bool"]
792 assert defaults["bool"]["bool_inverted_index"]["enabled"] is True
793 assert defaults["bool"]["bool_inverted_index"]["config"] == {}
794
795 # Verify key overrides structure in detail
796 keys = json_data["keys"]
797
798 # Check #document
799 assert "#document" in keys
800 assert "string" in keys["#document"]
801 assert "fts_index" in keys["#document"]["string"]
802 assert keys["#document"]["string"]["fts_index"]["enabled"] is True
803 assert keys["#document"]["string"]["fts_index"]["config"] == {}
804 assert "string_inverted_index" in keys["#document"]["string"]
805 assert keys["#document"]["string"]["string_inverted_index"]["enabled"] is False
806 assert keys["#document"]["string"]["string_inverted_index"]["config"] == {}
807
808 # Check #embedding
809 assert "#embedding" in keys
810 assert "float_list" in keys["#embedding"]
811 assert "vector_index" in keys["#embedding"]["float_list"]
812 assert keys["#embedding"]["float_list"]["vector_index"]["enabled"] is True
813 embedding_vector_config = keys["#embedding"]["float_list"]["vector_index"][
814 "config"
815 ]
816 assert "space" in embedding_vector_config
817 assert embedding_vector_config["space"] == "l2" # Default space
818 assert "source_key" in embedding_vector_config
819 assert embedding_vector_config["source_key"] == "#document"
820 assert "embedding_function" in embedding_vector_config
821 assert embedding_vector_config["embedding_function"]["type"] == "known"
822 assert embedding_vector_config["embedding_function"]["name"] == "default"
823 assert embedding_vector_config["embedding_function"]["config"] == {}
824
825 # Deserialize back to Schema
826 deserialized = Schema.deserialize_from_json(json_data)
827
828 # Verify deserialized schema matches original - exhaustive validation
829 # Check defaults.string
830 assert deserialized.defaults.string is not None
831 assert deserialized.defaults.string.fts_index is not None
832 assert deserialized.defaults.string.fts_index.enabled is False
833 assert (
834 deserialized.defaults.string.fts_index.enabled
835 == original.defaults.string.fts_index.enabled
836 ) # type: ignore[union-attr]
837 assert deserialized.defaults.string.string_inverted_index is not None
838 assert deserialized.defaults.string.string_inverted_index.enabled is True
839 assert (
840 deserialized.defaults.string.string_inverted_index.enabled
841 == original.defaults.string.string_inverted_index.enabled
842 ) # type: ignore[union-attr]
843
844 # Check defaults.float_list (vector index)
845 assert deserialized.defaults.float_list is not None
846 assert deserialized.defaults.float_list.vector_index is not None
847 assert deserialized.defaults.float_list.vector_index.enabled is False
848 assert (
849 deserialized.defaults.float_list.vector_index.enabled
850 == original.defaults.float_list.vector_index.enabled
851 ) # type: ignore[union-attr]
852 # Space is resolved during serialization, so deserialized has explicit value
853 assert deserialized.defaults.float_list.vector_index.config.space == "l2"
854 # Check embedding function is preserved
855 assert (
856 deserialized.defaults.float_list.vector_index.config.embedding_function
857 is not None
858 )
859 assert (
860 deserialized.defaults.float_list.vector_index.config.embedding_function.name()
861 == "default"
862 )
863 assert (
864 original.defaults.float_list.vector_index.config.embedding_function.name()
865 == "default"
866 ) # type: ignore[union-attr]
867
868 # Check defaults.sparse_vector
869 assert deserialized.defaults.sparse_vector is not None
870 assert deserialized.defaults.sparse_vector.sparse_vector_index is not None
871 assert deserialized.defaults.sparse_vector.sparse_vector_index.enabled is False
872 assert (
873 deserialized.defaults.sparse_vector.sparse_vector_index.enabled
874 == original.defaults.sparse_vector.sparse_vector_index.enabled
875 ) # type: ignore[union-attr]
876
877 # Check defaults.int_value
878 assert deserialized.defaults.int_value is not None
879 assert deserialized.defaults.int_value.int_inverted_index is not None
880 assert deserialized.defaults.int_value.int_inverted_index.enabled is True
881 assert (
882 deserialized.defaults.int_value.int_inverted_index.enabled
883 == original.defaults.int_value.int_inverted_index.enabled
884 ) # type: ignore[union-attr]
885
886 # Check defaults.float_value
887 assert deserialized.defaults.float_value is not None
888 assert deserialized.defaults.float_value.float_inverted_index is not None
889 assert deserialized.defaults.float_value.float_inverted_index.enabled is True
890 assert (
891 deserialized.defaults.float_value.float_inverted_index.enabled
892 == original.defaults.float_value.float_inverted_index.enabled
893 ) # type: ignore[union-attr]
894
895 # Check defaults.boolean
896 assert deserialized.defaults.boolean is not None
897 assert deserialized.defaults.boolean.bool_inverted_index is not None
898 assert deserialized.defaults.boolean.bool_inverted_index.enabled is True
899 assert (
900 deserialized.defaults.boolean.bool_inverted_index.enabled
901 == original.defaults.boolean.bool_inverted_index.enabled
902 ) # type: ignore[union-attr]
903
904 # Check keys.#document
905 assert "#document" in deserialized.keys
906 assert deserialized.keys["#document"].string is not None
907 assert deserialized.keys["#document"].string.fts_index is not None
908 assert deserialized.keys["#document"].string.fts_index.enabled is True
909 assert (
910 deserialized.keys["#document"].string.fts_index.enabled
911 == original.keys["#document"].string.fts_index.enabled
912 ) # type: ignore[union-attr]
913 assert deserialized.keys["#document"].string.string_inverted_index is not None
914 assert (
915 deserialized.keys["#document"].string.string_inverted_index.enabled is False
916 )
917 assert (
918 deserialized.keys["#document"].string.string_inverted_index.enabled
919 == original.keys["#document"].string.string_inverted_index.enabled
920 ) # type: ignore[union-attr]
921
922 # Check keys.#embedding
923 assert "#embedding" in deserialized.keys
924 assert deserialized.keys["#embedding"].float_list is not None
925 assert deserialized.keys["#embedding"].float_list.vector_index is not None
926 assert deserialized.keys["#embedding"].float_list.vector_index.enabled is True
927 assert (
928 deserialized.keys["#embedding"].float_list.vector_index.enabled
929 == original.keys["#embedding"].float_list.vector_index.enabled
930 ) # type: ignore[union-attr]
931 # Verify source_key is preserved
932 assert (
933 deserialized.keys["#embedding"].float_list.vector_index.config.source_key
934 == "#document"
935 )
936 assert (
937 original.keys["#embedding"].float_list.vector_index.config.source_key
938 == "#document"
939 ) # type: ignore[union-attr]
940 # Verify space is preserved (resolved during serialization)
941 assert (
942 deserialized.keys["#embedding"].float_list.vector_index.config.space == "l2"
943 )
944 # Verify embedding function is preserved
945 assert (
946 deserialized.keys[
947 "#embedding"
948 ].float_list.vector_index.config.embedding_function
949 is not None
950 )
951 assert (
952 deserialized.keys[
953 "#embedding"
954 ].float_list.vector_index.config.embedding_function.name()
955 == "default"
956 )
957 assert (
958 original.keys[
959 "#embedding"
960 ].float_list.vector_index.config.embedding_function.name()
961 == "default"
962 ) # type: ignore[union-attr]
963
964 def test_serialize_deserialize_with_vector_config_no_ef(self) -> None:
965 """Test serialization/deserialization of Schema with vector config where embedding_function=None."""
966 # Create a default schema and modify vector config with ef=None
967 original = Schema()
968 vector_config = VectorIndexConfig(
969 space="cosine",
970 embedding_function=None, # Explicitly set to None
971 )
972 original.create_index(config=vector_config)
973
974 # Serialize to JSON
975 json_data = original.serialize_to_json()
976
977 # Verify defaults structure - vector index should reflect the changes
978 defaults = json_data["defaults"]
979 assert "float_list" in defaults
980 assert "vector_index" in defaults["float_list"]
981 vector_json = defaults["float_list"]["vector_index"]
982 assert vector_json["enabled"] is False # Still disabled in defaults
983 assert vector_json["config"]["space"] == "cosine" # User-specified space
984 # When ef=None, it should serialize as legacy
985 assert vector_json["config"]["embedding_function"]["type"] == "legacy"
986
987 # Verify #embedding also has the updated config
988 keys = json_data["keys"]
989 assert "#embedding" in keys
990 embedding_vector_json = keys["#embedding"]["float_list"]["vector_index"]
991 assert embedding_vector_json["enabled"] is True # Still enabled on #embedding
992 assert (
993 embedding_vector_json["config"]["space"] == "cosine"
994 ) # User-specified space
995 assert embedding_vector_json["config"]["source_key"] == "#document" # Preserved
996 # When ef=None, it should serialize as legacy
997 assert embedding_vector_json["config"]["embedding_function"]["type"] == "legacy"
998
999 # Deserialize back to Schema
1000 deserialized = Schema.deserialize_from_json(json_data)
1001
1002 # Verify deserialized schema has the correct values
1003 # Check defaults.float_list (vector index)
1004 assert deserialized.defaults.float_list is not None
1005 assert deserialized.defaults.float_list.vector_index is not None
1006 assert deserialized.defaults.float_list.vector_index.enabled is False
1007 assert (
1008 deserialized.defaults.float_list.vector_index.config.space == "cosine"
1009 ) # User space preserved
1010 # ef=None should deserialize as None (legacy)
1011 assert (
1012 deserialized.defaults.float_list.vector_index.config.embedding_function
1013 is None
1014 )
1015
1016 # Check #embedding vector index
1017 assert "#embedding" in deserialized.keys
1018 assert deserialized.keys["#embedding"].float_list is not None
1019 assert deserialized.keys["#embedding"].float_list.vector_index is not None
1020 assert deserialized.keys["#embedding"].float_list.vector_index.enabled is True
1021 assert (
1022 deserialized.keys["#embedding"].float_list.vector_index.config.space
1023 == "cosine"
1024 ) # User space preserved
1025 assert (
1026 deserialized.keys["#embedding"].float_list.vector_index.config.source_key
1027 == "#document"
1028 ) # Preserved
1029 # ef=None should deserialize as None (legacy)
1030 assert (
1031 deserialized.keys[
1032 "#embedding"
1033 ].float_list.vector_index.config.embedding_function
1034 is None
1035 )
1036
1037 def test_serialize_deserialize_with_custom_ef(self) -> None:
1038 """Test serialization/deserialization of Schema with custom embedding function."""
1039 # Register the mock embedding function so it can be deserialized
1040 from chromadb.utils.embedding_functions import known_embedding_functions
1041
1042 known_embedding_functions["mock_embedding"] = MockEmbeddingFunction
1043
1044 try:
1045 # Create a default schema and modify vector config with custom EF
1046 original = Schema()
1047 custom_ef = MockEmbeddingFunction(model_name="custom_model_v3")
1048 hnsw_config = HnswIndexConfig(
1049 ef_construction=256, max_neighbors=48, ef_search=128
1050 )
1051 vector_config = VectorIndexConfig(
1052 embedding_function=custom_ef,
1053 space="ip", # Inner product
1054 hnsw=hnsw_config,
1055 )
1056 original.create_index(config=vector_config)
1057
1058 # Serialize to JSON
1059 json_data = original.serialize_to_json()
1060
1061 # Verify defaults structure - vector index should reflect the changes
1062 defaults = json_data["defaults"]
1063 assert "float_list" in defaults
1064 assert "vector_index" in defaults["float_list"]
1065 vector_json = defaults["float_list"]["vector_index"]
1066 assert vector_json["enabled"] is False # Still disabled in defaults
1067 assert vector_json["config"]["space"] == "ip" # User-specified space
1068 # Custom EF should serialize as known type
1069 assert vector_json["config"]["embedding_function"]["type"] == "known"
1070 assert (
1071 vector_json["config"]["embedding_function"]["name"] == "mock_embedding"
1072 )
1073 assert (
1074 vector_json["config"]["embedding_function"]["config"]["model_name"]
1075 == "custom_model_v3"
1076 )
1077 # HNSW config should be present
1078 assert "hnsw" in vector_json["config"]
1079 assert vector_json["config"]["hnsw"]["ef_construction"] == 256
1080 assert vector_json["config"]["hnsw"]["max_neighbors"] == 48
1081 assert vector_json["config"]["hnsw"]["ef_search"] == 128
1082
1083 # Verify #embedding also has the updated config
1084 keys = json_data["keys"]
1085 assert "#embedding" in keys
1086 embedding_vector_json = keys["#embedding"]["float_list"]["vector_index"]
1087 assert (
1088 embedding_vector_json["enabled"] is True
1089 ) # Still enabled on #embedding
1090 assert (
1091 embedding_vector_json["config"]["space"] == "ip"
1092 ) # User-specified space
1093 assert (
1094 embedding_vector_json["config"]["source_key"] == "#document"
1095 ) # Preserved
1096 # Custom EF should serialize as known type
1097 assert (
1098 embedding_vector_json["config"]["embedding_function"]["type"] == "known"
1099 )
1100 assert (
1101 embedding_vector_json["config"]["embedding_function"]["name"]
1102 == "mock_embedding"
1103 )
1104 assert (
1105 embedding_vector_json["config"]["embedding_function"]["config"][
1106 "model_name"
1107 ]
1108 == "custom_model_v3"
1109 )
1110 # HNSW config should be present
1111 assert "hnsw" in embedding_vector_json["config"]
1112 assert embedding_vector_json["config"]["hnsw"]["ef_construction"] == 256
1113 assert embedding_vector_json["config"]["hnsw"]["max_neighbors"] == 48
1114 assert embedding_vector_json["config"]["hnsw"]["ef_search"] == 128
1115
1116 # Deserialize back to Schema
1117 deserialized = Schema.deserialize_from_json(json_data)
1118
1119 # Verify deserialized schema has the correct values
1120 # Check defaults.float_list (vector index)
1121 assert deserialized.defaults.float_list is not None
1122 assert deserialized.defaults.float_list.vector_index is not None
1123 assert deserialized.defaults.float_list.vector_index.enabled is False
1124 assert (
1125 deserialized.defaults.float_list.vector_index.config.space == "ip"
1126 ) # User space preserved
1127 # Custom EF should be reconstructed
1128 assert (
1129 deserialized.defaults.float_list.vector_index.config.embedding_function
1130 is not None
1131 )
1132 assert (
1133 deserialized.defaults.float_list.vector_index.config.embedding_function.name()
1134 == "mock_embedding"
1135 )
1136 # Verify the EF config is correct
1137 ef_config = deserialized.defaults.float_list.vector_index.config.embedding_function.get_config()
1138 assert ef_config["model_name"] == "custom_model_v3"
1139 # HNSW config should be preserved
1140 assert deserialized.defaults.float_list.vector_index.config.hnsw is not None
1141 assert (
1142 deserialized.defaults.float_list.vector_index.config.hnsw.ef_construction
1143 == 256
1144 )
1145 assert (
1146 deserialized.defaults.float_list.vector_index.config.hnsw.max_neighbors
1147 == 48
1148 )
1149 assert (
1150 deserialized.defaults.float_list.vector_index.config.hnsw.ef_search
1151 == 128
1152 )
1153
1154 # Check #embedding vector index
1155 assert "#embedding" in deserialized.keys
1156 assert deserialized.keys["#embedding"].float_list is not None
1157 assert deserialized.keys["#embedding"].float_list.vector_index is not None
1158 assert (
1159 deserialized.keys["#embedding"].float_list.vector_index.enabled is True
1160 )
1161 assert (
1162 deserialized.keys["#embedding"].float_list.vector_index.config.space
1163 == "ip"
1164 ) # User space preserved
1165 assert (
1166 deserialized.keys[
1167 "#embedding"
1168 ].float_list.vector_index.config.source_key
1169 == "#document"
1170 ) # Preserved
1171 # Custom EF should be reconstructed
1172 assert (
1173 deserialized.keys[
1174 "#embedding"
1175 ].float_list.vector_index.config.embedding_function
1176 is not None
1177 )
1178 assert (
1179 deserialized.keys[
1180 "#embedding"
1181 ].float_list.vector_index.config.embedding_function.name()
1182 == "mock_embedding"
1183 )
1184 # Verify the EF config is correct
1185 ef_config_embedding = deserialized.keys[
1186 "#embedding"
1187 ].float_list.vector_index.config.embedding_function.get_config()
1188 assert ef_config_embedding["model_name"] == "custom_model_v3"
1189 # HNSW config should be preserved
1190 assert (
1191 deserialized.keys["#embedding"].float_list.vector_index.config.hnsw
1192 is not None
1193 )
1194 assert (
1195 deserialized.keys[
1196 "#embedding"
1197 ].float_list.vector_index.config.hnsw.ef_construction
1198 == 256
1199 )
1200 assert (