From e5956a102ab3bb23a4efb2fd0b55a2445bcf15f8 Mon Sep 17 00:00:00 2001
From: George Panchuk <george.panchuk@qdrant.tech>
Date: Tue, 12 Nov 2024 16:59:11 +0100
Subject: [PATCH] fix: fix model json schema compat, fix fields inspected for
 inference

---
 qdrant_client/_pydantic_compat.py    | 9 +++++++++
 qdrant_client/embed/schema_parser.py | 5 +++--
 2 files changed, 12 insertions(+), 2 deletions(-)

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)