diff --git a/qdrant_client/_pydantic_compat.py b/qdrant_client/_pydantic_compat.py index 90ef5d7e..0149a2f2 100644 --- a/qdrant_client/_pydantic_compat.py +++ b/qdrant_client/_pydantic_compat.py @@ -1,3 +1,5 @@ +import json + from typing import Any, Dict, Type, TypeVar from pydantic import BaseModel @@ -44,3 +46,10 @@ def model_fields_set(model: BaseModel) -> set: return model.model_fields_set else: return model.__fields_set__ + + +def model_json_schema(model: BaseModel, *args: Any, **kwargs: Any) -> dict[str, Any]: + if PYDANTIC_V2: + return model.model_json_schema(*args, **kwargs) + else: + return json.loads(model.schema_json(*args, **kwargs)) diff --git a/qdrant_client/embed/schema_parser.py b/qdrant_client/embed/schema_parser.py index c1722088..e25b67eb 100644 --- a/qdrant_client/embed/schema_parser.py +++ b/qdrant_client/embed/schema_parser.py @@ -2,6 +2,7 @@ from typing import List, Type, Dict, Union, Any, Set, Optional from pydantic import BaseModel +from pydantic.json_schema import model_json_schema from qdrant_client.embed.utils import FieldPath, convert_paths @@ -67,7 +68,7 @@ class ModelSchemaParser: """ CACHE_PATH = "_inspection_cache.py" - INFERENCE_OBJECT_NAMES = {"Document", "Image"} + INFERENCE_OBJECT_NAMES = {"Document", "Image", "InferenceObject"} def __init__(self) -> None: self._defs: Dict[str, Union[Dict[str, Any], List[Dict[str, Any]]]] = deepcopy(DEFS) # type: ignore[arg-type] @@ -223,7 +224,7 @@ def parse_model(self, model: Type[BaseModel]) -> None: if model_name in self._cache: return None - schema = model.model_json_schema() + schema = model_json_schema(model) self._defs.update(schema.get("$defs", {})) defs = self._replace_refs(schema)