From 131e8dd6ea493380ae197f59065b8499936b30c0 Mon Sep 17 00:00:00 2001 From: Trayan Azarov Date: Wed, 19 Jun 2024 22:19:13 +0200 Subject: [PATCH] [BUG]: Fix cross version compatibility 0.5.0 or lower with 0.5.1+ (#2378) Closes #2377 #2379 ## Description of changes *Summarize the changes made by this PR.* - Improvements & Bug fixes - Making dimension and version lookup optional in the Collection model creation in fastapi client ## Test plan *How are these changes tested?* - [x] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust ## Documentation Changes N/A --- chromadb/api/__init__.py | 15 +++++++++++++++ chromadb/api/async_fastapi.py | 33 ++++----------------------------- chromadb/api/fastapi.py | 33 ++++----------------------------- 3 files changed, 23 insertions(+), 58 deletions(-) diff --git a/chromadb/api/__init__.py b/chromadb/api/__init__.py index 9481b92f272..47869ce71fe 100644 --- a/chromadb/api/__init__.py +++ b/chromadb/api/__init__.py @@ -25,6 +25,7 @@ from chromadb.config import Component, Settings from chromadb.types import Database, Tenant import chromadb.utils.embedding_functions as ef +from chromadb.types import Collection as CollectionModel # Re-export the async version from chromadb.api.async_api import ( # noqa: F401 @@ -600,3 +601,17 @@ def delete_collection( database: str = DEFAULT_DATABASE, ) -> None: pass + + +def json_to_collection_model(json_collection: dict) -> CollectionModel: + return CollectionModel( + id=json_collection["id"], + name=json_collection["name"], + metadata=json_collection["metadata"], + dimension=json_collection["dimension"] + if "dimension" in json_collection + else None, + tenant=json_collection["tenant"], + database=json_collection["database"], + version=json_collection["version"] if "version" in json_collection else None, + ) diff --git a/chromadb/api/async_fastapi.py b/chromadb/api/async_fastapi.py index 7a742b5198d..a7e50b00b4d 100644 --- a/chromadb/api/async_fastapi.py +++ b/chromadb/api/async_fastapi.py @@ -6,7 +6,7 @@ import logging import httpx from overrides import override -from chromadb.api import AsyncServerAPI +from chromadb.api import AsyncServerAPI, json_to_collection_model from chromadb.api.base_http_client import BaseHTTPClient from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, System, Settings from chromadb.telemetry.opentelemetry import ( @@ -19,7 +19,6 @@ import chromadb.utils.embedding_functions as ef from chromadb.types import Database, Tenant -from chromadb.types import Collection as CollectionModel from chromadb.api.models.AsyncCollection import AsyncCollection from chromadb.api.types import ( @@ -204,15 +203,7 @@ async def list_collections( collections = [] for json_collection in resp_json: - model = CollectionModel( - id=json_collection["id"], - name=json_collection["name"], - metadata=json_collection["metadata"], - dimension=json_collection["dimension"], - tenant=json_collection["tenant"], - database=json_collection["database"], - version=json_collection["version"], - ) + model = json_to_collection_model(json_collection) collections.append(AsyncCollection(client=self, model=model)) @@ -257,15 +248,7 @@ async def create_collection( params={"tenant": tenant, "database": database}, ) - model = CollectionModel( - id=resp_json["id"], - name=resp_json["name"], - metadata=resp_json["metadata"], - dimension=resp_json["dimension"], - tenant=resp_json["tenant"], - database=resp_json["database"], - version=resp_json["version"], - ) + model = json_to_collection_model(resp_json) return AsyncCollection( client=self, @@ -300,15 +283,7 @@ async def get_collection( params=params, ) - model = CollectionModel( - id=resp_json["id"], - name=resp_json["name"], - metadata=resp_json["metadata"], - dimension=resp_json["dimension"], - tenant=resp_json["tenant"], - database=resp_json["database"], - version=resp_json["version"], - ) + model = json_to_collection_model(resp_json) return AsyncCollection( client=self, diff --git a/chromadb/api/fastapi.py b/chromadb/api/fastapi.py index 18beec10a11..e06706c0264 100644 --- a/chromadb/api/fastapi.py +++ b/chromadb/api/fastapi.py @@ -10,7 +10,7 @@ from chromadb.api.base_http_client import BaseHTTPClient from chromadb.types import Database, Tenant import chromadb.utils.embedding_functions as ef -from chromadb.api import ServerAPI +from chromadb.api import ServerAPI, json_to_collection_model from chromadb.api.models.Collection import Collection from chromadb.api.types import ( DataLoader, @@ -40,7 +40,6 @@ trace_method, ) from chromadb.telemetry.product import ProductTelemetryClient -from chromadb.types import Collection as CollectionModel logger = logging.getLogger(__name__) @@ -159,15 +158,7 @@ def list_collections( ) collections = [] for json_collection in json_collections: - model = CollectionModel( - id=json_collection["id"], - name=json_collection["name"], - metadata=json_collection["metadata"], - dimension=json_collection["dimension"], - tenant=json_collection["tenant"], - database=json_collection["database"], - version=json_collection["version"], - ) + model = json_to_collection_model(json_collection) collections.append(Collection(self, model=model)) return collections @@ -211,15 +202,7 @@ def create_collection( params={"tenant": tenant, "database": database}, ) - model = CollectionModel( - id=resp_json["id"], - name=resp_json["name"], - metadata=resp_json["metadata"], - dimension=resp_json["dimension"], - tenant=resp_json["tenant"], - database=resp_json["database"], - version=resp_json["version"], - ) + model = json_to_collection_model(resp_json) return Collection( client=self, model=model, @@ -254,15 +237,7 @@ def get_collection( params=_params, ) - model = CollectionModel( - id=resp_json["id"], - name=resp_json["name"], - metadata=resp_json["metadata"], - dimension=resp_json["dimension"], - tenant=resp_json["tenant"], - database=resp_json["database"], - version=resp_json["version"], - ) + model = json_to_collection_model(resp_json) return Collection( client=self, model=model,