feat: use Pydantic V2 models when possible instead of defaulting to Pydantic V1 models

.flake8

@@ -12,4 +12,5 @@ exclude =
     .nox,
     .pytest_cache
     __pycache__,
+    _pydantic.py,
     __init__.py

.python-version

@@ -0,0 +1 @@
+3.8

Makefile

@@ -34,6 +34,9 @@ format: install
 test: install
 	$(POETRY) run -- nox -s test $(if $(PYTHON),--python=$(PYTHON),) -- $(ARGS)
 
+test_fast: install
+	$(POETRY) run -- nox -s test $(if $(PYTHON),--python=$(PYTHON),) -- -m 'not slow' $(ARGS)
+
 benchmark: install
 	python -m benchmark.benchmark_merge_chunks
 
@@ -47,5 +50,6 @@ help:
 	@echo 'lint                         - run linters'
 	@echo '-- TESTS --'
 	@echo 'test                         - run unit tests'
+	@echo 'test_fast                    - run unit tests without slow tests'
 	@echo 'test PYTHON=<python_version> - run unit tests with the specific python version'
 	@echo 'benchmark                    - run benchmarks'

README.md

@@ -122,6 +122,7 @@ install PyCharm>=2023.2 with [built-in Black support](https://blog.jetbrains.com
 |Variable|Default|Description|
 |---|---|---|
 |DIAL_SDK_LOG|WARNING|DIAL SDK log level|
+|PYDANTIC_V2|False|When `True` and Pydantic V2 is installed, DIAL SDK classes for requests/responses will be based on Pydantic V2 `BaseModel`. Otherwise, they will be based on Pydantic V1 `BaseModel`.|
 
 ## Lint
 

aidial_sdk/_errors.py

@@ -1,9 +1,9 @@
 from fastapi import HTTPException, Request
 from fastapi.responses import JSONResponse
 
+from aidial_sdk._pydantic import ValidationError
 from aidial_sdk.exceptions import HTTPException as DIALException
 from aidial_sdk.exceptions import InvalidRequestError
-from aidial_sdk.pydantic_v1 import ValidationError
 
 
 def pydantic_validation_exception_handler(

aidial_sdk/_pydantic/__init__.py

@@ -0,0 +1,50 @@
+"""
+The entry point for all Pydantic definition that unifies v1 and v2 APIs for SDK internals.
+
+It's private, since it expected that the SDK client
+will either import `aidial_sdk.pydantic_v1` or `pydantic`.
+
+This is the only place where `pydantic` imports
+are allowed in the DIAL SDK package.
+"""
+
+from typing import TYPE_CHECKING
+
+import pydantic
+
+from aidial_sdk.utils.env import env_bool
+
+INSTALLED_PYDANTIC_V2 = pydantic.VERSION.startswith("2.")
+USE_PYDANTIC_V2 = env_bool("PYDANTIC_V2", False)
+PYDANTIC_V2 = INSTALLED_PYDANTIC_V2 and USE_PYDANTIC_V2
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+    from pydantic import ConfigDict as ConfigDict
+    from pydantic import Field, ValidationError
+    from pydantic import field_validator as validator
+    from pydantic import model_validator
+    from pydantic._internal._model_construction import ModelMetaclass
+    from pydantic.fields import FieldInfo
+else:
+
+    if PYDANTIC_V2:
+        from pydantic import BaseModel, ConfigDict, Field, ValidationError
+        from pydantic import field_validator as validator
+        from pydantic import model_validator
+        from pydantic._internal._model_construction import ModelMetaclass
+        from pydantic.fields import FieldInfo
+    else:
+        from pydantic.v1 import BaseModel, Field, validator
+
+        try:
+            from pydantic.v1.main import ModelMetaclass
+        except ImportError:
+            from pydantic.main import ModelMetaclass
+        from pydantic.v1 import ValidationError, root_validator
+        from pydantic.v1.fields import FieldInfo
+
+        def _fail(*args, **kwargs):
+            raise ImportError("ConfigDict is only supported in Pydantic v2")
+
+        ConfigDict = _fail

aidial_sdk/_pydantic/_compat.py

@@ -0,0 +1,97 @@
+"""
+The module provide the basic Pydantic BaseModel extended
+with `model_dump` method mimicking the one from Pydantic V2.
+
+All SDK models inherit from this class.
+
+It proves to be useful since
+1. `model_dump` method is used extensively in the SDK,
+2. the SDK client may call this method on SDK models even if the client uses Pydantic V1.
+"""
+
+from datetime import date, datetime
+from typing import Any, Dict, Iterable, Mapping, Optional, Set, Union, cast
+
+from typing_extensions import Literal
+
+import aidial_sdk._pydantic as pydantic
+from aidial_sdk._pydantic import PYDANTIC_V2
+
+_IncEx = Union[Set[int], Set[str], Dict[int, Any], Dict[str, Any], None]
+
+
+class BaseModel(pydantic.BaseModel):
+    if not PYDANTIC_V2:
+        # we define aliases for some of the new pydantic v2 methods so
+        # that we can just document these methods without having to specify
+        # a specific pydantic version as some users may not know which
+        # pydantic version they are currently using
+
+        def model_dump(
+            self,
+            *,
+            mode: Union[Literal["json", "python"], str] = "python",
+            include: _IncEx = None,
+            exclude: _IncEx = None,
+            by_alias: bool = False,
+            exclude_unset: bool = False,
+            exclude_defaults: bool = False,
+            exclude_none: bool = False,
+            round_trip: bool = False,
+            warnings: Union[bool, Literal["none", "warn", "error"]] = True,
+            context: Optional[Dict[str, Any]] = None,
+            serialize_as_any: bool = False,
+        ) -> Dict[str, Any]:
+            if mode not in {"json", "python"}:
+                raise ValueError("mode must be either 'json' or 'python'")
+            if round_trip is not False:
+                raise ValueError("round_trip is only supported in Pydantic v2")
+            if warnings is not True:
+                raise ValueError("warnings is only supported in Pydantic v2")
+            if context is not None:
+                raise ValueError("context is only supported in Pydantic v2")
+            if serialize_as_any is not False:
+                raise ValueError(
+                    "serialize_as_any is only supported in Pydantic v2"
+                )
+            dumped = super().dict(  # pyright: ignore[reportDeprecated]
+                include=include,
+                exclude=exclude,
+                by_alias=by_alias,
+                exclude_unset=exclude_unset,
+                exclude_defaults=exclude_defaults,
+                exclude_none=exclude_none,
+            )
+
+            return (
+                cast(Dict[str, Any], _json_safe(dumped))
+                if mode == "json"
+                else dumped
+            )
+
+
+def _json_safe(data: object) -> object:
+    """Translates a mapping / sequence recursively in the same fashion
+    as `pydantic` v2's `model_dump(mode="json")`.
+    """
+    if isinstance(data, Mapping):
+        return {
+            _json_safe(key): _json_safe(value) for key, value in data.items()
+        }
+
+    if isinstance(data, Iterable) and not isinstance(
+        data, (str, bytes, bytearray)
+    ):
+        return [_json_safe(item) for item in data]
+
+    if isinstance(data, (datetime, date)):
+        return data.isoformat()
+
+    return data
+
+
+def model_validator(*, mode: Literal["before", "after"]) -> Any:
+    if PYDANTIC_V2:
+        return pydantic.model_validator(mode=mode)
+    else:
+        return pydantic.root_validator(pre=(mode == "before"))  # type: ignore

aidial_sdk/_pydantic/_model_config.py

@@ -0,0 +1,144 @@
+"""
+Helper classes that unify model configuration between Pydantic v1 and v2.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Dict, Optional, Type, TypeVar
+
+from pydantic import BaseModel
+
+from aidial_sdk._pydantic import PYDANTIC_V2
+
+_Model = TypeVar("_Model", bound=BaseModel)
+
+
+class ModelConfigWrapper:
+    _model_config: ModelConfigBase
+
+    def __init__(self, model_config: ModelConfigBase):
+        self._model_config = model_config
+
+    def __getitem__(self, field: str) -> Any:
+        return self._model_config.get_field(field, None)
+
+    def __setitem__(self, field: str, value: Any) -> None:
+        self._model_config.set_field(field, value)
+
+    def post_process_schema(
+        self, on_schema: Callable[[Dict[str, Any]], None]
+    ) -> None:
+        attr_name = self._model_config.schema_extra_field
+        old_schema_extra = self[attr_name]
+
+        def _schema_extra(
+            schema: Dict[str, Any], model: Type[BaseModel]
+        ) -> None:
+            if old_schema_extra:
+                old_schema_extra(schema, model)
+            on_schema(schema)
+
+        self[attr_name] = _schema_extra
+
+    @classmethod
+    def create(
+        cls, base_cls: Optional[Type[_Model]], namespace: Dict[str, Any]
+    ) -> ModelConfigWrapper:
+        if PYDANTIC_V2:
+            return cls(_ConfigV2.create(base_cls, namespace))
+        else:
+            return cls(_ConfigV1.create(base_cls, namespace))
+
+
+class ModelConfigBase(ABC):
+    @abstractmethod
+    def set_field(self, field: str, value: Any) -> None:
+        pass
+
+    @abstractmethod
+    def get_field(self, field: str, default: Any) -> Any:
+        pass
+
+    @property
+    @abstractmethod
+    def schema_extra_field(self) -> str:
+        pass
+
+    @classmethod
+    @abstractmethod
+    def create(
+        cls, base_cls: Optional[Type[_Model]], namespace: Dict[str, Any]
+    ) -> ModelConfigBase:
+        pass
+
+
+class _ConfigV1(ModelConfigBase):
+    config_cls: type
+
+    def __init__(self, config_cls: type):
+        self.config_cls = config_cls
+
+    def set_field(self, field: str, value: Any) -> None:
+        setattr(self.config_cls, field, value)
+
+    def get_field(self, field: str, default: Any) -> Any:
+        return getattr(self.config_cls, field, default)
+
+    @property
+    def schema_extra_field(self) -> str:
+        return "schema_extra"
+
+    @classmethod
+    def create(
+        cls, base_cls: Optional[Type[_Model]], namespace: Dict[str, Any]
+    ) -> ModelConfigBase:
+        if (config_cls := namespace.get("Config")) is None:
+            conf_base_cls = (
+                None if base_cls is None else getattr(base_cls, "Config", None)
+            )
+
+            config_cls = type("Config", (conf_base_cls or object,), {})
+
+            if module := namespace.get("__module__"):
+                config_cls.__module__ = module
+            if qualname := namespace.get("__qualname__"):
+                config_cls.__qualname__ = f"{qualname}.{config_cls.__name__}"
+
+            namespace["Config"] = config_cls
+
+        return cls(config_cls)
+
+
+class _ConfigV2(ModelConfigBase):
+    model_config: dict
+
+    def __init__(self, model_config: dict):
+        self.model_config = model_config
+
+    def set_field(self, field: str, value: Any) -> None:
+        self.model_config[field] = value
+
+    def get_field(self, field: str, default: Any) -> Any:
+        return self.model_config.get(field, default)
+
+    @property
+    def schema_extra_field(self) -> str:
+        return "json_schema_extra"
+
+    @classmethod
+    def create(
+        cls, base_cls: Optional[Type[_Model]], namespace: Dict[str, Any]
+    ) -> ModelConfigBase:
+        base_model_config = (
+            {} if base_cls is None else getattr(base_cls, "model_config", {})
+        )
+
+        curr_model_config = namespace.get("model_config") or {}
+
+        model_config = namespace["model_config"] = {
+            **base_model_config,
+            **curr_model_config,
+        }
+
+        return cls(model_config)

aidial_sdk/application.py

@@ -12,6 +12,8 @@
     fastapi_exception_handler,
     pydantic_validation_exception_handler,
 )
+from aidial_sdk._pydantic import ValidationError
+from aidial_sdk._pydantic._compat import BaseModel
 from aidial_sdk.chat_completion.base import ChatCompletion
 from aidial_sdk.chat_completion.request import Request as ChatCompletionRequest
 from aidial_sdk.chat_completion.response import (
@@ -26,7 +28,6 @@
 from aidial_sdk.embeddings.request import Request as EmbeddingsRequest
 from aidial_sdk.exceptions import HTTPException as DIALException
 from aidial_sdk.header_propagator import HeaderPropagator
-from aidial_sdk.pydantic_v1 import BaseModel, ValidationError
 from aidial_sdk.telemetry.types import TelemetryConfig
 from aidial_sdk.utils._reflection import get_method_implementation
 from aidial_sdk.utils.log_config import LogConfig
@@ -38,7 +39,7 @@
     to_streaming_response,
 )
 
-logging.config.dictConfig(LogConfig().dict())
+logging.config.dictConfig(LogConfig().model_dump())
 
 RequestType = TypeVar("RequestType", bound=FromRequestMixin)
 
@@ -207,7 +208,7 @@ async def _handler(original_request: Request) -> Response:
             if isinstance(response, dict):
                 response_json = response
             elif isinstance(response, BaseModel):
-                response_json = response.dict()
+                response_json = response.model_dump()
             else:
                 raise ValueError(
                     f"Unexpected response type from {endpoint}: {type(response)}"
@@ -291,7 +292,7 @@ async def _handler(original_request: Request):
                 EmbeddingsRequest, original_request, deployment_id
             )
             response = await impl.embeddings(request)
-            response_json = response.dict()
+            response_json = response.model_dump()
             return JSONResponse(content=response_json)
 
         return _handler