codekingpro/portable-devtools
114k
1"""
2Integration test for Chroma's Task API
3
4Tests the task creation, execution, and removal functionality
5for automatically processing collections.
6"""
7
8import pytest
9from chromadb.api.client import Client as ClientCreator
10from chromadb.api.functions import (
11 RECORD_COUNTER_FUNCTION,
12 STATISTICS_FUNCTION,
13 Function,
14)
15from chromadb.config import System
16from chromadb.errors import ChromaError, NotFoundError
17from chromadb.test.conftest import skip_if_not_cluster
18from chromadb.test.utils.wait_for_version_increase import (
19 get_collection_version,
20 wait_for_version_increase,
21)
22from time import sleep
23
24pytestmark = [skip_if_not_cluster()]
25
26
27def test_count_function_attach_and_detach(basic_http_client: System) -> None:
28 """Test creating and removing a function with the record_counter operator"""
29 client = ClientCreator.from_system(basic_http_client)
30 client.reset()
31
32 # Create a collection
33 collection = client.get_or_create_collection(
34 name="my_document",
35 metadata={"description": "Sample documents for task processing"},
36 )
37
38 # Create a task that counts records in the collection
39 attached_fn, created = collection.attach_function(
40 name="count_my_docs",
41 function=RECORD_COUNTER_FUNCTION,
42 output_collection="my_documents_counts",
43 params=None,
44 )
45
46 # Verify task creation succeeded
47 assert attached_fn is not None
48 assert created is True
49 initial_version = get_collection_version(client, collection.name)
50
51 # Add documents
52 collection.add(
53 ids=["doc_{}".format(i) for i in range(0, 300)],
54 documents=["test document"] * 300,
55 )
56
57 # Verify documents were added
58 assert collection.count() == 300
59
60 wait_for_version_increase(client, collection.name, initial_version)
61 # Give some time to invalidate the frontend query cache
62 sleep(60)
63
64 result = client.get_collection("my_documents_counts").get("function_output")
65 assert result["metadatas"] is not None
66 assert result["metadatas"][0]["total_count"] == 300
67
68 # Remove the task
69 success = collection.detach_function(
70 attached_fn.name,
71 delete_output_collection=True,
72 )
73
74 # Verify task removal succeeded
75 assert success is True
76
77
78def test_task_with_invalid_function(basic_http_client: System) -> None:
79 """Test that creating a task with an invalid function raises an error"""
80 client = ClientCreator.from_system(basic_http_client)
81 client.reset()
82
83 collection = client.get_or_create_collection(name="test_invalid_function")
84 collection.add(ids=["id1"], documents=["test document"])
85
86 # Attempt to create task with non-existent function should raise ChromaError
87 with pytest.raises(ChromaError, match="function not found"):
88 collection.attach_function(
89 function=Function._NONEXISTENT_TEST_ONLY,
90 name="invalid_task",
91 output_collection="output_collection",
92 params=None,
93 )
94
95
96def test_attach_function_returns_function_name(basic_http_client: System) -> None:
97 """Test that attach_function and get_attached_function return function_name field instead of UUID"""
98 client = ClientCreator.from_system(basic_http_client)
99 client.reset()
100
101 collection = client.create_collection(name="test_function_name")
102 collection.add(ids=["id1"], documents=["doc1"])
103
104 # Attach a function and verify function_name field in response
105 attached_fn, created = collection.attach_function(
106 function=RECORD_COUNTER_FUNCTION,
107 name="my_counter",
108 output_collection="output_collection",
109 params=None,
110 )
111
112 # Verify the attached function has function_name (not function_id UUID)
113 assert created is True
114 assert attached_fn.function_name == "record_counter"
115 assert attached_fn.name == "my_counter"
116
117 # Get the attached function and verify function_name field is also present
118 retrieved_fn = collection.get_attached_function("my_counter")
119 assert retrieved_fn == attached_fn
120
121 # Clean up
122 collection.detach_function(attached_fn.name, delete_output_collection=True)
123
124
125def test_function_multiple_collections(basic_http_client: System) -> None:
126 """Test attaching functions on multiple collections"""
127 client = ClientCreator.from_system(basic_http_client)
128 client.reset()
129
130 # Create first collection and task
131 collection1 = client.create_collection(name="collection_1")
132 collection1.add(ids=["id1", "id2"], documents=["doc1", "doc2"])
133
134 attached_fn1, created1 = collection1.attach_function(
135 function=RECORD_COUNTER_FUNCTION,
136 name="task_1",
137 output_collection="output_1",
138 params=None,
139 )
140
141 assert attached_fn1 is not None
142 assert created1 is True
143
144 # Create second collection and task
145 collection2 = client.create_collection(name="collection_2")
146 collection2.add(ids=["id3", "id4"], documents=["doc3", "doc4"])
147
148 attached_fn2, created2 = collection2.attach_function(
149 function=RECORD_COUNTER_FUNCTION,
150 name="task_2",
151 output_collection="output_2",
152 params=None,
153 )
154
155 assert attached_fn2 is not None
156 assert created2 is True
157
158 # Task IDs should be different
159 assert attached_fn1.id != attached_fn2.id
160
161 # Clean up
162 assert (
163 collection1.detach_function(attached_fn1.name, delete_output_collection=True)
164 is True
165 )
166 assert (
167 collection2.detach_function(attached_fn2.name, delete_output_collection=True)
168 is True
169 )
170
171
172def test_functions_one_attached_function_per_collection(
173 basic_http_client: System,
174) -> None:
175 """Test that only one attached function is allowed per collection"""
176 client = ClientCreator.from_system(basic_http_client)
177 client.reset()
178
179 # Create a single collection
180 collection = client.create_collection(name="single_task_collection")
181 collection.add(ids=["id1", "id2", "id3"], documents=["doc1", "doc2", "doc3"])
182
183 # Create first task on the collection
184 attached_fn1, created = collection.attach_function(
185 function=RECORD_COUNTER_FUNCTION,
186 name="task_1",
187 output_collection="output_1",
188 params=None,
189 )
190
191 assert attached_fn1 is not None
192 assert created is True
193
194 # Attempt to create a second task with a different name should fail
195 # (only one attached function allowed per collection)
196 with pytest.raises(
197 ChromaError,
198 match="collection already has an attached function: name=task_1, function=record_counter, output_collection=output_1",
199 ):
200 collection.attach_function(
201 function=RECORD_COUNTER_FUNCTION,
202 name="task_2",
203 output_collection="output_2",
204 params=None,
205 )
206
207 # Attempt to create a task with the same name but different function_id should also fail
208 with pytest.raises(
209 ChromaError,
210 match=r"collection already has an attached function: name=task_1, function=record_counter, output_collection=output_1",
211 ):
212 collection.attach_function(
213 function=STATISTICS_FUNCTION,
214 name="task_1",
215 output_collection="output_different", # Different output collection
216 params=None,
217 )
218
219 # Detach the first function
220 assert (
221 collection.detach_function(attached_fn1.name, delete_output_collection=True)
222 is True
223 )
224
225 # Now we should be able to attach a new function
226 attached_fn2, created2 = collection.attach_function(
227 function=RECORD_COUNTER_FUNCTION,
228 name="task_2",
229 output_collection="output_2",
230 params=None,
231 )
232
233 assert attached_fn2 is not None
234 assert created2 is True
235 assert attached_fn2.id != attached_fn1.id
236
237 # Clean up
238 assert (
239 collection.detach_function(attached_fn2.name, delete_output_collection=True)
240 is True
241 )
242
243
244def test_attach_function_with_invalid_params(basic_http_client: System) -> None:
245 """Test that attach_function with non-empty params raises an error"""
246 client = ClientCreator.from_system(basic_http_client)
247 client.reset()
248
249 collection = client.create_collection(name="test_invalid_params")
250 collection.add(ids=["id1"], documents=["test document"])
251
252 # Attempt to create task with non-empty params should fail
253 # (no functions currently accept parameters)
254 with pytest.raises(
255 ChromaError,
256 match="params must be empty - no functions currently accept parameters",
257 ):
258 collection.attach_function(
259 name="invalid_params_task",
260 function=RECORD_COUNTER_FUNCTION,
261 output_collection="output_collection",
262 params={"some_key": "some_value"},
263 )
264
265
266def test_attach_function_output_collection_already_exists(
267 basic_http_client: System,
268) -> None:
269 """Test that attach_function fails when output collection name already exists"""
270 client = ClientCreator.from_system(basic_http_client)
271 client.reset()
272
273 # Create a collection that will be used as input
274 input_collection = client.create_collection(name="input_collection")
275 input_collection.add(ids=["id1"], documents=["test document"])
276
277 # Create another collection with the name we want to use for output
278 client.create_collection(name="existing_output_collection")
279
280 # Attempt to create task with output collection name that already exists
281 with pytest.raises(
282 ChromaError,
283 match=r"Output collection \[existing_output_collection\] already exists",
284 ):
285 input_collection.attach_function(
286 name="my_task",
287 function=RECORD_COUNTER_FUNCTION,
288 output_collection="existing_output_collection",
289 params=None,
290 )
291
292
293def test_function_remove_nonexistent(basic_http_client: System) -> None:
294 """Test removing a task that doesn't exist raises NotFoundError"""
295 client = ClientCreator.from_system(basic_http_client)
296 client.reset()
297
298 collection = client.create_collection(name="test_collection")
299 collection.add(ids=["id1"], documents=["test"])
300 attached_fn, _ = collection.attach_function(
301 function=RECORD_COUNTER_FUNCTION,
302 name="test_function",
303 output_collection="output_collection",
304 params=None,
305 )
306
307 collection.detach_function(attached_fn.name, delete_output_collection=True)
308
309 # Trying to detach this function again should raise NotFoundError
310 with pytest.raises(NotFoundError, match="does not exist"):
311 collection.detach_function(attached_fn.name, delete_output_collection=True)
312
313
314def test_attach_to_output_collection_fails(basic_http_client: System) -> None:
315 """Test that attaching a function to an output collection fails"""
316 client = ClientCreator.from_system(basic_http_client)
317 client.reset()
318
319 # Create input collection
320 input_collection = client.create_collection(name="input_collection")
321 input_collection.add(ids=["id1"], documents=["test"])
322
323 _, _ = input_collection.attach_function(
324 name="test_function",
325 function=RECORD_COUNTER_FUNCTION,
326 output_collection="output_collection",
327 params=None,
328 )
329 output_collection = client.get_collection(name="output_collection")
330
331 with pytest.raises(
332 ChromaError, match="cannot attach function to an output collection"
333 ):
334 _ = output_collection.attach_function(
335 name="test_function_2",
336 function=RECORD_COUNTER_FUNCTION,
337 output_collection="output_collection_2",
338 params=None,
339 )
340
341
342def test_delete_output_collection_detaches_function(basic_http_client: System) -> None:
343 """Test that deleting an output collection also detaches the attached function"""
344 client = ClientCreator.from_system(basic_http_client)
345 client.reset()
346
347 # Create input collection and attach a function
348 input_collection = client.create_collection(name="input_collection")
349 input_collection.add(ids=["id1"], documents=["test"])
350
351 attached_fn, created = input_collection.attach_function(
352 name="my_function",
353 function=RECORD_COUNTER_FUNCTION,
354 output_collection="output_collection",
355 params=None,
356 )
357 assert attached_fn is not None
358 assert created is True
359
360 # Delete the output collection directly
361 client.delete_collection("output_collection")
362
363 # The attached function should now be gone - trying to get it should raise NotFoundError
364 with pytest.raises(NotFoundError):
365 input_collection.get_attached_function("my_function")
366
367
368def test_delete_orphaned_output_collection(basic_http_client: System) -> None:
369 """Test that deleting an output collection from a recently detached function works"""
370 client = ClientCreator.from_system(basic_http_client)
371 client.reset()
372
373 # Create input collection and attach a function
374 input_collection = client.create_collection(name="input_collection")
375 input_collection.add(ids=["id1"], documents=["test"])
376
377 attached_fn, created = input_collection.attach_function(
378 name="my_function",
379 function=RECORD_COUNTER_FUNCTION,
380 output_collection="output_collection",
381 params=None,
382 )
383 assert attached_fn is not None
384 assert created is True
385
386 input_collection.detach_function(attached_fn.name, delete_output_collection=False)
387
388 # Delete the output collection directly
389 client.delete_collection("output_collection")
390
391 # The attached function should still exist but be marked as detached
392 with pytest.raises(NotFoundError):
393 input_collection.get_attached_function("my_function")
394
395 with pytest.raises(NotFoundError):
396 # Try to use the function - it should fail since it's detached
397 client.get_collection("output_collection")
398
399
400def test_partial_attach_function_repair(
401 basic_http_client: System,
402) -> None:
403 """Test creating and removing a function with the record_counter operator"""
404 client = ClientCreator.from_system(basic_http_client)
405 client.reset()
406
407 # Create a collection
408 collection = client.get_or_create_collection(
409 name="my_document",
410 )
411
412 # Create a task that counts records in the collection
413 attached_fn, created = collection.attach_function(
414 name="count_my_docs",
415 function=RECORD_COUNTER_FUNCTION,
416 output_collection="my_documents_counts",
417 params=None,
418 )
419 assert created is True
420
421 # Verify task creation succeeded
422 assert attached_fn is not None
423
424 collection2 = client.get_or_create_collection(
425 name="my_document2",
426 )
427
428 # Create a task that counts records in the collection
429 # This should fail
430 with pytest.raises(
431 ChromaError, match=r"Output collection \[my_documents_counts\] already exists"
432 ):
433 attached_fn, _ = collection2.attach_function(
434 name="count_my_docs",
435 function=RECORD_COUNTER_FUNCTION,
436 output_collection="my_documents_counts",
437 params=None,
438 )
439
440 # Detach the function
441 assert (
442 collection.detach_function(attached_fn.name, delete_output_collection=True)
443 is True
444 )
445
446 # Create a task that counts records in the collection
447 attached_fn, created = collection2.attach_function(
448 name="count_my_docs",
449 function=RECORD_COUNTER_FUNCTION,
450 output_collection="my_documents_counts",
451 params=None,
452 )
453 assert attached_fn is not None
454 assert created is True
455
456
457def test_output_collection_created_with_schema(basic_http_client: System) -> None:
458 """Test that output collections are created with the source_attached_function_id in the schema"""
459 client = ClientCreator.from_system(basic_http_client)
460 client.reset()
461
462 # Create input collection and attach a function
463 input_collection = client.create_collection(name="input_collection")
464 input_collection.add(ids=["id1"], documents=["test"])
465
466 attached_fn, created = input_collection.attach_function(
467 name="my_function",
468 function=RECORD_COUNTER_FUNCTION,
469 output_collection="output_collection",
470 params=None,
471 )
472 assert attached_fn is not None
473 assert created is True
474
475 # Get the output collection - it should exist
476 output_collection = client.get_collection(name="output_collection")
477 assert output_collection is not None
478
479 # The source_attached_function_id is stored in the schema (not metadata)
480 # We can't directly access the schema from the client, but we verify the collection exists
481 # and the attached function orchestrator will use this field internally
482 assert "source_attached_function_id" in output_collection._model.pretty_schema()
483
484 # Clean up
485 input_collection.detach_function(attached_fn.name, delete_output_collection=True)
486
487
488def test_count_function_attach_and_detach_attach_attach(
489 basic_http_client: System,
490) -> None:
491 """Test creating and removing a function with the record_counter operator"""
492 client = ClientCreator.from_system(basic_http_client)
493 client.reset()
494
495 # Create a collection
496 collection = client.get_or_create_collection(
497 name="my_document",
498 metadata={"description": "Sample documents for task processing"},
499 )
500
501 # Create a task that counts records in the collection
502 attached_fn, created = collection.attach_function(
503 name="count_my_docs",
504 function=RECORD_COUNTER_FUNCTION,
505 output_collection="my_documents_counts",
506 params=None,
507 )
508
509 # Verify task creation succeeded
510 assert created is True
511 assert attached_fn is not None
512 initial_version = get_collection_version(client, collection.name)
513
514 # Add documents
515 collection.add(
516 ids=["doc_{}".format(i) for i in range(0, 300)],
517 documents=["test document"] * 300,
518 )
519
520 # Verify documents were added
521 assert collection.count() == 300
522
523 wait_for_version_increase(client, collection.name, initial_version)
524 # Give some time to invalidate the frontend query cache
525 sleep(60)
526
527 result = client.get_collection("my_documents_counts").get("function_output")
528 assert result["metadatas"] is not None
529 assert result["metadatas"][0]["total_count"] == 300
530
531 # Remove the task
532 success = collection.detach_function(
533 attached_fn.name, delete_output_collection=True
534 )
535
536 # Verify task removal succeeded
537 assert success is True
538
539 # Attach a function that counts records in the collection
540 attached_fn, created = collection.attach_function(
541 name="count_my_docs",
542 function=RECORD_COUNTER_FUNCTION,
543 output_collection="my_documents_counts",
544 params=None,
545 )
546 assert attached_fn is not None
547 assert created is True
548
549 # Attach a function that counts records in the collection
550 attached_fn, created = collection.attach_function(
551 name="count_my_docs",
552 function=RECORD_COUNTER_FUNCTION,
553 output_collection="my_documents_counts",
554 params=None,
555 )
556 assert created is False
557 assert attached_fn is not None
558
559
560def test_attach_function_idempotency(basic_http_client: System) -> None:
561 """Test that attach_function is idempotent - calling it twice with same params returns created=False"""
562 client = ClientCreator.from_system(basic_http_client)
563 client.reset()
564
565 collection = client.create_collection(name="idempotency_test")
566 collection.add(ids=["id1"], documents=["test document"])
567
568 # First attach - should be newly created
569 attached_fn1, created1 = collection.attach_function(
570 name="my_function",
571 function=RECORD_COUNTER_FUNCTION,
572 output_collection="output_collection",
573 params=None,
574 )
575 assert attached_fn1 is not None
576 assert created1 is True
577
578 # Second attach with identical params - should be idempotent (created=False)
579 attached_fn2, created2 = collection.attach_function(
580 name="my_function",
581 function=RECORD_COUNTER_FUNCTION,
582 output_collection="output_collection",
583 params=None,
584 )
585 assert attached_fn2 is not None
586 assert created2 is False
587
588 # Both should return the same function ID
589 assert attached_fn1.id == attached_fn2.id
590
591 # Clean up
592 collection.detach_function(attached_fn1.name, delete_output_collection=True)
593