pixeltable 0.2.28__py3-none-any.whl → 0.2.30__py3-none-any.whl

pixeltable/func/udf.py

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
-from typing import Any, Callable, Optional, overload
+from typing import Any, Callable, Optional, Sequence, overload
 
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
 
 from .callable_function import CallableFunction
-from .expr_template_function import ExprTemplateFunction
+from .expr_template_function import ExprTemplateFunction, ExprTemplate
 from .function import Function
 from .function_registry import FunctionRegistry
 from .globals import validate_symbol_path
@@ -21,13 +21,14 @@ def udf(decorated_fn: Callable) -> Function: ...
 # Decorator schema invoked with parentheses: @pxt.udf(**kwargs)
 @overload
 def udf(
-        *,
-        batch_size: Optional[int] = None,
-        substitute_fn: Optional[Callable] = None,
-        is_method: bool = False,
-        is_property: bool = False,
-        _force_stored: bool = False
-) -> Callable[[Callable], Function]: ...
+    *,
+    batch_size: Optional[int] = None,
+    substitute_fn: Optional[Callable] = None,
+    is_method: bool = False,
+    is_property: bool = False,
+    type_substitutions: Optional[Sequence[dict]] = None,
+    _force_stored: bool = False
+) -> Callable[[Callable], CallableFunction]: ...
 
 
 def udf(*args, **kwargs):
@@ -52,6 +53,7 @@ def udf(*args, **kwargs):
         substitute_fn = kwargs.pop('substitute_fn', None)
         is_method = kwargs.pop('is_method', None)
         is_property = kwargs.pop('is_property', None)
+        type_substitutions = kwargs.pop('type_substitutions', None)
         force_stored = kwargs.pop('_force_stored', False)
         if len(kwargs) > 0:
             raise excs.Error(f'Invalid @udf decorator kwargs: {", ".join(kwargs.keys())}')
@@ -65,6 +67,7 @@ def udf(*args, **kwargs):
                 substitute_fn=substitute_fn,
                 is_method=is_method,
                 is_property=is_property,
+                type_substitutions=type_substitutions,
                 force_stored=force_stored
             )
 
@@ -79,9 +82,10 @@ def make_function(
     substitute_fn: Optional[Callable] = None,
     is_method: bool = False,
     is_property: bool = False,
+    type_substitutions: Optional[Sequence[dict]] = None,
     function_name: Optional[str] = None,
     force_stored: bool = False
-) -> Function:
+) -> CallableFunction:
     """
     Constructs a `CallableFunction` from the specified parameters.
     If `substitute_fn` is specified, then `decorated_fn`
@@ -104,25 +108,43 @@ def make_function(
     # Display name to use for error messages
     errmsg_name = function_name if function_path is None else function_path
 
-    sig = Signature.create(decorated_fn, param_types, return_type)
-
-    # batched functions must have a batched return type
-    # TODO: remove 'Python' from the error messages when we have full inference with Annotated types
-    if batch_size is not None and not sig.is_batched:
-        raise excs.Error(f'{errmsg_name}(): batch_size is specified; Python return type must be a `Batch`')
-    if batch_size is not None and len(sig.batched_parameters) == 0:
-        raise excs.Error(f'{errmsg_name}(): batch_size is specified; at least one Python parameter must be `Batch`')
-    if batch_size is None and len(sig.batched_parameters) > 0:
-        raise excs.Error(f'{errmsg_name}(): batched parameters in udf, but no `batch_size` given')
-
-    if is_method and is_property:
-        raise excs.Error(f'Cannot specify both `is_method` and `is_property` (in function `{function_name}`)')
-    if is_property and len(sig.parameters) != 1:
-        raise excs.Error(
-            f"`is_property=True` expects a UDF with exactly 1 parameter, but `{function_name}` has {len(sig.parameters)}"
-        )
-    if (is_method or is_property) and function_path is None:
-        raise excs.Error('Stored functions cannot be declared using `is_method` or `is_property`')
+    signatures: list[Signature]
+    if type_substitutions is None:
+        sig = Signature.create(decorated_fn, param_types, return_type)
+
+        # batched functions must have a batched return type
+        # TODO: remove 'Python' from the error messages when we have full inference with Annotated types
+        if batch_size is not None and not sig.is_batched:
+            raise excs.Error(f'{errmsg_name}(): batch_size is specified; Python return type must be a `Batch`')
+        if batch_size is not None and len(sig.batched_parameters) == 0:
+            raise excs.Error(f'{errmsg_name}(): batch_size is specified; at least one Python parameter must be `Batch`')
+        if batch_size is None and len(sig.batched_parameters) > 0:
+            raise excs.Error(f'{errmsg_name}(): batched parameters in udf, but no `batch_size` given')
+
+        if is_method and is_property:
+            raise excs.Error(f'Cannot specify both `is_method` and `is_property` (in function `{function_name}`)')
+        if is_property and len(sig.parameters) != 1:
+            raise excs.Error(
+                f"`is_property=True` expects a UDF with exactly 1 parameter, but `{function_name}` has {len(sig.parameters)}"
+            )
+        if (is_method or is_property) and function_path is None:
+            raise excs.Error('Stored functions cannot be declared using `is_method` or `is_property`')
+
+        signatures = [sig]
+    else:
+        if function_path is None:
+            raise excs.Error(
+                f'{errmsg_name}(): type substitutions can only be used with module UDFs (not locally defined UDFs)'
+            )
+        if batch_size is not None:
+            raise excs.Error(f'{errmsg_name}(): type substitutions cannot be used with batched functions')
+        if is_method is not None or is_property is not None:
+            # TODO: Support this for `is_method`?
+            raise excs.Error(f'{errmsg_name}(): type substitutions cannot be used with `is_method` or `is_property`')
+        signatures = [
+            Signature.create(decorated_fn, param_types, return_type, type_substitutions=subst)
+            for subst in type_substitutions
+        ]
 
     if substitute_fn is None:
         py_fn = decorated_fn
@@ -132,8 +154,8 @@ def make_function(
         py_fn = substitute_fn
 
     result = CallableFunction(
-        signature=sig,
-        py_fn=py_fn,
+        signatures=signatures,
+        py_fns=[py_fn] * len(signatures),  # All signatures share the same Python function
         self_path=function_path,
         self_name=function_name,
         batch_size=batch_size,
@@ -171,12 +193,12 @@ def expr_udf(*args: Any, **kwargs: Any) -> Any:
         import pixeltable.exprs as exprs
         var_exprs = [exprs.Variable(param.name, param.col_type) for param in sig.parameters.values()]
         # call the function with the parameter expressions to construct an Expr with parameters
-        template = py_fn(*var_exprs)
-        assert isinstance(template, exprs.Expr)
-        sig.return_type = template.col_type
+        expr = py_fn(*var_exprs)
+        assert isinstance(expr, exprs.Expr)
+        sig.return_type = expr.col_type
         if function_path is not None:
             validate_symbol_path(function_path)
-        return ExprTemplateFunction(template, sig, self_path=function_path, name=py_fn.__name__)
+        return ExprTemplateFunction([ExprTemplate(expr, sig)], self_path=function_path, name=py_fn.__name__)
 
     if len(args) == 1:
         assert len(kwargs) == 0 and callable(args[0])

pixeltable/functions/__init__.py

@@ -1,7 +1,7 @@
 from pixeltable.utils.code import local_public_names
 
-from . import (anthropic, audio, fireworks, gemini, huggingface, image, json, llama_cpp, mistralai, ollama, openai,
-               string, timestamp, together, video, vision, whisper)
+from . import (anthropic, audio, fireworks, gemini, huggingface, image, json, llama_cpp, math, mistralai, ollama,
+               openai, string, timestamp, together, video, vision, whisper)
 from .globals import *
 
 __all__ = local_public_names(__name__, exclude=['globals']) + local_public_names(globals.__name__)

pixeltable/functions/anthropic.py

@@ -10,7 +10,8 @@ from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Union
 import tenacity
 
 import pixeltable as pxt
-from pixeltable import env
+from pixeltable import env, exprs
+from pixeltable.func import Tools
 from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
@@ -47,7 +48,7 @@ def messages(
     system: Optional[str] = None,
     temperature: Optional[float] = None,
     tool_choice: Optional[list[dict]] = None,
-    tools: Optional[dict] = None,
+    tools: Optional[list[dict]] = None,
     top_k: Optional[int] = None,
     top_p: Optional[float] = None,
 ) -> dict:
@@ -77,6 +78,21 @@ def messages(
         >>> msgs = [{'role': 'user', 'content': tbl.prompt}]
         ... tbl['response'] = messages(msgs, model='claude-3-haiku-20240307')
     """
+    if tools is not None:
+        # Reformat `tools` into Anthropic format
+        tools = [
+            {
+                'name': tool['name'],
+                'description': tool['description'],
+                'input_schema': {
+                    'type': 'object',
+                    'properties': tool['parameters']['properties'],
+                    'required': tool['required'],
+                },
+            }
+            for tool in tools
+        ]
+
     return _retry(_anthropic_client().messages.create)(
         messages=messages,
         model=model,
@@ -92,6 +108,24 @@ def messages(
     ).dict()
 
 
+def invoke_tools(tools: Tools, response: exprs.Expr) -> exprs.InlineDict:
+    """Converts an Anthropic response dict to Pixeltable tool invocation format and calls `tools._invoke()`."""
+    return tools._invoke(_anthropic_response_to_pxt_tool_calls(response))
+
+
+@pxt.udf
+def _anthropic_response_to_pxt_tool_calls(response: dict) -> Optional[dict]:
+    anthropic_tool_calls = [r for r in response['content'] if r['type'] == 'tool_use']
+    if len(anthropic_tool_calls) > 0:
+        return {
+            tool_call['name']: {
+                'args': tool_call['input']
+            }
+            for tool_call in anthropic_tool_calls
+        }
+    return None
+
+
 _T = TypeVar('_T')
 
 

pixeltable/functions/globals.py

@@ -1,6 +1,7 @@
 import builtins
 from typing import _GenericAlias  # type: ignore[attr-defined]
 from typing import Optional, Union
+import typing
 
 import sqlalchemy as sql
 
@@ -16,23 +17,24 @@ def cast(expr: exprs.Expr, target_type: Union[ts.ColumnType, type, _GenericAlias
     return expr
 
 
-@func.uda(
-    update_types=[ts.IntType(nullable=True)], value_type=ts.IntType(nullable=False),
-    allows_window=True, requires_order_by=False)
-class sum(func.Aggregator):
+T = typing.TypeVar('T')
+
+
+@func.uda(allows_window=True, type_substitutions=({T: Optional[int]}, {T: Optional[float]}))  # type: ignore[misc]
+class sum(func.Aggregator, typing.Generic[T]):
     """Sums the selected integers or floats."""
     def __init__(self):
-        self.sum: Optional[int] = None
+        self.sum: T = None
 
-    def update(self, val: Optional[int]) -> None:
+    def update(self, val: T) -> None:
         if val is None:
             return
         if self.sum is None:
             self.sum = val
         else:
-            self.sum += val
+            self.sum += val  # type: ignore[operator]
 
-    def value(self) -> Union[int, float]:
+    def value(self) -> T:
         return self.sum
 
 
@@ -43,12 +45,22 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     return sql.sql.func.sum(val)
 
 
-@func.uda(update_types=[ts.IntType(nullable=True)], value_type=ts.IntType(), allows_window=True, requires_order_by=False)
-class count(func.Aggregator):
+@func.uda(
+    allows_window=True,
+    # Allow counting non-null values of any type
+    # TODO: I couldn't include "Array" because we don't have a way to represent a generic array (of arbitrary dimension).
+    # TODO: should we have an "Any" type that can be used here?
+    type_substitutions=tuple(
+        {T: Optional[t]}  # type: ignore[misc]
+        for t in (ts.String, ts.Int, ts.Float, ts.Bool, ts.Timestamp,
+                  ts.Json, ts.Image, ts.Video, ts.Audio, ts.Document)
+    ),
+)
+class count(func.Aggregator, typing.Generic[T]):
     def __init__(self):
         self.count = 0
 
-    def update(self, val: Optional[int]) -> None:
+    def update(self, val: T) -> None:
         if val is not None:
             self.count += 1
 
@@ -62,74 +74,82 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
 
 
 @func.uda(
-    update_types=[ts.IntType(nullable=True)], value_type=ts.IntType(nullable=True), allows_window=True,
-    requires_order_by=False)
-class min(func.Aggregator):
+    allows_window=True,
+    type_substitutions=tuple({T: Optional[t]} for t in (str, int, float, bool, ts.Timestamp))  # type: ignore[misc]
+)
+class min(func.Aggregator, typing.Generic[T]):
     def __init__(self):
-        self.val: Optional[int] = None
+        self.val: T = None
 
-    def update(self, val: Optional[int]) -> None:
+    def update(self, val: T) -> None:
         if val is None:
             return
         if self.val is None:
             self.val = val
         else:
-            self.val = builtins.min(self.val, val)
+            self.val = builtins.min(self.val, val)  # type: ignore[call-overload]
 
-    def value(self) -> Optional[int]:
+    def value(self) -> T:
         return self.val
 
 
 @min.to_sql
 def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
+    if val.type.python_type == bool:
+        # TODO: min/max aggregation of booleans is not supported in Postgres (but it is in Python).
+        # Right now we simply force the computation to be done in Python; we might consider implementing an alternate
+        # way of doing it in SQL. (min/max of booleans is simply logical and/or, respectively.)
+        return None
     return sql.sql.func.min(val)
 
 
 @func.uda(
-    update_types=[ts.IntType(nullable=True)], value_type=ts.IntType(nullable=True), allows_window=True,
-    requires_order_by=False)
-class max(func.Aggregator):
+    allows_window=True,
+    type_substitutions=tuple({T: Optional[t]} for t in (str, int, float, bool, ts.Timestamp))  # type: ignore[misc]
+)
+class max(func.Aggregator, typing.Generic[T]):
     def __init__(self):
-        self.val: Optional[int] = None
+        self.val: T = None
 
-    def update(self, val: Optional[int]) -> None:
+    def update(self, val: T) -> None:
         if val is None:
             return
         if self.val is None:
             self.val = val
         else:
-            self.val = builtins.max(self.val, val)
+            self.val = builtins.max(self.val, val)  # type: ignore[call-overload]
 
-    def value(self) -> Optional[int]:
+    def value(self) -> T:
         return self.val
 
 
 @max.to_sql
 def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
+    if val.type.python_type == bool:
+        # TODO: see comment in @min.to_sql.
+        return None
     return sql.sql.func.max(val)
 
 
-@func.uda(
-    update_types=[ts.IntType(nullable=True)], value_type=ts.FloatType(nullable=True), allows_window=False,
-    requires_order_by=False)
-class mean(func.Aggregator):
+@func.uda(type_substitutions=({T: Optional[int]}, {T: Optional[float]}))  # type: ignore[misc]
+class mean(func.Aggregator, typing.Generic[T]):
     def __init__(self):
-        self.sum: Optional[int] = None
+        self.sum: T = None
         self.count = 0
 
-    def update(self, val: Optional[int]) -> None:
+    def update(self, val: T) -> None:
         if val is None:
             return
         if self.sum is None:
             self.sum = val
         else:
-            self.sum += val
+            self.sum += val  # type: ignore[operator]
         self.count += 1
 
-    def value(self) -> Optional[float]:
+    def value(self) -> Optional[float]:  # Always a float
         if self.count == 0:
             return None
-        return self.sum / self.count
+        return self.sum / self.count  # type: ignore[operator]
 
 
 @mean.to_sql

pixeltable/functions/json.py

@@ -16,20 +16,15 @@ import pixeltable as pxt
 from pixeltable.utils.code import local_public_names
 
 
-@pxt.uda(
-    update_types=[pxt.JsonType(nullable=True)],
-    value_type=pxt.JsonType(),
-    requires_order_by=False,
-    allows_window=False,
-)
+@pxt.uda
 class make_list(pxt.Aggregator):
     """
     Collects arguments into a list.
     """
-    def __init__(self):
+    def __init__(self) -> None:
         self.output: list[Any] = []
 
-    def update(self, obj: Any) -> None:
+    def update(self, obj: pxt.Json) -> None:
         if obj is None:
             return
         self.output.append(obj)

pixeltable/functions/math.py

@@ -0,0 +1,67 @@
+import builtins
+import math
+from typing import Optional
+
+import sqlalchemy as sql
+
+import pixeltable as pxt
+from pixeltable.utils.code import local_public_names
+
+
+@pxt.udf(is_method=True)
+def abs(self: float) -> float:
+    return builtins.abs(self)
+
+
+@abs.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.func.abs(self)
+
+
+@pxt.udf(is_method=True)
+def ceil(self: float) -> float:
+    # This ensures the same behavior as SQL
+    if math.isfinite(self):
+        return float(math.ceil(self))
+    else:
+        return self
+
+
+@ceil.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.func.ceiling(self)
+
+
+@pxt.udf(is_method=True)
+def floor(self: float) -> float:
+    # This ensures the same behavior as SQL
+    if math.isfinite(self):
+        return float(math.floor(self))
+    else:
+        return self
+
+
+@floor.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.func.floor(self)
+
+
+@pxt.udf(is_method=True)
+def round(self: float, digits: Optional[int] = None) -> float:
+    # Set digits explicitly to 0 to guarantee a return type of float; this ensures the same behavior as SQL
+    return builtins.round(self, digits or 0)
+
+
+@round.to_sql
+def _(self: sql.ColumnElement, digits: Optional[sql.ColumnElement] = None) -> sql.ColumnElement:
+    if digits is None:
+        return sql.func.round(self)
+    else:
+        return sql.func.round(sql.cast(self, sql.Numeric), sql.cast(digits, sql.Integer))
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/functions/ollama.py

@@ -34,7 +34,7 @@ def generate(
     template: str = '',
     context: Optional[list[int]] = None,
     raw: bool = False,
-    format: str = '',
+    format: Optional[str] = None,
     options: Optional[dict] = None,
 ) -> dict:
     """
@@ -44,7 +44,7 @@ def generate(
         prompt: The prompt to generate a response for.
         model: The model name.
         suffix: The text after the model response.
-        format: The format of the response; must be one of `'json'` or `''` (the empty string).
+        format: The format of the response; must be one of `'json'` or `None`.
         system: System message.
         template: Prompt template to use.
         context: The context parameter returned from a previous call to `generate()`.
@@ -77,7 +77,7 @@ def chat(
     *,
     model: str,
     tools: Optional[list[dict]] = None,
-    format: str = '',
+    format: Optional[str] = None,
     options: Optional[dict] = None,
 ) -> dict:
     """
@@ -87,7 +87,7 @@ def chat(
         messages: The messages of the chat.
         model: The model name.
         tools: Tools for the model to use.
-        format: The format of the response; must be one of `'json'` or `''` (the empty string).
+        format: The format of the response; must be one of `'json'` or `None`.
         options: Additional options to pass to the `chat` call, such as `max_tokens`, `temperature`, `top_p`, and `top_k`.
             For details, see the
             [Valid Parameters and Values](https://github.com/ollama/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values)

pixeltable/functions/openai.py

@@ -7,6 +7,7 @@ the [Working with OpenAI](https://pixeltable.readme.io/docs/working-with-openai)
 
 import base64
 import io
+import json
 import pathlib
 import uuid
 from typing import TYPE_CHECKING, Callable, Optional, TypeVar, Union
@@ -16,8 +17,8 @@ import PIL.Image
 import tenacity
 
 import pixeltable as pxt
-from pixeltable import env
-from pixeltable.func import Batch
+from pixeltable import env, exprs
+from pixeltable.func import Batch, Tools
 from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
@@ -225,6 +226,16 @@ def chat_completions(
             ]
             tbl['response'] = chat_completions(messages, model='gpt-4o-mini')
     """
+
+    if tools is not None:
+        tools = [
+            {
+                'type': 'function',
+                'function': tool
+            }
+            for tool in tools
+        ]
+
     result = _retry(_openai_client().chat.completions.create)(
         messages=messages,
         model=model,
@@ -453,6 +464,24 @@ def moderations(input: str, *, model: Optional[str] = None) -> dict:
     return result.dict()
 
 
+def invoke_tools(tools: Tools, response: exprs.Expr) -> exprs.InlineDict:
+    """Converts an OpenAI response dict to Pixeltable tool invocation format and calls `tools._invoke()`."""
+    return tools._invoke(_openai_response_to_pxt_tool_calls(response))
+
+
+@pxt.udf
+def _openai_response_to_pxt_tool_calls(response: dict) -> Optional[dict]:
+    openai_tool_calls = response['choices'][0]['message']['tool_calls']
+    if openai_tool_calls is not None:
+        return {
+            tool_call['function']['name']: {
+                'args': json.loads(tool_call['function']['arguments'])
+            }
+            for tool_call in openai_tool_calls
+        }
+    return None
+
+
 _T = TypeVar('_T')
 
 

pixeltable/functions/timestamp.py

@@ -232,7 +232,7 @@ def _(
         sql.cast(day, sql.Integer),
         sql.cast(hour, sql.Integer),
         sql.cast(minute, sql.Integer),
-        sql.cast(second + microsecond / 1000000.0, sql.Double))
+        sql.cast(second + microsecond / 1000000.0, sql.Float))
 
 # @pxt.udf
 # def date(self: datetime) -> datetime:

pixeltable/functions/video.py

@@ -47,13 +47,7 @@ _format_defaults = {  # format -> (codec, ext)
 #         output_container.mux(packet)
 
 
-@pxt.uda(
-    init_types=[pxt.IntType()],
-    update_types=[pxt.ImageType()],
-    value_type=pxt.VideoType(),
-    requires_order_by=True,
-    allows_window=False,
-)
+@pxt.uda(requires_order_by=True)
 class make_video(pxt.Aggregator):
     """
     Aggregator that creates a video from a sequence of images.
@@ -80,7 +74,7 @@ class make_video(pxt.Aggregator):
         for packet in self.stream.encode(av_frame):
             self.container.mux(packet)
 
-    def value(self) -> str:
+    def value(self) -> pxt.Video:
         for packet in self.stream.encode():
             self.container.mux(packet)
         self.container.close()

pixeltable/functions/vision.py

@@ -220,7 +220,7 @@ def eval_detections(
     return result
 
 
-@pxt.uda(update_types=[pxt.JsonType()], value_type=pxt.JsonType(), allows_std_agg=True, allows_window=False)
+@pxt.uda
 class mean_ap(pxt.Aggregator):
     """
     Calculates the mean average precision (mAP) over