vellum-ai 0.9.16rc2__py3-none-any.whl → 0.9.16rc4__py3-none-any.whl

vellum/plugins/pydantic.py

@@ -0,0 +1,74 @@
+from typing import Any, Dict, Literal, Optional, Tuple, Union
+
+from pydantic.plugin import (
+    PydanticPluginProtocol,
+    SchemaKind,
+    SchemaTypePath,
+    ValidateJsonHandlerProtocol,
+    ValidatePythonHandlerProtocol,
+    ValidateStringsHandlerProtocol,
+)
+from pydantic_core import CoreSchema
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+
+
+# https://docs.pydantic.dev/2.8/concepts/plugins/#build-a-plugin
+class OnValidatePython(ValidatePythonHandlerProtocol):
+    tracked_descriptors: Dict[str, BaseDescriptor] = {}
+
+    def on_enter(
+        self,
+        input: Any,
+        *,
+        strict: Optional[bool] = None,
+        from_attributes: Optional[bool] = None,
+        context: Optional[Dict[str, Any]] = None,
+        self_instance: Optional[Any] = None,
+        allow_partial: Union[bool, Literal["off", "on", "trailing-strings"]] = False,
+    ) -> None:
+        if not isinstance(input, dict):
+            return
+
+        self.tracked_descriptors = {}
+
+        for key, value in input.items():
+            if isinstance(value, BaseDescriptor):
+                self.tracked_descriptors[key] = value
+                # TODO: This does not yet work for descriptors that map to more complex types
+                # https://app.shortcut.com/vellum/story/4636
+                input[key] = value.types[0]()
+
+    def on_success(self, result: Any) -> None:
+        if self.tracked_descriptors:
+            frozen = result.model_config.get("frozen")
+            if frozen:
+                result.model_config["frozen"] = False
+
+            for key, value in self.tracked_descriptors.items():
+                setattr(result, key, value)
+
+            if frozen:
+                result.model_config["frozen"] = True
+
+            self.tracked_descriptors = {}
+
+
+class VellumPydanticPlugin(PydanticPluginProtocol):
+    def new_schema_validator(
+        self,
+        schema: CoreSchema,
+        schema_type: Any,
+        schema_type_path: SchemaTypePath,
+        schema_kind: SchemaKind,
+        config: Any,
+        plugin_settings: Dict[str, Any],
+    ) -> Tuple[
+        Union[ValidatePythonHandlerProtocol, None],
+        Union[ValidateJsonHandlerProtocol, None],
+        Union[ValidateStringsHandlerProtocol, None],
+    ]:
+        return OnValidatePython(), None, None
+
+
+pydantic_plugin = VellumPydanticPlugin()

vellum/plugins/utils.py

@@ -0,0 +1,19 @@
+from pydantic.plugin import _loader as _pydantic_plugin_loader
+
+from vellum.plugins.pydantic import pydantic_plugin
+
+_loaded = False
+
+
+def load_runtime_plugins() -> None:
+    global _loaded
+    if _loaded:
+        return
+    _loaded = True
+
+    # TODO: This is a hack to get the Vellum plugin to load. We're supposed to use
+    # pyproject.toml, but I couldn't figure out after an hour
+    # https://app.shortcut.com/vellum/story/4635
+    _pydantic_plugin_loader.get_plugins()
+    if _pydantic_plugin_loader._plugins is not None:
+        _pydantic_plugin_loader._plugins["vellum_plugin"] = pydantic_plugin