pixeltable 0.2.26__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/func/signature.py

@@ -1,45 +1,45 @@
 from __future__ import annotations
 
 import dataclasses
-import enum
 import inspect
-import json
 import logging
 import typing
-from typing import Any, Callable, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, ClassVar
 
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
 
+if TYPE_CHECKING:
+    from pixeltable import exprs
+
 _logger = logging.getLogger('pixeltable')
 
 
 @dataclasses.dataclass
 class Parameter:
     name: str
-    col_type: Optional[ts.ColumnType]  # None for variable parameters
+    col_type: ts.ColumnType | None  # None for variable parameters
     kind: inspect._ParameterKind
     # for some reason, this needs to precede is_batched in the dataclass definition,
     # otherwise Python complains that an argument with a default is followed by an argument without a default
-    default: Any = inspect.Parameter.empty  # default value for the parameter
+    default: 'exprs.Literal' | None = None  # default value for the parameter
     is_batched: bool = False  # True if the parameter is a batched parameter (eg, Batch[dict])
 
     def __post_init__(self) -> None:
-        # make sure that default is json-serializable and of the correct type
-        if self.default is inspect.Parameter.empty or self.default is None:
-            return
-        try:
-            _ = json.dumps(self.default)
-        except TypeError:
-            raise excs.Error(f'Default value for parameter {self.name} is not JSON-serializable: {str(self.default)}')
-        if self.col_type is not None:
-            try:
-                self.col_type.validate_literal(self.default)
-            except TypeError as e:
-                raise excs.Error(f'Default value for parameter {self.name}: {str(e)}')
+        from pixeltable import exprs
+
+        if self.default is not None:
+            if self.col_type is None:
+                raise excs.Error(f'Cannot have a default value for variable parameter {self.name!r}')
+            if not isinstance(self.default, exprs.Literal):
+                raise excs.Error(f'Default value for parameter {self.name!r} is not a constant')
+            if not self.col_type.is_supertype_of(self.default.col_type):
+                raise excs.Error(
+                    f'Default value for parameter {self.name!r} is not of type {self.col_type!r}: {self.default}'
+                )
 
     def has_default(self) -> bool:
-        return self.default is not inspect.Parameter.empty
+        return self.default is not None
 
     def as_dict(self) -> dict[str, Any]:
         return {
@@ -47,27 +47,29 @@ class Parameter:
             'col_type': self.col_type.as_dict() if self.col_type is not None else None,
             'kind': self.kind.name,
             'is_batched': self.is_batched,
-            'has_default': self.has_default(),
-            'default': self.default if self.has_default() else None,
+            'default': None if self.default is None else self.default.as_dict(),
         }
 
     @classmethod
     def from_dict(cls, d: dict[str, Any]) -> Parameter:
-        has_default = d['has_default']
-        if has_default:
-            default = d['default']
-        else:
-            default = inspect.Parameter.empty
+        from pixeltable import exprs
+
+        assert d['default'] is None or isinstance(d['default'], dict), d
+        default = None if d['default'] is None else exprs.Literal.from_dict(d['default'])
         return cls(
             name=d['name'],
             col_type=ts.ColumnType.from_dict(d['col_type']) if d['col_type'] is not None else None,
             kind=getattr(inspect.Parameter, d['kind']),
             is_batched=d['is_batched'],
-            default=default
+            default=default,
         )
 
     def to_py_param(self) -> inspect.Parameter:
-        return inspect.Parameter(self.name, self.kind, default=self.default)
+        py_default = self.default.val if self.default is not None else inspect.Parameter.empty
+        return inspect.Parameter(self.name, self.kind, default=py_default)
+
+    def __hash__(self) -> int:
+        return hash((self.name, self.col_type, self.kind, self.default, self.is_batched))
 
 
 T = typing.TypeVar('T')
@@ -80,9 +82,30 @@ class Signature:
 
     - self.is_batched: return type is a Batch[...] type
     """
-    SPECIAL_PARAM_NAMES = ['group_by', 'order_by']
 
-    def __init__(self, return_type: ts.ColumnType, parameters: list[Parameter], is_batched: bool = False):
+    SPECIAL_PARAM_NAMES: ClassVar[list[str]] = ['group_by', 'order_by']
+    SYSTEM_PARAM_NAMES: ClassVar[list[str]] = ['_runtime_ctx']
+
+    return_type: ts.ColumnType
+    is_batched: bool
+    parameters: dict[str, Parameter]  # name -> Parameter
+    parameters_by_pos: list[Parameter]  # ordered by position in the signature
+    constant_parameters: list[Parameter]  # parameters that are not batched
+    batched_parameters: list[Parameter]  # parameters that are batched
+    required_parameters: list[Parameter]  # parameters that do not have a default value
+
+    # the names of recognized system parameters in the signature; these are excluded from self.parameters
+    system_parameters: list[str]
+
+    py_signature: inspect.Signature
+
+    def __init__(
+        self,
+        return_type: ts.ColumnType,
+        parameters: list[Parameter],
+        is_batched: bool = False,
+        system_parameters: list[str] | None = None,
+    ):
         assert isinstance(return_type, ts.ColumnType)
         self.return_type = return_type
         self.is_batched = is_batched
@@ -92,6 +115,7 @@ class Signature:
         self.constant_parameters = [p for p in parameters if not p.is_batched]
         self.batched_parameters = [p for p in parameters if p.is_batched]
         self.required_parameters = [p for p in parameters if not p.has_default()]
+        self.system_parameters = system_parameters if system_parameters is not None else []
         self.py_signature = inspect.Signature([p.to_py_param() for p in self.parameters_by_pos])
 
     def get_return_type(self) -> ts.ColumnType:
@@ -111,6 +135,74 @@ class Signature:
         parameters = [Parameter.from_dict(param_dict) for param_dict in d['parameters']]
         return cls(ts.ColumnType.from_dict(d['return_type']), parameters, d['is_batched'])
 
+    def is_consistent_with(self, other: Signature) -> bool:
+        """
+        Returns True if this signature is consistent with the other signature.
+        S is consistent with T if we could safely replace S by T in any call where S is used. Specifically:
+        (i) S.return_type is a supertype of T.return_type
+        (ii) For each parameter p in S, there is a parameter q in T such that:
+            - p and q have the same name and kind
+            - q.col_type is a supertype of p.col_type
+        (iii) For each *required* parameter q in T, there is a parameter p in S with the same name (in which
+            case the kinds and types must also match, by condition (ii)).
+        """
+        # Check (i)
+        if not self.get_return_type().is_supertype_of(other.get_return_type(), ignore_nullable=True):
+            return False
+
+        # Check (ii)
+        for param_name, param in self.parameters.items():
+            if param_name not in other.parameters:
+                return False
+            other_param = other.parameters[param_name]
+            if (
+                param.kind != other_param.kind
+                or (param.col_type is None) != (other_param.col_type is None)  # this can happen if they are varargs
+                or (
+                    param.col_type is not None
+                    and not other_param.col_type.is_supertype_of(param.col_type, ignore_nullable=True)
+                )
+            ):
+                return False
+
+        # Check (iii)
+        for other_param in other.required_parameters:  # noqa: SIM110
+            if other_param.name not in self.parameters:
+                return False
+
+        return True
+
+    def validate_args(self, bound_args: dict[str, 'exprs.Expr' | None], context: str = '') -> None:
+        if context:
+            context = f' ({context})'
+
+        for param_name, arg in bound_args.items():
+            assert param_name in self.parameters, f'{param_name!r} not in {list(self.parameters.keys())}'
+            param = self.parameters[param_name]
+            is_var_param = param.kind in {inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD}
+            if is_var_param:
+                continue
+            assert param.col_type is not None
+
+            if arg is None:
+                raise excs.Error(f'Parameter {param_name!r}{context}: invalid argument')
+
+            # Check that the argument is consistent with the expected parameter type, with the allowance that
+            # non-nullable parameters can still accept nullable arguments (since in that event, FunctionCall.eval()
+            # detects the Nones and skips evaluation).
+            if not (
+                param.col_type.is_supertype_of(arg.col_type, ignore_nullable=True)
+                # TODO: this is a hack to allow JSON columns to be passed to functions that accept scalar
+                # types. It's necessary to avoid littering notebooks with `apply(str)` calls or equivalent.
+                # (Previously, this wasn't necessary because `is_supertype_of()` was improperly implemented.)
+                # We need to think through the right way to handle this scenario.
+                or (arg.col_type.is_json_type() and param.col_type.is_scalar_type())
+            ):
+                raise excs.Error(
+                    f'Parameter {param_name!r}{context}: argument type {arg.col_type} does not'
+                    f' match parameter type {param.col_type}'
+                )
+
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, Signature):
             return False
@@ -124,6 +216,9 @@ class Signature:
                 return False
         return True
 
+    def __hash__(self) -> int:
+        return hash((self.return_type, self.parameters))
+
     def __str__(self) -> str:
         param_strs: list[str] = []
         for p in self.parameters.values():
@@ -132,21 +227,21 @@ class Signature:
             elif p.kind == inspect.Parameter.VAR_KEYWORD:
                 param_strs.append(f'**{p.name}')
             else:
-                param_strs.append(f'{p.name}: {str(p.col_type)}')
-        return f'({", ".join(param_strs)}) -> {str(self.get_return_type())}'
+                param_strs.append(f'{p.name}: pxt.{p.col_type}')
+        return f'({", ".join(param_strs)}) -> pxt.{self.get_return_type()}'
 
     @classmethod
-    def _infer_type(cls, annotation: Optional[type]) -> tuple[Optional[ts.ColumnType], Optional[bool]]:
+    def _infer_type(cls, annotation: type | None) -> tuple[ts.ColumnType | None, bool | None]:
         """Returns: (column type, is_batched) or (None, ...) if the type cannot be inferred"""
         if annotation is None:
             return (None, None)
-        py_type: Optional[type] = None
+        py_type: type | None = None
         is_batched = False
         if typing.get_origin(annotation) == typing.Annotated:
             type_args = typing.get_args(annotation)
             if len(type_args) == 2 and type_args[1] == 'pxt-batch':
                 # this is our Batch
-                assert typing.get_origin(type_args[0]) == list
+                assert typing.get_origin(type_args[0]) is list
                 is_batched = True
                 py_type = typing.get_args(type_args[0])[0]
         if py_type is None:
@@ -156,55 +251,89 @@ class Signature:
 
     @classmethod
     def create_parameters(
-            cls, py_fn: Optional[Callable] = None, py_params: Optional[list[inspect.Parameter]] = None,
-            param_types: Optional[list[ts.ColumnType]] = None
+        cls,
+        py_fn: Callable | None = None,
+        py_params: list[inspect.Parameter] | None = None,
+        param_types: list[ts.ColumnType] | None = None,
+        type_substitutions: dict | None = None,
+        is_cls_method: bool = False,
     ) -> list[Parameter]:
+        """Ignores parameters starting with '_'."""
+        from pixeltable import exprs
+
         assert (py_fn is None) != (py_params is None)
         if py_fn is not None:
             sig = inspect.signature(py_fn)
             py_params = list(sig.parameters.values())
         parameters: list[Parameter] = []
 
+        if type_substitutions is None:
+            type_substitutions = {}
+
         for idx, param in enumerate(py_params):
+            if is_cls_method and idx == 0:
+                continue  # skip 'self' or 'cls' parameter
+            if param.name in cls.SYSTEM_PARAM_NAMES:
+                continue  # skip system parameters
+            if param.name.startswith('_'):
+                raise excs.Error(f"{param.name!r}: parameters starting with '_' are reserved")
             if param.name in cls.SPECIAL_PARAM_NAMES:
-                raise excs.Error(f"'{param.name}' is a reserved parameter name")
-            if param.kind == inspect.Parameter.VAR_POSITIONAL or param.kind == inspect.Parameter.VAR_KEYWORD:
+                raise excs.Error(f'{param.name!r} is a reserved parameter name')
+            if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
                 parameters.append(Parameter(param.name, col_type=None, kind=param.kind))
                 continue
 
             # check non-var parameters for name collisions and default value compatibility
             if param_types is not None:
                 if idx >= len(param_types):
-                    raise excs.Error(f'Missing type for parameter {param.name}')
+                    raise excs.Error(f'Missing type for parameter {param.name!r}')
                 param_type = param_types[idx]
                 is_batched = False
             else:
-                param_type, is_batched = cls._infer_type(param.annotation)
+                # Look up the substitution for param.annotation, defaulting to param.annotation if there is none
+                py_type = type_substitutions.get(param.annotation, param.annotation)
+                param_type, is_batched = cls._infer_type(py_type)
                 if param_type is None:
-                    raise excs.Error(f'Cannot infer pixeltable type for parameter {param.name}')
+                    raise excs.Error(f'Cannot infer pixeltable type for parameter {param.name!r}')
 
-            parameters.append(Parameter(
-                param.name, col_type=param_type, kind=param.kind, is_batched=is_batched, default=param.default))
+            default = None if param.default is inspect.Parameter.empty else exprs.Expr.from_object(param.default)
+            if not (default is None or isinstance(default, exprs.Literal)):
+                raise excs.Error(f'Default value for parameter {param.name!r} must be a constant')
+
+            parameters.append(
+                Parameter(param.name, col_type=param_type, kind=param.kind, is_batched=is_batched, default=default)
+            )
 
         return parameters
 
     @classmethod
     def create(
-        cls, py_fn: Callable,
-        param_types: Optional[list[ts.ColumnType]] = None,
-        return_type: Optional[ts.ColumnType] = None
+        cls,
+        py_fn: Callable,
+        param_types: list[ts.ColumnType] | None = None,
+        return_type: ts.ColumnType | None = None,
+        type_substitutions: dict | None = None,
+        is_cls_method: bool = False,
     ) -> Signature:
         """Create a signature for the given Callable.
         Infer the parameter and return types, if none are specified.
         Raises an exception if the types cannot be inferred.
         """
-        parameters = cls.create_parameters(py_fn=py_fn, param_types=param_types)
+        if type_substitutions is None:
+            type_substitutions = {}
+
+        parameters = cls.create_parameters(
+            py_fn=py_fn, param_types=param_types, is_cls_method=is_cls_method, type_substitutions=type_substitutions
+        )
         sig = inspect.signature(py_fn)
         if return_type is None:
-            return_type, return_is_batched = cls._infer_type(sig.return_annotation)
+            # Look up the substitution for sig.return_annotation, defaulting to return_annotation if there is none
+            py_type = type_substitutions.get(sig.return_annotation, sig.return_annotation)
+            return_type, return_is_batched = cls._infer_type(py_type)
             if return_type is None:
                 raise excs.Error('Cannot infer pixeltable return type')
         else:
             _, return_is_batched = cls._infer_type(sig.return_annotation)
+        system_params = [param_name for param_name in sig.parameters if param_name in cls.SYSTEM_PARAM_NAMES]
 
-        return Signature(return_type, parameters, return_is_batched)
+        return Signature(return_type, parameters, return_is_batched, system_parameters=system_params)

pixeltable/func/tools.py

@@ -0,0 +1,164 @@
+import json
+import uuid
+from typing import TYPE_CHECKING, Any, Callable, TypeVar
+
+import pydantic
+
+from pixeltable import exceptions as excs, type_system as ts
+
+from .function import Function
+from .signature import Parameter
+from .udf import udf
+
+if TYPE_CHECKING:
+    from pixeltable import exprs
+
+
+# The Tool and Tools classes are containers that hold Pixeltable UDFs and related metadata, so that they can be
+# realized as LLM tools. They are implemented as Pydantic models in order to provide a canonical way of converting
+# to JSON, via the Pydantic `model_serializer` interface. In this way, they can be passed directly as UDF
+# parameters as described in the `pixeltable.tools` and `pixeltable.tool` docstrings.
+#
+# (The dataclass dict serializer is insufficiently flexible for this purpose: `Tool` contains a member of type
+# `Function`, which is not natively JSON-serializable; Pydantic provides a way of customizing its default
+# serialization behavior, whereas dataclasses do not.)
+
+
+class Tool(pydantic.BaseModel):
+    # Allow arbitrary types so that we can include a Pixeltable function in the schema.
+    # We will implement a model_serializer to ensure the Tool model can be serialized.
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
+    fn: Function
+    name: str | None = None
+    description: str | None = None
+
+    @property
+    def parameters(self) -> dict[str, Parameter]:
+        return self.fn.signature.parameters
+
+    @pydantic.model_serializer
+    def ser_model(self) -> dict[str, Any]:
+        return {
+            'name': self.name or self.fn.name,
+            'description': self.description or self.fn.comment(),
+            'parameters': {
+                'type': 'object',
+                'properties': {param.name: param.col_type._to_json_schema() for param in self.parameters.values()},
+            },
+            'required': [param.name for param in self.parameters.values() if not param.col_type.nullable],
+            'additionalProperties': False,  # TODO Handle kwargs?
+        }
+
+    # The output of `tool_calls` must be a dict in standardized tool invocation format:
+    # {tool_name: [{'args': {name1: value1, name2: value2, ...}}, ...], ...}
+    def invoke(self, tool_calls: 'exprs.Expr') -> 'exprs.Expr':
+        import pixeltable.functions as pxtf
+
+        func_name = self.name or self.fn.name
+        return pxtf.map(tool_calls[func_name]['*'], lambda x: self.__invoke_kwargs(x.args))
+
+    def __invoke_kwargs(self, kwargs: 'exprs.Expr') -> 'exprs.FunctionCall':
+        kwargs = {param.name: self.__extract_tool_arg(param, kwargs) for param in self.parameters.values()}
+        return self.fn(**kwargs)
+
+    def __extract_tool_arg(self, param: Parameter, kwargs: 'exprs.Expr') -> 'exprs.FunctionCall':
+        if param.col_type.is_string_type():
+            return _extract_str_tool_arg(kwargs, param_name=param.name)
+        if param.col_type.is_int_type():
+            return _extract_int_tool_arg(kwargs, param_name=param.name)
+        if param.col_type.is_float_type():
+            return _extract_float_tool_arg(kwargs, param_name=param.name)
+        if param.col_type.is_bool_type():
+            return _extract_bool_tool_arg(kwargs, param_name=param.name)
+        if param.col_type.is_json_type():
+            return _extract_json_tool_arg(kwargs, param_name=param.name)
+        if param.col_type.is_uuid_type():
+            return _extract_uuid_tool_arg(kwargs, param_name=param.name)
+        raise AssertionError(param.col_type)
+
+
+class ToolChoice(pydantic.BaseModel):
+    auto: bool
+    required: bool
+    tool: str | None
+    parallel_tool_calls: bool
+
+
+class Tools(pydantic.BaseModel):
+    tools: list[Tool]
+
+    @pydantic.model_serializer
+    def ser_model(self) -> list[dict[str, Any]]:
+        return [tool.ser_model() for tool in self.tools]
+
+    # `tool_calls` must be in standardized tool invocation format:
+    # {tool_name: {'args': {name1: value1, name2: value2, ...}}, ...}
+    def _invoke(self, tool_calls: 'exprs.Expr') -> 'exprs.InlineDict':
+        from pixeltable import exprs
+
+        return exprs.InlineDict({tool.name or tool.fn.name: tool.invoke(tool_calls) for tool in self.tools})
+
+    def choice(
+        self,
+        auto: bool = False,
+        required: bool = False,
+        tool: str | Function | None = None,
+        parallel_tool_calls: bool = True,
+    ) -> ToolChoice:
+        if sum([auto, required, tool is not None]) != 1:
+            raise excs.Error('Exactly one of `auto`, `required`, or `tool` must be specified.')
+        tool_name: str | None = None
+        if tool is not None:
+            try:
+                tool_obj = next(
+                    t
+                    for t in self.tools
+                    if (isinstance(tool, Function) and t.fn == tool)
+                    or (isinstance(tool, str) and (t.name or t.fn.name) == tool)
+                )
+                tool_name = tool_obj.name or tool_obj.fn.name
+            except StopIteration:
+                raise excs.Error(f'That tool is not in the specified list of tools: {tool}') from None
+        return ToolChoice(auto=auto, required=required, tool=tool_name, parallel_tool_calls=parallel_tool_calls)
+
+
+@udf
+def _extract_str_tool_arg(kwargs: dict[str, Any], param_name: str) -> str | None:
+    return _extract_arg(str, kwargs, param_name)
+
+
+@udf
+def _extract_int_tool_arg(kwargs: dict[str, Any], param_name: str) -> int | None:
+    return _extract_arg(int, kwargs, param_name)
+
+
+@udf
+def _extract_float_tool_arg(kwargs: dict[str, Any], param_name: str) -> float | None:
+    return _extract_arg(float, kwargs, param_name)
+
+
+@udf
+def _extract_bool_tool_arg(kwargs: dict[str, Any], param_name: str) -> bool | None:
+    return _extract_arg(bool, kwargs, param_name)
+
+
+@udf
+def _extract_json_tool_arg(kwargs: dict[str, Any], param_name: str) -> ts.Json | None:
+    if param_name in kwargs:
+        return json.loads(kwargs[param_name])
+    return None
+
+
+@udf
+def _extract_uuid_tool_arg(kwargs: dict[str, Any], param_name: str) -> uuid.UUID | None:
+    return _extract_arg(uuid.UUID, kwargs, param_name)
+
+
+T = TypeVar('T')
+
+
+def _extract_arg(eval_fn: Callable[[Any], T], kwargs: dict[str, Any], param_name: str) -> T | None:
+    if param_name in kwargs:
+        return eval_fn(kwargs[param_name])
+    return None