diff --git a/chromadb/api/__init__.py b/chromadb/api/__init__.py index ae39b7b13f1..b68d7a27d10 100644 --- a/chromadb/api/__init__.py +++ b/chromadb/api/__init__.py @@ -90,7 +90,8 @@ def create_collection( @abstractmethod def get_collection( self, - name: str, + name: Optional[str] = None, + id: Optional[UUID] = None, embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(), ) -> Collection: """Get a collection with the given name. @@ -496,7 +497,8 @@ def create_collection( @override def get_collection( self, - name: str, + name: Optional[str] = None, + id: Optional[UUID] = None, embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(), tenant: str = DEFAULT_TENANT, database: str = DEFAULT_DATABASE, diff --git a/chromadb/api/client.py b/chromadb/api/client.py index 6e436cf499f..18067eafbd8 100644 --- a/chromadb/api/client.py +++ b/chromadb/api/client.py @@ -188,10 +188,12 @@ def create_collection( @override def get_collection( self, - name: str, + name: Optional[str] = None, + id: Optional[UUID] = None, embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(), ) -> Collection: return self._server.get_collection( + id=id, name=name, embedding_function=embedding_function, tenant=self.tenant, diff --git a/chromadb/api/fastapi.py b/chromadb/api/fastapi.py index fc5298c0720..38af7e52a91 100644 --- a/chromadb/api/fastapi.py +++ b/chromadb/api/fastapi.py @@ -250,15 +250,21 @@ def create_collection( @override def get_collection( self, - name: str, + name: Optional[str] = None, + id: Optional[UUID] = None, embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(), tenant: str = DEFAULT_TENANT, database: str = DEFAULT_DATABASE, ) -> Collection: """Returns a collection""" + if (name is None and id is None) or (name is not None and id is not None): + raise ValueError("Name or id must be specified, but not both") + + _params = {"tenant": tenant, "database": database} + if id is not None: + _params["type"] = str(id) resp = self._session.get( - self._api_url + "/collections/" + name, - params={"tenant": tenant, "database": database}, + self._api_url + "/collections/" + name if name else str(id), params=_params ) raise_chroma_error(resp) resp_json = resp.json() diff --git a/chromadb/api/models/Collection.py b/chromadb/api/models/Collection.py index f605d9d9d84..ef7c66139d2 100644 --- a/chromadb/api/models/Collection.py +++ b/chromadb/api/models/Collection.py @@ -40,6 +40,8 @@ class Collection(BaseModel): name: str id: UUID metadata: Optional[CollectionMetadata] = None + tenant: Optional[str] = None + database: Optional[str] = None _client: "ServerAPI" = PrivateAttr() _embedding_function: Optional[EmbeddingFunction] = PrivateAttr() @@ -49,9 +51,13 @@ def __init__( name: str, id: UUID, embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(), + tenant: Optional[str] = None, + database: Optional[str] = None, metadata: Optional[CollectionMetadata] = None, ): - super().__init__(name=name, metadata=metadata, id=id) + super().__init__( + name=name, metadata=metadata, id=id, tenant=tenant, database=database + ) self._client = client self._embedding_function = embedding_function diff --git a/chromadb/api/segment.py b/chromadb/api/segment.py index a411a125fce..eaaee31606d 100644 --- a/chromadb/api/segment.py +++ b/chromadb/api/segment.py @@ -184,6 +184,8 @@ def create_collection( name=name, metadata=coll["metadata"], # type: ignore embedding_function=embedding_function, + tenant=tenant, + database=database, ) @trace_method( @@ -214,13 +216,16 @@ def get_or_create_collection( @override def get_collection( self, - name: str, + name: Optional[str] = None, + id: Optional[UUID] = None, embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(), tenant: str = DEFAULT_TENANT, database: str = DEFAULT_DATABASE, ) -> Collection: + if id is None and name is None or (id is not None and name is not None): + raise ValueError("Name or id must be specified, but not both") existing = self._sysdb.get_collections( - name=name, tenant=tenant, database=database + id=id, name=name, tenant=tenant, database=database ) if existing: @@ -230,6 +235,8 @@ def get_collection( name=existing[0]["name"], metadata=existing[0]["metadata"], # type: ignore embedding_function=embedding_function, + tenant=tenant, + database=database, ) else: raise ValueError(f"Collection {name} does not exist.") @@ -250,6 +257,8 @@ def list_collections( id=db_collection["id"], name=db_collection["name"], metadata=db_collection["metadata"], # type: ignore + tenant=db_collection["tenant"], + database=db_collection["database"], ) ) return collections @@ -486,7 +495,9 @@ def _get( embeddings=[r["embedding"] for r in vectors] if "embeddings" in include else None, - metadatas=_clean_metadatas(metadatas) if "metadatas" in include else None, # type: ignore + metadatas=_clean_metadatas(metadatas) + if "metadatas" in include + else None, # type: ignore documents=documents if "documents" in include else None, # type: ignore ) diff --git a/chromadb/db/impl/grpc/server.py b/chromadb/db/impl/grpc/server.py index 1a71929214e..257aa80f0e7 100644 --- a/chromadb/db/impl/grpc/server.py +++ b/chromadb/db/impl/grpc/server.py @@ -318,6 +318,8 @@ def CreateCollection( metadata=from_proto_metadata(request.metadata), dimension=request.dimension, topic=self._assignment_policy.assign_collection(id), + database=database, + tenant=tenant, ) collections[request.id] = new_collection return CreateCollectionResponse( diff --git a/chromadb/db/mixins/sysdb.py b/chromadb/db/mixins/sysdb.py index bfae0d07692..d22960d03f4 100644 --- a/chromadb/db/mixins/sysdb.py +++ b/chromadb/db/mixins/sysdb.py @@ -226,7 +226,13 @@ def create_collection( topic = self._assignment_policy.assign_collection(id) collection = Collection( - id=id, topic=topic, name=name, metadata=metadata, dimension=dimension + id=id, + topic=topic, + name=name, + metadata=metadata, + dimension=dimension, + tenant=tenant, + database=database, ) with self.tx() as cur: @@ -379,6 +385,7 @@ def get_collections( collections_t = Table("collections") metadata_t = Table("collection_metadata") + databases_t = Table("databases") q = ( self.querybuilder() .from_(collections_t) @@ -387,6 +394,8 @@ def get_collections( collections_t.name, collections_t.topic, collections_t.dimension, + databases_t.name, + databases_t.tenant_id, metadata_t.key, metadata_t.str_value, metadata_t.int_value, @@ -394,6 +403,8 @@ def get_collections( ) .left_join(metadata_t) .on(collections_t.id == metadata_t.collection_id) + .left_join(databases_t) + .on(collections_t.database_id == databases_t.id) .orderby(collections_t.id) ) if id: @@ -433,6 +444,8 @@ def get_collections( name=name, metadata=metadata, dimension=dimension, + tenant=str(rows[0][5]), + database=str(rows[0][4]), ) ) diff --git a/chromadb/proto/chroma_pb2.py b/chromadb/proto/chroma_pb2.py index 3644a70c74d..a61d26ecfc6 100644 --- a/chromadb/proto/chroma_pb2.py +++ b/chromadb/proto/chroma_pb2.py @@ -13,21 +13,22 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1b\x63hromadb/proto/chroma.proto\x12\x06\x63hroma\"&\n\x06Status\x12\x0e\n\x06reason\x18\x01 \x01(\t\x12\x0c\n\x04\x63ode\x18\x02 \x01(\x05\"0\n\x0e\x43hromaResponse\x12\x1e\n\x06status\x18\x01 \x01(\x0b\x32\x0e.chroma.Status\"U\n\x06Vector\x12\x11\n\tdimension\x18\x01 \x01(\x05\x12\x0e\n\x06vector\x18\x02 \x01(\x0c\x12(\n\x08\x65ncoding\x18\x03 \x01(\x0e\x32\x16.chroma.ScalarEncoding\"\xca\x01\n\x07Segment\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12#\n\x05scope\x18\x03 \x01(\x0e\x32\x14.chroma.SegmentScope\x12\x12\n\x05topic\x18\x04 \x01(\tH\x00\x88\x01\x01\x12\x17\n\ncollection\x18\x05 \x01(\tH\x01\x88\x01\x01\x12-\n\x08metadata\x18\x06 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x02\x88\x01\x01\x42\x08\n\x06_topicB\r\n\x0b_collectionB\x0b\n\t_metadata\"\x97\x01\n\nCollection\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\r\n\x05topic\x18\x03 \x01(\t\x12-\n\x08metadata\x18\x04 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x00\x88\x01\x01\x12\x16\n\tdimension\x18\x05 \x01(\x05H\x01\x88\x01\x01\x42\x0b\n\t_metadataB\x0c\n\n_dimension\"4\n\x08\x44\x61tabase\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0e\n\x06tenant\x18\x03 \x01(\t\"\x16\n\x06Tenant\x12\x0c\n\x04name\x18\x01 \x01(\t\"b\n\x13UpdateMetadataValue\x12\x16\n\x0cstring_value\x18\x01 \x01(\tH\x00\x12\x13\n\tint_value\x18\x02 \x01(\x03H\x00\x12\x15\n\x0b\x66loat_value\x18\x03 \x01(\x01H\x00\x42\x07\n\x05value\"\x96\x01\n\x0eUpdateMetadata\x12\x36\n\x08metadata\x18\x01 \x03(\x0b\x32$.chroma.UpdateMetadata.MetadataEntry\x1aL\n\rMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12*\n\x05value\x18\x02 \x01(\x0b\x32\x1b.chroma.UpdateMetadataValue:\x02\x38\x01\"\xb5\x01\n\x15SubmitEmbeddingRecord\x12\n\n\x02id\x18\x01 \x01(\t\x12#\n\x06vector\x18\x02 \x01(\x0b\x32\x0e.chroma.VectorH\x00\x88\x01\x01\x12-\n\x08metadata\x18\x03 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x01\x88\x01\x01\x12$\n\toperation\x18\x04 \x01(\x0e\x32\x11.chroma.OperationB\t\n\x07_vectorB\x0b\n\t_metadata\"S\n\x15VectorEmbeddingRecord\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0e\n\x06seq_id\x18\x02 \x01(\x0c\x12\x1e\n\x06vector\x18\x03 \x01(\x0b\x32\x0e.chroma.Vector\"q\n\x11VectorQueryResult\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0e\n\x06seq_id\x18\x02 \x01(\x0c\x12\x10\n\x08\x64istance\x18\x03 \x01(\x01\x12#\n\x06vector\x18\x04 \x01(\x0b\x32\x0e.chroma.VectorH\x00\x88\x01\x01\x42\t\n\x07_vector\"@\n\x12VectorQueryResults\x12*\n\x07results\x18\x01 \x03(\x0b\x32\x19.chroma.VectorQueryResult\"(\n\x15SegmentServerResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\"4\n\x11GetVectorsRequest\x12\x0b\n\x03ids\x18\x01 \x03(\t\x12\x12\n\nsegment_id\x18\x02 \x01(\t\"D\n\x12GetVectorsResponse\x12.\n\x07records\x18\x01 \x03(\x0b\x32\x1d.chroma.VectorEmbeddingRecord\"\x86\x01\n\x13QueryVectorsRequest\x12\x1f\n\x07vectors\x18\x01 \x03(\x0b\x32\x0e.chroma.Vector\x12\t\n\x01k\x18\x02 \x01(\x05\x12\x13\n\x0b\x61llowed_ids\x18\x03 \x03(\t\x12\x1a\n\x12include_embeddings\x18\x04 \x01(\x08\x12\x12\n\nsegment_id\x18\x05 \x01(\t\"C\n\x14QueryVectorsResponse\x12+\n\x07results\x18\x01 \x03(\x0b\x32\x1a.chroma.VectorQueryResults*8\n\tOperation\x12\x07\n\x03\x41\x44\x44\x10\x00\x12\n\n\x06UPDATE\x10\x01\x12\n\n\x06UPSERT\x10\x02\x12\n\n\x06\x44\x45LETE\x10\x03*(\n\x0eScalarEncoding\x12\x0b\n\x07\x46LOAT32\x10\x00\x12\t\n\x05INT32\x10\x01*(\n\x0cSegmentScope\x12\n\n\x06VECTOR\x10\x00\x12\x0c\n\x08METADATA\x10\x01\x32\x94\x01\n\rSegmentServer\x12?\n\x0bLoadSegment\x12\x0f.chroma.Segment\x1a\x1d.chroma.SegmentServerResponse\"\x00\x12\x42\n\x0eReleaseSegment\x12\x0f.chroma.Segment\x1a\x1d.chroma.SegmentServerResponse\"\x00\x32\xa2\x01\n\x0cVectorReader\x12\x45\n\nGetVectors\x12\x19.chroma.GetVectorsRequest\x1a\x1a.chroma.GetVectorsResponse\"\x00\x12K\n\x0cQueryVectors\x12\x1b.chroma.QueryVectorsRequest\x1a\x1c.chroma.QueryVectorsResponse\"\x00\x62\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1b\x63hromadb/proto/chroma.proto\x12\x06\x63hroma\"&\n\x06Status\x12\x0e\n\x06reason\x18\x01 \x01(\t\x12\x0c\n\x04\x63ode\x18\x02 \x01(\x05\"0\n\x0e\x43hromaResponse\x12\x1e\n\x06status\x18\x01 \x01(\x0b\x32\x0e.chroma.Status\"U\n\x06Vector\x12\x11\n\tdimension\x18\x01 \x01(\x05\x12\x0e\n\x06vector\x18\x02 \x01(\x0c\x12(\n\x08\x65ncoding\x18\x03 \x01(\x0e\x32\x16.chroma.ScalarEncoding\"\xca\x01\n\x07Segment\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12#\n\x05scope\x18\x03 \x01(\x0e\x32\x14.chroma.SegmentScope\x12\x12\n\x05topic\x18\x04 \x01(\tH\x00\x88\x01\x01\x12\x17\n\ncollection\x18\x05 \x01(\tH\x01\x88\x01\x01\x12-\n\x08metadata\x18\x06 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x02\x88\x01\x01\x42\x08\n\x06_topicB\r\n\x0b_collectionB\x0b\n\t_metadata\"\xb9\x01\n\nCollection\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\r\n\x05topic\x18\x03 \x01(\t\x12-\n\x08metadata\x18\x04 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x00\x88\x01\x01\x12\x16\n\tdimension\x18\x05 \x01(\x05H\x01\x88\x01\x01\x12\x0e\n\x06tenant\x18\x06 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x07 \x01(\tB\x0b\n\t_metadataB\x0c\n\n_dimension\"4\n\x08\x44\x61tabase\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0e\n\x06tenant\x18\x03 \x01(\t\"\x16\n\x06Tenant\x12\x0c\n\x04name\x18\x01 \x01(\t\"b\n\x13UpdateMetadataValue\x12\x16\n\x0cstring_value\x18\x01 \x01(\tH\x00\x12\x13\n\tint_value\x18\x02 \x01(\x03H\x00\x12\x15\n\x0b\x66loat_value\x18\x03 \x01(\x01H\x00\x42\x07\n\x05value\"\x96\x01\n\x0eUpdateMetadata\x12\x36\n\x08metadata\x18\x01 \x03(\x0b\x32$.chroma.UpdateMetadata.MetadataEntry\x1aL\n\rMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12*\n\x05value\x18\x02 \x01(\x0b\x32\x1b.chroma.UpdateMetadataValue:\x02\x38\x01\"\xb5\x01\n\x15SubmitEmbeddingRecord\x12\n\n\x02id\x18\x01 \x01(\t\x12#\n\x06vector\x18\x02 \x01(\x0b\x32\x0e.chroma.VectorH\x00\x88\x01\x01\x12-\n\x08metadata\x18\x03 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x01\x88\x01\x01\x12$\n\toperation\x18\x04 \x01(\x0e\x32\x11.chroma.OperationB\t\n\x07_vectorB\x0b\n\t_metadata\"S\n\x15VectorEmbeddingRecord\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0e\n\x06seq_id\x18\x02 \x01(\x0c\x12\x1e\n\x06vector\x18\x03 \x01(\x0b\x32\x0e.chroma.Vector\"q\n\x11VectorQueryResult\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0e\n\x06seq_id\x18\x02 \x01(\x0c\x12\x10\n\x08\x64istance\x18\x03 \x01(\x01\x12#\n\x06vector\x18\x04 \x01(\x0b\x32\x0e.chroma.VectorH\x00\x88\x01\x01\x42\t\n\x07_vector\"@\n\x12VectorQueryResults\x12*\n\x07results\x18\x01 \x03(\x0b\x32\x19.chroma.VectorQueryResult\"(\n\x15SegmentServerResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\"4\n\x11GetVectorsRequest\x12\x0b\n\x03ids\x18\x01 \x03(\t\x12\x12\n\nsegment_id\x18\x02 \x01(\t\"D\n\x12GetVectorsResponse\x12.\n\x07records\x18\x01 \x03(\x0b\x32\x1d.chroma.VectorEmbeddingRecord\"\x86\x01\n\x13QueryVectorsRequest\x12\x1f\n\x07vectors\x18\x01 \x03(\x0b\x32\x0e.chroma.Vector\x12\t\n\x01k\x18\x02 \x01(\x05\x12\x13\n\x0b\x61llowed_ids\x18\x03 \x03(\t\x12\x1a\n\x12include_embeddings\x18\x04 \x01(\x08\x12\x12\n\nsegment_id\x18\x05 \x01(\t\"C\n\x14QueryVectorsResponse\x12+\n\x07results\x18\x01 \x03(\x0b\x32\x1a.chroma.VectorQueryResults*8\n\tOperation\x12\x07\n\x03\x41\x44\x44\x10\x00\x12\n\n\x06UPDATE\x10\x01\x12\n\n\x06UPSERT\x10\x02\x12\n\n\x06\x44\x45LETE\x10\x03*(\n\x0eScalarEncoding\x12\x0b\n\x07\x46LOAT32\x10\x00\x12\t\n\x05INT32\x10\x01*(\n\x0cSegmentScope\x12\n\n\x06VECTOR\x10\x00\x12\x0c\n\x08METADATA\x10\x01\x32\x94\x01\n\rSegmentServer\x12?\n\x0bLoadSegment\x12\x0f.chroma.Segment\x1a\x1d.chroma.SegmentServerResponse\"\x00\x12\x42\n\x0eReleaseSegment\x12\x0f.chroma.Segment\x1a\x1d.chroma.SegmentServerResponse\"\x00\x32\xa2\x01\n\x0cVectorReader\x12\x45\n\nGetVectors\x12\x19.chroma.GetVectorsRequest\x1a\x1a.chroma.GetVectorsResponse\"\x00\x12K\n\x0cQueryVectors\x12\x1b.chroma.QueryVectorsRequest\x1a\x1c.chroma.QueryVectorsResponse\"\x00\x42\x43ZAgithub.com/chroma/chroma-coordinator/internal/proto/coordinatorpbb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'chromadb.proto.chroma_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = b'ZAgithub.com/chroma/chroma-coordinator/internal/proto/coordinatorpb' _UPDATEMETADATA_METADATAENTRY._options = None _UPDATEMETADATA_METADATAENTRY._serialized_options = b'8\001' - _globals['_OPERATION']._serialized_start=1728 - _globals['_OPERATION']._serialized_end=1784 - _globals['_SCALARENCODING']._serialized_start=1786 - _globals['_SCALARENCODING']._serialized_end=1826 - _globals['_SEGMENTSCOPE']._serialized_start=1828 - _globals['_SEGMENTSCOPE']._serialized_end=1868 + _globals['_OPERATION']._serialized_start=1762 + _globals['_OPERATION']._serialized_end=1818 + _globals['_SCALARENCODING']._serialized_start=1820 + _globals['_SCALARENCODING']._serialized_end=1860 + _globals['_SEGMENTSCOPE']._serialized_start=1862 + _globals['_SEGMENTSCOPE']._serialized_end=1902 _globals['_STATUS']._serialized_start=39 _globals['_STATUS']._serialized_end=77 _globals['_CHROMARESPONSE']._serialized_start=79 @@ -37,37 +38,37 @@ _globals['_SEGMENT']._serialized_start=217 _globals['_SEGMENT']._serialized_end=419 _globals['_COLLECTION']._serialized_start=422 - _globals['_COLLECTION']._serialized_end=573 - _globals['_DATABASE']._serialized_start=575 - _globals['_DATABASE']._serialized_end=627 - _globals['_TENANT']._serialized_start=629 - _globals['_TENANT']._serialized_end=651 - _globals['_UPDATEMETADATAVALUE']._serialized_start=653 - _globals['_UPDATEMETADATAVALUE']._serialized_end=751 - _globals['_UPDATEMETADATA']._serialized_start=754 - _globals['_UPDATEMETADATA']._serialized_end=904 - _globals['_UPDATEMETADATA_METADATAENTRY']._serialized_start=828 - _globals['_UPDATEMETADATA_METADATAENTRY']._serialized_end=904 - _globals['_SUBMITEMBEDDINGRECORD']._serialized_start=907 - _globals['_SUBMITEMBEDDINGRECORD']._serialized_end=1088 - _globals['_VECTOREMBEDDINGRECORD']._serialized_start=1090 - _globals['_VECTOREMBEDDINGRECORD']._serialized_end=1173 - _globals['_VECTORQUERYRESULT']._serialized_start=1175 - _globals['_VECTORQUERYRESULT']._serialized_end=1288 - _globals['_VECTORQUERYRESULTS']._serialized_start=1290 - _globals['_VECTORQUERYRESULTS']._serialized_end=1354 - _globals['_SEGMENTSERVERRESPONSE']._serialized_start=1356 - _globals['_SEGMENTSERVERRESPONSE']._serialized_end=1396 - _globals['_GETVECTORSREQUEST']._serialized_start=1398 - _globals['_GETVECTORSREQUEST']._serialized_end=1450 - _globals['_GETVECTORSRESPONSE']._serialized_start=1452 - _globals['_GETVECTORSRESPONSE']._serialized_end=1520 - _globals['_QUERYVECTORSREQUEST']._serialized_start=1523 - _globals['_QUERYVECTORSREQUEST']._serialized_end=1657 - _globals['_QUERYVECTORSRESPONSE']._serialized_start=1659 - _globals['_QUERYVECTORSRESPONSE']._serialized_end=1726 - _globals['_SEGMENTSERVER']._serialized_start=1871 - _globals['_SEGMENTSERVER']._serialized_end=2019 - _globals['_VECTORREADER']._serialized_start=2022 - _globals['_VECTORREADER']._serialized_end=2184 + _globals['_COLLECTION']._serialized_end=607 + _globals['_DATABASE']._serialized_start=609 + _globals['_DATABASE']._serialized_end=661 + _globals['_TENANT']._serialized_start=663 + _globals['_TENANT']._serialized_end=685 + _globals['_UPDATEMETADATAVALUE']._serialized_start=687 + _globals['_UPDATEMETADATAVALUE']._serialized_end=785 + _globals['_UPDATEMETADATA']._serialized_start=788 + _globals['_UPDATEMETADATA']._serialized_end=938 + _globals['_UPDATEMETADATA_METADATAENTRY']._serialized_start=862 + _globals['_UPDATEMETADATA_METADATAENTRY']._serialized_end=938 + _globals['_SUBMITEMBEDDINGRECORD']._serialized_start=941 + _globals['_SUBMITEMBEDDINGRECORD']._serialized_end=1122 + _globals['_VECTOREMBEDDINGRECORD']._serialized_start=1124 + _globals['_VECTOREMBEDDINGRECORD']._serialized_end=1207 + _globals['_VECTORQUERYRESULT']._serialized_start=1209 + _globals['_VECTORQUERYRESULT']._serialized_end=1322 + _globals['_VECTORQUERYRESULTS']._serialized_start=1324 + _globals['_VECTORQUERYRESULTS']._serialized_end=1388 + _globals['_SEGMENTSERVERRESPONSE']._serialized_start=1390 + _globals['_SEGMENTSERVERRESPONSE']._serialized_end=1430 + _globals['_GETVECTORSREQUEST']._serialized_start=1432 + _globals['_GETVECTORSREQUEST']._serialized_end=1484 + _globals['_GETVECTORSRESPONSE']._serialized_start=1486 + _globals['_GETVECTORSRESPONSE']._serialized_end=1554 + _globals['_QUERYVECTORSREQUEST']._serialized_start=1557 + _globals['_QUERYVECTORSREQUEST']._serialized_end=1691 + _globals['_QUERYVECTORSRESPONSE']._serialized_start=1693 + _globals['_QUERYVECTORSRESPONSE']._serialized_end=1760 + _globals['_SEGMENTSERVER']._serialized_start=1905 + _globals['_SEGMENTSERVER']._serialized_end=2053 + _globals['_VECTORREADER']._serialized_start=2056 + _globals['_VECTORREADER']._serialized_end=2218 # @@protoc_insertion_point(module_scope) diff --git a/chromadb/proto/chroma_pb2.pyi b/chromadb/proto/chroma_pb2.pyi index 386f253e776..690a81c8301 100644 --- a/chromadb/proto/chroma_pb2.pyi +++ b/chromadb/proto/chroma_pb2.pyi @@ -72,18 +72,22 @@ class Segment(_message.Message): def __init__(self, id: _Optional[str] = ..., type: _Optional[str] = ..., scope: _Optional[_Union[SegmentScope, str]] = ..., topic: _Optional[str] = ..., collection: _Optional[str] = ..., metadata: _Optional[_Union[UpdateMetadata, _Mapping]] = ...) -> None: ... class Collection(_message.Message): - __slots__ = ["id", "name", "topic", "metadata", "dimension"] + __slots__ = ["id", "name", "topic", "metadata", "dimension", "tenant", "database"] ID_FIELD_NUMBER: _ClassVar[int] NAME_FIELD_NUMBER: _ClassVar[int] TOPIC_FIELD_NUMBER: _ClassVar[int] METADATA_FIELD_NUMBER: _ClassVar[int] DIMENSION_FIELD_NUMBER: _ClassVar[int] + TENANT_FIELD_NUMBER: _ClassVar[int] + DATABASE_FIELD_NUMBER: _ClassVar[int] id: str name: str topic: str metadata: UpdateMetadata dimension: int - def __init__(self, id: _Optional[str] = ..., name: _Optional[str] = ..., topic: _Optional[str] = ..., metadata: _Optional[_Union[UpdateMetadata, _Mapping]] = ..., dimension: _Optional[int] = ...) -> None: ... + tenant: str + database: str + def __init__(self, id: _Optional[str] = ..., name: _Optional[str] = ..., topic: _Optional[str] = ..., metadata: _Optional[_Union[UpdateMetadata, _Mapping]] = ..., dimension: _Optional[int] = ..., tenant: _Optional[str] = ..., database: _Optional[str] = ...) -> None: ... class Database(_message.Message): __slots__ = ["id", "name", "tenant"] diff --git a/chromadb/proto/convert.py b/chromadb/proto/convert.py index 129d7e3ff2a..5919071a275 100644 --- a/chromadb/proto/convert.py +++ b/chromadb/proto/convert.py @@ -69,7 +69,8 @@ def from_proto_operation(operation: proto.Operation) -> Operation: elif operation == proto.Operation.DELETE: return Operation.DELETE else: - raise RuntimeError(f"Unknown operation {operation}") # TODO: full error + # TODO: full error + raise RuntimeError(f"Unknown operation {operation}") def from_proto_metadata(metadata: proto.UpdateMetadata) -> Optional[Metadata]: @@ -200,6 +201,8 @@ def from_proto_collection(collection: proto.Collection) -> Collection: dimension=collection.dimension if collection.HasField("dimension") and collection.dimension else None, + database=collection.database, + tenant=collection.tenant, ) @@ -212,6 +215,8 @@ def to_proto_collection(collection: Collection) -> proto.Collection: if collection["metadata"] is None else to_proto_update_metadata(collection["metadata"]), dimension=collection["dimension"], + tenant=collection["tenant"], + database=collection["database"], ) diff --git a/chromadb/proto/coordinator_pb2.py b/chromadb/proto/coordinator_pb2.py index 42039c2d23f..fda6a099867 100644 --- a/chromadb/proto/coordinator_pb2.py +++ b/chromadb/proto/coordinator_pb2.py @@ -15,13 +15,14 @@ from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2 -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n chromadb/proto/coordinator.proto\x12\x06\x63hroma\x1a\x1b\x63hromadb/proto/chroma.proto\x1a\x1bgoogle/protobuf/empty.proto\"A\n\x15\x43reateDatabaseRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0e\n\x06tenant\x18\x03 \x01(\t\"2\n\x12GetDatabaseRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0e\n\x06tenant\x18\x02 \x01(\t\"Y\n\x13GetDatabaseResponse\x12\"\n\x08\x64\x61tabase\x18\x01 \x01(\x0b\x32\x10.chroma.Database\x12\x1e\n\x06status\x18\x02 \x01(\x0b\x32\x0e.chroma.Status\"#\n\x13\x43reateTenantRequest\x12\x0c\n\x04name\x18\x02 \x01(\t\" \n\x10GetTenantRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\"S\n\x11GetTenantResponse\x12\x1e\n\x06tenant\x18\x01 \x01(\x0b\x32\x0e.chroma.Tenant\x12\x1e\n\x06status\x18\x02 \x01(\x0b\x32\x0e.chroma.Status\"8\n\x14\x43reateSegmentRequest\x12 \n\x07segment\x18\x01 \x01(\x0b\x32\x0f.chroma.Segment\"\"\n\x14\x44\x65leteSegmentRequest\x12\n\n\x02id\x18\x01 \x01(\t\"\xc2\x01\n\x12GetSegmentsRequest\x12\x0f\n\x02id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x11\n\x04type\x18\x02 \x01(\tH\x01\x88\x01\x01\x12(\n\x05scope\x18\x03 \x01(\x0e\x32\x14.chroma.SegmentScopeH\x02\x88\x01\x01\x12\x12\n\x05topic\x18\x04 \x01(\tH\x03\x88\x01\x01\x12\x17\n\ncollection\x18\x05 \x01(\tH\x04\x88\x01\x01\x42\x05\n\x03_idB\x07\n\x05_typeB\x08\n\x06_scopeB\x08\n\x06_topicB\r\n\x0b_collection\"X\n\x13GetSegmentsResponse\x12!\n\x08segments\x18\x01 \x03(\x0b\x32\x0f.chroma.Segment\x12\x1e\n\x06status\x18\x02 \x01(\x0b\x32\x0e.chroma.Status\"\xfa\x01\n\x14UpdateSegmentRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0f\n\x05topic\x18\x02 \x01(\tH\x00\x12\x15\n\x0breset_topic\x18\x03 \x01(\x08H\x00\x12\x14\n\ncollection\x18\x04 \x01(\tH\x01\x12\x1a\n\x10reset_collection\x18\x05 \x01(\x08H\x01\x12*\n\x08metadata\x18\x06 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x02\x12\x18\n\x0ereset_metadata\x18\x07 \x01(\x08H\x02\x42\x0e\n\x0ctopic_updateB\x13\n\x11\x63ollection_updateB\x11\n\x0fmetadata_update\"\xe5\x01\n\x17\x43reateCollectionRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12-\n\x08metadata\x18\x03 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x00\x88\x01\x01\x12\x16\n\tdimension\x18\x04 \x01(\x05H\x01\x88\x01\x01\x12\x1a\n\rget_or_create\x18\x05 \x01(\x08H\x02\x88\x01\x01\x12\x0e\n\x06tenant\x18\x06 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x07 \x01(\tB\x0b\n\t_metadataB\x0c\n\n_dimensionB\x10\n\x0e_get_or_create\"s\n\x18\x43reateCollectionResponse\x12&\n\ncollection\x18\x01 \x01(\x0b\x32\x12.chroma.Collection\x12\x0f\n\x07\x63reated\x18\x02 \x01(\x08\x12\x1e\n\x06status\x18\x03 \x01(\x0b\x32\x0e.chroma.Status\"G\n\x17\x44\x65leteCollectionRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0e\n\x06tenant\x18\x02 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x03 \x01(\t\"\x8b\x01\n\x15GetCollectionsRequest\x12\x0f\n\x02id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x11\n\x04name\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x12\n\x05topic\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x0e\n\x06tenant\x18\x04 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x05 \x01(\tB\x05\n\x03_idB\x07\n\x05_nameB\x08\n\x06_topic\"a\n\x16GetCollectionsResponse\x12\'\n\x0b\x63ollections\x18\x01 \x03(\x0b\x32\x12.chroma.Collection\x12\x1e\n\x06status\x18\x02 \x01(\x0b\x32\x0e.chroma.Status\"\xde\x01\n\x17UpdateCollectionRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x12\n\x05topic\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x11\n\x04name\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x16\n\tdimension\x18\x04 \x01(\x05H\x03\x88\x01\x01\x12*\n\x08metadata\x18\x05 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x00\x12\x18\n\x0ereset_metadata\x18\x06 \x01(\x08H\x00\x42\x11\n\x0fmetadata_updateB\x08\n\x06_topicB\x07\n\x05_nameB\x0c\n\n_dimension2\xd6\x07\n\x05SysDB\x12I\n\x0e\x43reateDatabase\x12\x1d.chroma.CreateDatabaseRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12H\n\x0bGetDatabase\x12\x1a.chroma.GetDatabaseRequest\x1a\x1b.chroma.GetDatabaseResponse\"\x00\x12\x45\n\x0c\x43reateTenant\x12\x1b.chroma.CreateTenantRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12\x42\n\tGetTenant\x12\x18.chroma.GetTenantRequest\x1a\x19.chroma.GetTenantResponse\"\x00\x12G\n\rCreateSegment\x12\x1c.chroma.CreateSegmentRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12G\n\rDeleteSegment\x12\x1c.chroma.DeleteSegmentRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12H\n\x0bGetSegments\x12\x1a.chroma.GetSegmentsRequest\x1a\x1b.chroma.GetSegmentsResponse\"\x00\x12G\n\rUpdateSegment\x12\x1c.chroma.UpdateSegmentRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12W\n\x10\x43reateCollection\x12\x1f.chroma.CreateCollectionRequest\x1a .chroma.CreateCollectionResponse\"\x00\x12M\n\x10\x44\x65leteCollection\x12\x1f.chroma.DeleteCollectionRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12Q\n\x0eGetCollections\x12\x1d.chroma.GetCollectionsRequest\x1a\x1e.chroma.GetCollectionsResponse\"\x00\x12M\n\x10UpdateCollection\x12\x1f.chroma.UpdateCollectionRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12>\n\nResetState\x12\x16.google.protobuf.Empty\x1a\x16.chroma.ChromaResponse\"\x00\x62\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n chromadb/proto/coordinator.proto\x12\x06\x63hroma\x1a\x1b\x63hromadb/proto/chroma.proto\x1a\x1bgoogle/protobuf/empty.proto\"A\n\x15\x43reateDatabaseRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0e\n\x06tenant\x18\x03 \x01(\t\"2\n\x12GetDatabaseRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0e\n\x06tenant\x18\x02 \x01(\t\"Y\n\x13GetDatabaseResponse\x12\"\n\x08\x64\x61tabase\x18\x01 \x01(\x0b\x32\x10.chroma.Database\x12\x1e\n\x06status\x18\x02 \x01(\x0b\x32\x0e.chroma.Status\"#\n\x13\x43reateTenantRequest\x12\x0c\n\x04name\x18\x02 \x01(\t\" \n\x10GetTenantRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\"S\n\x11GetTenantResponse\x12\x1e\n\x06tenant\x18\x01 \x01(\x0b\x32\x0e.chroma.Tenant\x12\x1e\n\x06status\x18\x02 \x01(\x0b\x32\x0e.chroma.Status\"8\n\x14\x43reateSegmentRequest\x12 \n\x07segment\x18\x01 \x01(\x0b\x32\x0f.chroma.Segment\"\"\n\x14\x44\x65leteSegmentRequest\x12\n\n\x02id\x18\x01 \x01(\t\"\xc2\x01\n\x12GetSegmentsRequest\x12\x0f\n\x02id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x11\n\x04type\x18\x02 \x01(\tH\x01\x88\x01\x01\x12(\n\x05scope\x18\x03 \x01(\x0e\x32\x14.chroma.SegmentScopeH\x02\x88\x01\x01\x12\x12\n\x05topic\x18\x04 \x01(\tH\x03\x88\x01\x01\x12\x17\n\ncollection\x18\x05 \x01(\tH\x04\x88\x01\x01\x42\x05\n\x03_idB\x07\n\x05_typeB\x08\n\x06_scopeB\x08\n\x06_topicB\r\n\x0b_collection\"X\n\x13GetSegmentsResponse\x12!\n\x08segments\x18\x01 \x03(\x0b\x32\x0f.chroma.Segment\x12\x1e\n\x06status\x18\x02 \x01(\x0b\x32\x0e.chroma.Status\"\xfa\x01\n\x14UpdateSegmentRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0f\n\x05topic\x18\x02 \x01(\tH\x00\x12\x15\n\x0breset_topic\x18\x03 \x01(\x08H\x00\x12\x14\n\ncollection\x18\x04 \x01(\tH\x01\x12\x1a\n\x10reset_collection\x18\x05 \x01(\x08H\x01\x12*\n\x08metadata\x18\x06 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x02\x12\x18\n\x0ereset_metadata\x18\x07 \x01(\x08H\x02\x42\x0e\n\x0ctopic_updateB\x13\n\x11\x63ollection_updateB\x11\n\x0fmetadata_update\"\xe5\x01\n\x17\x43reateCollectionRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12-\n\x08metadata\x18\x03 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x00\x88\x01\x01\x12\x16\n\tdimension\x18\x04 \x01(\x05H\x01\x88\x01\x01\x12\x1a\n\rget_or_create\x18\x05 \x01(\x08H\x02\x88\x01\x01\x12\x0e\n\x06tenant\x18\x06 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x07 \x01(\tB\x0b\n\t_metadataB\x0c\n\n_dimensionB\x10\n\x0e_get_or_create\"s\n\x18\x43reateCollectionResponse\x12&\n\ncollection\x18\x01 \x01(\x0b\x32\x12.chroma.Collection\x12\x0f\n\x07\x63reated\x18\x02 \x01(\x08\x12\x1e\n\x06status\x18\x03 \x01(\x0b\x32\x0e.chroma.Status\"G\n\x17\x44\x65leteCollectionRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0e\n\x06tenant\x18\x02 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x03 \x01(\t\"\x8b\x01\n\x15GetCollectionsRequest\x12\x0f\n\x02id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x11\n\x04name\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x12\n\x05topic\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x0e\n\x06tenant\x18\x04 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x05 \x01(\tB\x05\n\x03_idB\x07\n\x05_nameB\x08\n\x06_topic\"a\n\x16GetCollectionsResponse\x12\'\n\x0b\x63ollections\x18\x01 \x03(\x0b\x32\x12.chroma.Collection\x12\x1e\n\x06status\x18\x02 \x01(\x0b\x32\x0e.chroma.Status\"\xde\x01\n\x17UpdateCollectionRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x12\n\x05topic\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x11\n\x04name\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x16\n\tdimension\x18\x04 \x01(\x05H\x03\x88\x01\x01\x12*\n\x08metadata\x18\x05 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x00\x12\x18\n\x0ereset_metadata\x18\x06 \x01(\x08H\x00\x42\x11\n\x0fmetadata_updateB\x08\n\x06_topicB\x07\n\x05_nameB\x0c\n\n_dimension2\xd6\x07\n\x05SysDB\x12I\n\x0e\x43reateDatabase\x12\x1d.chroma.CreateDatabaseRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12H\n\x0bGetDatabase\x12\x1a.chroma.GetDatabaseRequest\x1a\x1b.chroma.GetDatabaseResponse\"\x00\x12\x45\n\x0c\x43reateTenant\x12\x1b.chroma.CreateTenantRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12\x42\n\tGetTenant\x12\x18.chroma.GetTenantRequest\x1a\x19.chroma.GetTenantResponse\"\x00\x12G\n\rCreateSegment\x12\x1c.chroma.CreateSegmentRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12G\n\rDeleteSegment\x12\x1c.chroma.DeleteSegmentRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12H\n\x0bGetSegments\x12\x1a.chroma.GetSegmentsRequest\x1a\x1b.chroma.GetSegmentsResponse\"\x00\x12G\n\rUpdateSegment\x12\x1c.chroma.UpdateSegmentRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12W\n\x10\x43reateCollection\x12\x1f.chroma.CreateCollectionRequest\x1a .chroma.CreateCollectionResponse\"\x00\x12M\n\x10\x44\x65leteCollection\x12\x1f.chroma.DeleteCollectionRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12Q\n\x0eGetCollections\x12\x1d.chroma.GetCollectionsRequest\x1a\x1e.chroma.GetCollectionsResponse\"\x00\x12M\n\x10UpdateCollection\x12\x1f.chroma.UpdateCollectionRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12>\n\nResetState\x12\x16.google.protobuf.Empty\x1a\x16.chroma.ChromaResponse\"\x00\x42\x43ZAgithub.com/chroma/chroma-coordinator/internal/proto/coordinatorpbb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'chromadb.proto.coordinator_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = b'ZAgithub.com/chroma/chroma-coordinator/internal/proto/coordinatorpb' _globals['_CREATEDATABASEREQUEST']._serialized_start=102 _globals['_CREATEDATABASEREQUEST']._serialized_end=167 _globals['_GETDATABASEREQUEST']._serialized_start=169 diff --git a/chromadb/test/db/test_system.py b/chromadb/test/db/test_system.py index acdf2c6704c..b8206f2bd21 100644 --- a/chromadb/test/db/test_system.py +++ b/chromadb/test/db/test_system.py @@ -20,6 +20,8 @@ from pytest import FixtureRequest import uuid +# These are the sample collections that are used in the tests below. Tests can override +# the fields as needed. sample_collections = [ Collection( id=uuid.UUID(int=1), @@ -27,6 +29,8 @@ topic="persistent://test-tenant/test-topic/00000000-0000-0000-0000-000000000001", metadata={"test_str": "str1", "test_int": 1, "test_float": 1.3}, dimension=128, + database=DEFAULT_DATABASE, + tenant=DEFAULT_TENANT, ), Collection( id=uuid.UUID(int=2), @@ -34,6 +38,8 @@ topic="persistent://test-tenant/test-topic/00000000-0000-0000-0000-000000000002", metadata={"test_str": "str2", "test_int": 2, "test_float": 2.3}, dimension=None, + database=DEFAULT_DATABASE, + tenant=DEFAULT_TENANT, ), Collection( id=uuid.UUID(int=3), @@ -41,6 +47,8 @@ topic="persistent://test-tenant/test-topic/00000000-0000-0000-0000-000000000003", metadata={"test_str": "str3", "test_int": 3, "test_float": 3.3}, dimension=None, + database=DEFAULT_DATABASE, + tenant=DEFAULT_TENANT, ), ] @@ -141,6 +149,8 @@ def test_create_get_delete_collections(sysdb: SysDB) -> None: metadata=collection["metadata"], dimension=collection["dimension"], ) + collection["database"] = DEFAULT_DATABASE + collection["tenant"] = DEFAULT_TENANT results = sysdb.get_collections() results = sorted(results, key=lambda c: c["name"]) @@ -202,6 +212,8 @@ def test_update_collections(sysdb: SysDB) -> None: topic=sample_collections[0]["topic"], metadata=sample_collections[0]["metadata"], dimension=sample_collections[0]["dimension"], + database=DEFAULT_DATABASE, + tenant=DEFAULT_TENANT, ) sysdb.reset_state() @@ -360,6 +372,7 @@ def test_create_get_delete_database_and_collection(sysdb: SysDB) -> None: name=sample_collections[0]["name"], database="new_database" ) assert len(result) == 1 + sample_collections[0]["database"] = "new_database" assert result[0] == sample_collections[0] # Check that the collection in the default database exists @@ -471,6 +484,7 @@ def test_get_multiple_with_database(sysdb: SysDB) -> None: dimension=collection["dimension"], database="new_database", ) + collection["database"] = "new_database" # Get all collections in the new database result = sysdb.get_collections(database="new_database") @@ -508,6 +522,8 @@ def test_create_database_with_tenants(sysdb: SysDB) -> None: database="new_database", tenant="tenant1", ) + sample_collections[0]["tenant"] = "tenant1" + sample_collections[0]["database"] = "new_database" # Create a new collection in the default tenant sysdb.create_collection( @@ -518,6 +534,8 @@ def test_create_database_with_tenants(sysdb: SysDB) -> None: database="new_database", ) + sample_collections[1]["database"] = "new_database" + # Check that both tenants have the correct collections result = sysdb.get_collections(database="new_database", tenant="tenant1") assert len(result) == 1 diff --git a/chromadb/types.py b/chromadb/types.py index ed74b4f5e20..b100bb8f862 100644 --- a/chromadb/types.py +++ b/chromadb/types.py @@ -3,6 +3,7 @@ from uuid import UUID from enum import Enum + Metadata = Mapping[str, Union[str, int, float, bool]] UpdateMetadata = Mapping[str, Union[int, float, str, bool, None]] @@ -27,6 +28,8 @@ class Collection(TypedDict): topic: str metadata: Optional[Metadata] dimension: Optional[int] + tenant: str + database: str class Database(TypedDict): diff --git a/clients/js/test/collection.client.test.ts b/clients/js/test/collection.client.test.ts index 2e5ebd1c478..0045067c7ea 100644 --- a/clients/js/test/collection.client.test.ts +++ b/clients/js/test/collection.client.test.ts @@ -22,7 +22,7 @@ test("it should create a collection", async () => { expect(collection).toHaveProperty('id') expect(collection.name).toBe("test"); let collections = await chroma.listCollections(); - expect([{ name: "test", metadata: null, id: collection.id }]).toEqual( + expect([{ name: "test", metadata: null, id: collection.id, database: "default_database", tenant: "default_tenant" }]).toEqual( expect.arrayContaining(collections) ); expect([{ name: "test2", metadata: null }]).not.toEqual( @@ -39,7 +39,7 @@ test("it should create a collection", async () => { expect(collection2.metadata).toHaveProperty("test"); expect(collection2.metadata).toEqual({ test: "test" }); let collections2 = await chroma.listCollections(); - expect([{ name: "test2", metadata: { test: "test" }, id: collection2.id }]).toEqual( + expect([{ name: "test2", metadata: { test: "test" }, id: collection2.id, database: "default_database", tenant: "default_tenant" }]).toEqual( expect.arrayContaining(collections2) ); }); diff --git a/examples/basic_functionality/test_get_collection_by_id.ipynb b/examples/basic_functionality/test_get_collection_by_id.ipynb new file mode 100644 index 00000000000..99e31eec599 --- /dev/null +++ b/examples/basic_functionality/test_get_collection_by_id.ipynb @@ -0,0 +1,73 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "current_dir = os.getcwd()\n", + "os.chdir(\"../../\")" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "name='test_collection' id=UUID('97c86f7b-6197-4ebe-bf97-9b73fceccdef') metadata=None tenant='default_tenant' database='default_database'\n", + "default_tenant\n", + "default_database\n" + ] + } + ], + "source": [ + "\n", + "import chromadb\n", + "\n", + "client = chromadb.PersistentClient()\n", + "\n", + "col = client.get_or_create_collection('test_collection')\n", + "print(col)\n", + "\n", + "\n", + "col1=client.get_collection(id=col.id)\n", + "\n", + "print(col1.tenant)\n", + "assert col1.id == col.id\n", + "\n", + "assert col1.tenant == col.tenant\n", + "assert col1.name == col.name\n", + "assert col1.database == col.database\n", + "print(col1.database)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.2" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/idl/chromadb/proto/chroma.proto b/idl/chromadb/proto/chroma.proto index 848cd6f080c..a3cd6d67124 100644 --- a/idl/chromadb/proto/chroma.proto +++ b/idl/chromadb/proto/chroma.proto @@ -55,6 +55,8 @@ message Collection { string topic = 3; optional UpdateMetadata metadata = 4; optional int32 dimension = 5; + string tenant = 6; + string database = 7; } message Database {