vellum-ai 0.14.72__py3-none-any.whl → 0.14.74__py3-none-any.whl

vellum/workflows/utils/functions.py

@@ -26,27 +26,27 @@ type_map = {
 }
 
 
-def _compile_annotation(annotation: Optional[Any], defs: dict[str, Any]) -> dict:
+def compile_annotation(annotation: Optional[Any], defs: dict[str, Any]) -> dict:
     if annotation is None:
         return {"type": "null"}
 
     if get_origin(annotation) is Union:
-        return {"anyOf": [_compile_annotation(a, defs) for a in get_args(annotation)]}
+        return {"anyOf": [compile_annotation(a, defs) for a in get_args(annotation)]}
 
     if get_origin(annotation) is dict:
         _, value_type = get_args(annotation)
-        return {"type": "object", "additionalProperties": _compile_annotation(value_type, defs)}
+        return {"type": "object", "additionalProperties": compile_annotation(value_type, defs)}
 
     if get_origin(annotation) is list:
         item_type = get_args(annotation)[0]
-        return {"type": "array", "items": _compile_annotation(item_type, defs)}
+        return {"type": "array", "items": compile_annotation(item_type, defs)}
 
     if dataclasses.is_dataclass(annotation):
         if annotation.__name__ not in defs:
             properties = {}
             required = []
             for field in dataclasses.fields(annotation):
-                properties[field.name] = _compile_annotation(field.type, defs)
+                properties[field.name] = compile_annotation(field.type, defs)
                 if field.default is dataclasses.MISSING:
                     required.append(field.name)
                 else:
@@ -61,7 +61,7 @@ def _compile_annotation(annotation: Optional[Any], defs: dict[str, Any]) -> dict
             for field_name, field in annotation.model_fields.items():
                 # Mypy is incorrect here, the `annotation` attribute is defined on `FieldInfo`
                 field_annotation = field.annotation  # type: ignore[attr-defined]
-                properties[field_name] = _compile_annotation(field_annotation, defs)
+                properties[field_name] = compile_annotation(field_annotation, defs)
                 if field.default is PydanticUndefined:
                     required.append(field_name)
                 else:
@@ -88,7 +88,7 @@ def _compile_default_value(default: Any) -> Any:
     return default
 
 
-def _compile_deployment_workflow_input(input_var: Any) -> dict[str, Any]:
+def _compile_workflow_deployment_input(input_var: Any) -> dict[str, Any]:
     """
     Converts a deployment workflow input variable to a JSON schema type definition.
     """
@@ -115,7 +115,7 @@ def compile_function_definition(function: Callable) -> FunctionDefinition:
     required = []
     defs: dict[str, Any] = {}
     for param in signature.parameters.values():
-        properties[param.name] = _compile_annotation(param.annotation, defs)
+        properties[param.name] = compile_annotation(param.annotation, defs)
         if param.default is inspect.Parameter.empty:
             required.append(param.name)
         else:
@@ -132,7 +132,7 @@ def compile_function_definition(function: Callable) -> FunctionDefinition:
     )
 
 
-def compile_workflow_function_definition(workflow_class: Type["BaseWorkflow"]) -> FunctionDefinition:
+def compile_inline_workflow_function_definition(workflow_class: Type["BaseWorkflow"]) -> FunctionDefinition:
     """
     Converts a base workflow class into our Vellum-native FunctionDefinition type.
     """
@@ -148,7 +148,7 @@ def compile_workflow_function_definition(workflow_class: Type["BaseWorkflow"]) -
         if name.startswith("__"):
             continue
 
-        properties[name] = _compile_annotation(field_type, defs)
+        properties[name] = compile_annotation(field_type, defs)
 
         # Check if the field has a default value
         if name not in vars_inputs_class:
@@ -168,7 +168,7 @@ def compile_workflow_function_definition(workflow_class: Type["BaseWorkflow"]) -
     )
 
 
-def compile_deployment_workflow_function_definition(
+def compile_workflow_deployment_function_definition(
     deployment_config: Dict[str, str],
     vellum_client: Vellum,
 ) -> FunctionDefinition:
@@ -193,7 +193,7 @@ def compile_deployment_workflow_function_definition(
     required = []
 
     for input_var in input_variables:
-        properties[input_var.key] = _compile_deployment_workflow_input(input_var)
+        properties[input_var.key] = _compile_workflow_deployment_input(input_var)
 
         if input_var.required and input_var.default is None:
             required.append(input_var.key)

vellum/workflows/utils/pydantic_schema.py

@@ -0,0 +1,38 @@
+import inspect
+from typing import Any, Dict
+
+from pydantic import BaseModel
+
+from vellum.workflows.utils.functions import compile_annotation
+
+
+def normalize_json(schema_input: Any) -> Any:
+    """
+    Recursively normalize JSON data by converting Pydantic models to JSON schema.
+
+    This function processes dictionaries recursively to find and convert any
+    Pydantic model classes or instances to their JSON schema representation.
+
+    Args:
+        schema_input: Can be a Pydantic model class, instance, dict, or any other value
+
+    Returns:
+        Normalized JSON data with Pydantic models converted to JSON schema
+    """
+    if isinstance(schema_input, dict):
+        return {key: normalize_json(value) for key, value in schema_input.items()}
+
+    if inspect.isclass(schema_input) and issubclass(schema_input, BaseModel):
+        defs: Dict[str, Any] = {}
+        result = compile_annotation(schema_input, defs)
+
+        if "$ref" in result and defs:
+            ref_name = result["$ref"].split("/")[-1]
+            if ref_name in defs:
+                return defs[ref_name]
+
+        return result
+    elif isinstance(schema_input, BaseModel):
+        return {key: normalize_json(getattr(schema_input, key)) for key in schema_input.__class__.model_fields.keys()}
+    else:
+        return schema_input

vellum/workflows/utils/tests/test_functions.py

@@ -12,9 +12,9 @@ from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.state.base import BaseState
 from vellum.workflows.utils.functions import (
-    compile_deployment_workflow_function_definition,
     compile_function_definition,
-    compile_workflow_function_definition,
+    compile_inline_workflow_function_definition,
+    compile_workflow_deployment_function_definition,
 )
 
 
@@ -306,7 +306,7 @@ def test_compile_function_definition__lambda():
     )
 
 
-def test_compile_workflow_function_definition():
+def test_compile_inline_workflow_function_definition():
     class MyNode(BaseNode):
         pass
 
@@ -314,7 +314,7 @@ def test_compile_workflow_function_definition():
         graph = MyNode
 
     # WHEN compiling the function
-    compiled_function = compile_workflow_function_definition(MyWorkflow)
+    compiled_function = compile_inline_workflow_function_definition(MyWorkflow)
 
     # THEN it should return the compiled function definition
     assert compiled_function == FunctionDefinition(
@@ -323,7 +323,7 @@ def test_compile_workflow_function_definition():
     )
 
 
-def test_compile_workflow_function_definition__docstring():
+def test_compile_inline_workflow_function_definition__docstring():
     class MyNode(BaseNode):
         pass
 
@@ -335,7 +335,7 @@ def test_compile_workflow_function_definition__docstring():
         graph = MyNode
 
     # WHEN compiling the function
-    compiled_function = compile_workflow_function_definition(MyWorkflow)
+    compiled_function = compile_inline_workflow_function_definition(MyWorkflow)
 
     # THEN it should return the compiled function definition
     assert compiled_function == FunctionDefinition(
@@ -345,7 +345,7 @@ def test_compile_workflow_function_definition__docstring():
     )
 
 
-def test_compile_workflow_function_definition__all_args():
+def test_compile_inline_workflow_function_definition__all_args():
     class MyInputs(BaseInputs):
         a: str
         b: int
@@ -361,7 +361,7 @@ def test_compile_workflow_function_definition__all_args():
         graph = MyNode
 
     # WHEN compiling the workflow
-    compiled_function = compile_workflow_function_definition(MyWorkflow)
+    compiled_function = compile_inline_workflow_function_definition(MyWorkflow)
 
     # THEN it should return the compiled function definition
     assert compiled_function == FunctionDefinition(
@@ -381,7 +381,7 @@ def test_compile_workflow_function_definition__all_args():
     )
 
 
-def test_compile_workflow_function_definition__unions():
+def test_compile_inline_workflow_function_definition__unions():
     # GIVEN a workflow with a union
     class MyInputs(BaseInputs):
         a: Union[str, int]
@@ -393,7 +393,7 @@ def test_compile_workflow_function_definition__unions():
         graph = MyNode
 
     # WHEN compiling the workflow
-    compiled_function = compile_workflow_function_definition(MyWorkflow)
+    compiled_function = compile_inline_workflow_function_definition(MyWorkflow)
 
     # THEN it should return the compiled function definition
     assert compiled_function == FunctionDefinition(
@@ -406,7 +406,7 @@ def test_compile_workflow_function_definition__unions():
     )
 
 
-def test_compile_workflow_function_definition__optionals():
+def test_compile_inline_workflow_function_definition__optionals():
     class MyInputs(BaseInputs):
         a: str
         b: Optional[str]
@@ -421,7 +421,7 @@ def test_compile_workflow_function_definition__optionals():
         graph = MyNode
 
     # WHEN compiling the workflow
-    compiled_function = compile_workflow_function_definition(MyWorkflow)
+    compiled_function = compile_inline_workflow_function_definition(MyWorkflow)
 
     # THEN it should return the compiled function definition
     assert compiled_function == FunctionDefinition(
@@ -440,7 +440,7 @@ def test_compile_workflow_function_definition__optionals():
     )
 
 
-def test_compile_deployment_workflow_function_definition__just_name():
+def test_compile_workflow_deployment_function_definition__just_name():
     # GIVEN a mock Vellum client and deployment
     mock_client = Mock()
     mock_release = Mock()
@@ -451,7 +451,7 @@ def test_compile_deployment_workflow_function_definition__just_name():
     deployment_config = {"deployment": "my_deployment", "release_tag": "latest"}
 
     # WHEN compiling the deployment workflow function
-    compiled_function = compile_deployment_workflow_function_definition(deployment_config, mock_client)
+    compiled_function = compile_workflow_deployment_function_definition(deployment_config, mock_client)
 
     # THEN it should return the compiled function definition (same structure as function test)
     assert compiled_function == FunctionDefinition(
@@ -461,7 +461,7 @@ def test_compile_deployment_workflow_function_definition__just_name():
     )
 
 
-def test_compile_deployment_workflow_function_definition__all_args():
+def test_compile_workflow_deployment_function_definition__all_args():
     # GIVEN a mock Vellum client and deployment
     mock_client = Mock()
     mock_release = Mock()
@@ -497,7 +497,7 @@ def test_compile_deployment_workflow_function_definition__all_args():
     deployment_config = {"deployment": "my_deployment", "release_tag": "latest"}
 
     # WHEN compiling the deployment workflow function
-    compiled_function = compile_deployment_workflow_function_definition(deployment_config, mock_client)
+    compiled_function = compile_workflow_deployment_function_definition(deployment_config, mock_client)
 
     # THEN it should return the compiled function definition
     assert compiled_function == FunctionDefinition(
@@ -524,7 +524,7 @@ def test_compile_deployment_workflow_function_definition__all_args():
     )
 
 
-def test_compile_deployment_workflow_function_definition__defaults():
+def test_compile_workflow_deployment_function_definition__defaults():
     # GIVEN a mock Vellum client and deployment
     mock_client = Mock()
     mock_release = Mock()
@@ -565,7 +565,7 @@ def test_compile_deployment_workflow_function_definition__defaults():
     deployment_config = {"deployment": "my_deployment", "release_tag": "latest"}
 
     # WHEN compiling the deployment workflow function
-    compiled_function = compile_deployment_workflow_function_definition(deployment_config, mock_client)
+    compiled_function = compile_workflow_deployment_function_definition(deployment_config, mock_client)
 
     # THEN it should return the compiled function definition with proper default handling
     assert compiled_function == FunctionDefinition(

{vellum_ai-0.14.72.dist-info → vellum_ai-0.14.74.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.14.72
+Version: 0.14.74
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0