vellum-ai 0.9.16rc2__py3-none-any.whl → 0.10.0__py3-none-any.whl

vellum/workflows/types/tests/test_utils.py

@@ -0,0 +1,76 @@
+import pytest
+from typing import ClassVar, Generic, List, TypeVar, Union
+
+from vellum.workflows.outputs.base import BaseOutputs
+from vellum.workflows.references.output import OutputReference
+from vellum.workflows.types.utils import get_class_attr_names, infer_types
+
+
+class ExampleClass:
+    alpha: str
+    beta = 3
+    gamma: Union[str, int]
+    epsilon = OutputReference(
+        name="epsilon",
+        types=(List[str],),
+        instance=None,
+        outputs_class=BaseOutputs,
+    )
+    zeta: ClassVar[str]
+    eta: List[str]
+
+
+T = TypeVar("T")
+
+
+class ExampleGenericClass(Generic[T]):
+    delta: T
+
+
+class ExampleInheritedClass(ExampleClass):
+    theta: int
+
+
+@pytest.mark.parametrize(
+    "cls, attr_name, expected_type",
+    [
+        (ExampleClass, "alpha", (str,)),
+        (ExampleClass, "beta", (int,)),
+        (ExampleClass, "gamma", (str, int)),
+        (ExampleGenericClass, "delta", ()),
+        (ExampleGenericClass[str], "delta", (str,)),
+        (ExampleClass, "epsilon", (List[str],)),
+        (ExampleClass, "zeta", (str,)),
+        (ExampleClass, "eta", (List[str],)),
+        (ExampleInheritedClass, "theta", (int,)),
+        (ExampleInheritedClass, "alpha", (str,)),
+        (ExampleInheritedClass, "beta", (int,)),
+    ],
+    ids=[
+        "str",
+        "int",
+        "str_or_int",
+        "generic_blank",
+        "generic_str",
+        "descriptor",
+        "class_var",
+        "list_str",
+        "inherited_int",
+        "inherited_parent_annotation",
+        "inherited_parent_class_var",
+    ],
+)
+def test_infer_types(cls, attr_name, expected_type):
+    assert infer_types(cls, attr_name) == expected_type
+
+
+@pytest.mark.parametrize(
+    "cls, expected_attr_names",
+    [
+        (ExampleClass, {"alpha", "beta", "gamma", "epsilon", "zeta", "eta"}),
+        (ExampleGenericClass, {"delta"}),
+        (ExampleInheritedClass, {"alpha", "beta", "gamma", "epsilon", "zeta", "eta", "theta"}),
+    ],
+)
+def test_class_attr_names(cls, expected_attr_names):
+    assert get_class_attr_names(cls) == expected_attr_names

vellum/workflows/types/utils.py

@@ -0,0 +1,164 @@
+from copy import deepcopy
+from datetime import datetime
+import importlib
+from typing import (
+    Any,
+    ClassVar,
+    Dict,
+    Generic,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+from vellum import ArrayVellumValue, ArrayVellumValueRequest, ChatMessagePromptBlock
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.types.core import Json, SpecialGenericAlias, UnderGenericAlias, UnionGenericAlias
+
+_T = TypeVar("_T")
+_LHS = TypeVar("_LHS")
+_RHS = TypeVar("_RHS")
+
+LOCAL_NS = {
+    "Json": Json,
+    "ArrayVellumValueRequest": ArrayVellumValueRequest,
+    "ArrayVellumValue": ArrayVellumValue,
+    "ChatMessagePromptBlock": ChatMessagePromptBlock,
+}
+
+
+def resolve_types(value: Union[BaseDescriptor[_T], _T]) -> Tuple[Type[_T], ...]:
+    if isinstance(value, BaseDescriptor):
+        return value.types
+
+    return (value.__class__,)
+
+
+def resolve_combined_types(
+    lhs: Union[BaseDescriptor[_LHS], _LHS], rhs: Union[BaseDescriptor[_RHS], _RHS]
+) -> Tuple[Union[Type[_LHS], Type[_RHS]], ...]:
+    lhs_types = resolve_types(lhs)
+    rhs_types = resolve_types(rhs)
+
+    unique_results = set(lhs_types) | set(rhs_types)
+
+    return tuple(unique_results)
+
+
+def infer_types(object_: Type, attr_name: str, localns: Optional[Dict[str, Any]] = None) -> Tuple[Type, ...]:
+    try:
+        class_ = object_
+        type_var_mapping = {}
+        if isinstance(object_, UnderGenericAlias):
+            origin = get_origin(object_)
+            if origin and Generic in origin.__bases__:
+                class_ = origin
+                args = get_args(object_)
+                type_var_mapping = {t: a for t, a in zip(origin.__parameters__, args)}
+
+        type_hints = get_type_hints(class_, localns=LOCAL_NS if localns is None else {**LOCAL_NS, **localns})
+        if attr_name in type_hints:
+            type_hint = type_hints[attr_name]
+            if get_origin(type_hint) is ClassVar:
+                return get_args(type_hint)
+            if isinstance(type_hint, type):
+                return (type_hint,)
+            if isinstance(type_hint, UnionGenericAlias):
+                if get_origin(type_hint) is Union:
+                    return get_args(type_hint)
+            if isinstance(type_hint, UnderGenericAlias):
+                return (type_hint,)
+            if isinstance(type_hint, SpecialGenericAlias):
+                return (type_hint,)
+            if isinstance(type_hint, TypeVar):
+                if type_hint in type_var_mapping:
+                    return (type_var_mapping[type_hint],)
+                return type_hint.__constraints__
+
+        for base in reversed(class_.__mro__):
+            class_attributes = vars(base)
+            if attr_name in class_attributes:
+                class_attribute = class_attributes[attr_name]
+                return resolve_types(class_attribute)
+
+        raise AttributeError(f"Failed to infer type from attribute {attr_name} on {object_.__name__}")
+    except TypeError:
+        raise AttributeError(
+            f"Found 3.9+ typing syntax for field '{attr_name}' on class '{object_.__name__}' – {object_.__annotations__[attr_name]}. Type annotations must be compatible with python version 3.8. "  # noqa: E501
+        )
+
+
+def get_class_attr_names(cls: Type) -> Set[str]:
+    # gets type-annotated attributes `foo: int`
+    type_annotated_attributes: Set[str] = set()
+
+    # gets attributes declared `foo = 1`
+    class_attributes: Set[str] = set()
+
+    for base in reversed(cls.__mro__):
+        ann = base.__dict__.get("__annotations__", {})
+        type_annotated_attributes.update(ann.keys())
+
+        base_vars = vars(base).keys()
+        class_attributes.update(base_vars)
+
+    # combine and filter out private attributes
+    return {a for a in list(set(class_attributes) | set(type_annotated_attributes)) if not a.startswith("_")}
+
+
+def deepcopy_with_exclusions(
+    obj: _T,
+    memo: Any,
+    exclusions: Optional[Dict[str, Any]] = None,
+) -> _T:
+    cls = obj.__class__
+    new_instance = cls.__new__(cls)
+    new_instance.__dict__.update(obj.__dict__)
+
+    exclusions = exclusions or {}
+
+    for key, value in obj.__dict__.items():
+        if key in exclusions:
+            continue
+        new_instance.__dict__[key] = deepcopy(value, memo)
+
+    for key, value in exclusions.items():
+        new_instance.__dict__[key] = value
+
+    return new_instance
+
+
+def get_class_by_qualname(qualname: str) -> Type:
+    module_name, class_name = qualname.rsplit(".", 1)
+    module = importlib.import_module(module_name)
+    imported_class = getattr(module, class_name)
+    if not isinstance(imported_class, type):
+        raise ValueError(f"Class {qualname} is not a valid type")
+
+    return imported_class
+
+
+def datetime_now() -> datetime:
+    """
+    There's a race condition between freezegun and pydantic that causes `PydanticSchemaGenerationError`
+    because freezegun monkey-patches `datetime.now` to return a frozen datetime. This helper provides
+    an alternative way to facilitate testing that doesn't rely on freezegun.
+    """
+
+    return datetime.now()
+
+
+def get_original_base(cls: Type) -> Type:
+    if not hasattr(cls, "__orig_bases__"):
+        return type(None)
+
+    # in Python 3.12, there is `from types import get_original_bases`, making this future proof
+    # https://docs.python.org/3/library/types.html#types.get_original_bases
+    return cls.__orig_bases__[0]  # type: ignore[attr-defined]

vellum/workflows/utils/names.py

@@ -0,0 +1,13 @@
+import re
+
+
+def pascal_to_title_case(pascal_str: str) -> str:
+    # Insert space before each capital letter (except the first one)
+    title_case_str = re.sub(r"(?<!^)(?=[A-Z])", " ", pascal_str)
+
+    # Return the result in title case
+    return title_case_str.title()
+
+
+def snake_to_title_case(snake_str: str) -> str:
+    return pascal_to_title_case(snake_str.replace("_", " "))

vellum/workflows/utils/tests/test_names.py

@@ -0,0 +1,15 @@
+import pytest
+
+from vellum.workflows.utils.names import pascal_to_title_case
+
+
+@pytest.mark.parametrize(
+    ["input_str", "expected"],
+    [
+        ("MyPascalCaseString", "My Pascal Case String"),
+        ("AnotherPascalCaseString", "Another Pascal Case String"),
+    ],
+)
+def test_pascal_to_title_case(input_str, expected):
+    actual = pascal_to_title_case(input_str)
+    assert actual == expected

vellum/workflows/utils/tests/test_vellum_variables.py

@@ -0,0 +1,25 @@
+import pytest
+from typing import List, Optional
+
+from vellum import ChatMessage, SearchResult
+
+from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
+
+
+@pytest.mark.parametrize(
+    "type_, expected",
+    [
+        (str, "STRING"),
+        (Optional[str], "STRING"),
+        (int, "NUMBER"),
+        (Optional[int], "NUMBER"),
+        (float, "NUMBER"),
+        (Optional[float], "NUMBER"),
+        (List[ChatMessage], "CHAT_HISTORY"),
+        (Optional[List[ChatMessage]], "CHAT_HISTORY"),
+        (List[SearchResult], "SEARCH_RESULTS"),
+        (Optional[List[SearchResult]], "SEARCH_RESULTS"),
+    ],
+)
+def test_primitive_type_to_vellum_variable_type(type_, expected):
+    assert primitive_type_to_vellum_variable_type(type_) == expected

vellum/workflows/utils/vellum_variables.py

@@ -0,0 +1,81 @@
+from typing import List, Tuple, Type, Union, get_args, get_origin
+
+from vellum import (
+    ChatMessage,
+    ChatMessageRequest,
+    FunctionCall,
+    FunctionCallRequest,
+    SearchResult,
+    SearchResultRequest,
+    VellumAudio,
+    VellumAudioRequest,
+    VellumError,
+    VellumErrorRequest,
+    VellumImage,
+    VellumImageRequest,
+    VellumValue,
+    VellumValueRequest,
+    VellumVariableType,
+)
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+
+
+def primitive_type_to_vellum_variable_type(type_: Union[Type, BaseDescriptor]) -> VellumVariableType:
+    """Converts a python primitive to a VellumVariableType"""
+    if isinstance(type_, BaseDescriptor):
+        # Ignore None because those just make types optional
+        types = [t for t in type_.types if t is not type(None)]
+
+        # default to JSON for typevars where the types is empty tuple
+        if len(types) == 0:
+            return "JSON"
+
+        if len(types) != 1:
+            raise ValueError(f"Expected Descriptor to only have one type, got {types}")
+
+        type_ = type_.types[0]
+
+    if _is_type_optionally_equal(type_, str):
+        return "STRING"
+    elif _is_type_optionally_in(type_, (int, float)):
+        return "NUMBER"
+    elif _is_type_optionally_in(type_, (FunctionCall, FunctionCallRequest)):
+        return "FUNCTION_CALL"
+    elif _is_type_optionally_in(type_, (VellumImage, VellumImageRequest)):
+        return "IMAGE"
+    elif _is_type_optionally_in(type_, (VellumAudio, VellumAudioRequest)):
+        return "AUDIO"
+    elif _is_type_optionally_in(type_, (VellumError, VellumErrorRequest)):
+        return "ERROR"
+    elif _is_type_optionally_in(type_, (List[ChatMessage], List[ChatMessageRequest])):
+        return "CHAT_HISTORY"
+    elif _is_type_optionally_in(type_, (List[SearchResult], List[SearchResultRequest])):
+        return "SEARCH_RESULTS"
+    elif _is_type_optionally_in(type_, (List[VellumValue], List[VellumValueRequest])):
+        return "ARRAY"
+
+    return "JSON"
+
+
+def _is_type_optionally_equal(type_: Type, target_type: Type) -> bool:
+    if type_ == target_type:
+        return True
+
+    origin = get_origin(type_)
+    if origin is not Union:
+        return False
+
+    args = get_args(type_)
+    if len(args) != 2:
+        return False
+
+    source_type, none_type = args
+    if none_type is not type(None):
+        return False
+
+    return _is_type_optionally_equal(source_type, target_type)
+
+
+def _is_type_optionally_in(type_: Type, target_types: Tuple[Type, ...]) -> bool:
+    return any(_is_type_optionally_equal(type_, target_type) for target_type in target_types)

vellum/workflows/vellum_client.py

@@ -0,0 +1,18 @@
+import os
+from typing import Optional
+
+from vellum import Vellum, VellumEnvironment
+
+
+def create_vellum_client(api_key: Optional[str] = None) -> Vellum:
+    if api_key is None:
+        api_key = os.getenv("VELLUM_API_KEY", default="")
+
+    return Vellum(
+        api_key=api_key,
+        environment=VellumEnvironment(
+            default=os.getenv("VELLUM_DEFAULT_API_URL", os.getenv("VELLUM_API_URL", "https://api.vellum.ai")),
+            documents=os.getenv("VELLUM_DOCUMENTS_API_URL", os.getenv("VELLUM_API_URL", "https://documents.vellum.ai")),
+            predict=os.getenv("VELLUM_PREDICT_API_URL", os.getenv("VELLUM_API_URL", "https://predict.vellum.ai")),
+        ),
+    )

vellum/workflows/workflows/__init__.py

@@ -0,0 +1,5 @@
+from .base import BaseWorkflow
+
+__all__ = [
+    "BaseWorkflow",
+]