pixeltable 0.2.24__py3-none-any.whl → 0.3.0__py3-none-any.whl

pixeltable/func/tools.py

@@ -0,0 +1,153 @@
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, TypeVar, Union
+
+import pydantic
+
+import pixeltable.exceptions as excs
+
+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: Optional[str] = None
+    description: Optional[str] = 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._docstring(),
+            '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?
+        }
+
+    # `tool_calls` must be in standardized tool invocation format:
+    # {tool_name: {'args': {name1: value1, name2: value2, ...}}, ...}
+    def invoke(self, tool_calls: 'exprs.Expr') -> 'exprs.FunctionCall':
+        kwargs = {
+            param.name: self.__extract_tool_arg(param, tool_calls)
+            for param in self.parameters.values()
+        }
+        return self.fn(**kwargs)
+
+    def __extract_tool_arg(self, param: Parameter, tool_calls: 'exprs.Expr') -> 'exprs.Expr':
+        func_name = self.name or self.fn.name
+        if param.col_type.is_string_type():
+            return _extract_str_tool_arg(tool_calls, func_name=func_name, param_name=param.name)
+        if param.col_type.is_int_type():
+            return _extract_int_tool_arg(tool_calls, func_name=func_name, param_name=param.name)
+        if param.col_type.is_float_type():
+            return _extract_float_tool_arg(tool_calls, func_name=func_name, param_name=param.name)
+        if param.col_type.is_bool_type():
+            return _extract_bool_tool_arg(tool_calls, func_name=func_name, param_name=param.name)
+        assert False
+
+
+class ToolChoice(pydantic.BaseModel):
+    auto: bool
+    required: bool
+    tool: Optional[str]
+    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: Union[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: Optional[str] = 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}')
+        return ToolChoice(auto=auto, required=required, tool=tool_name, parallel_tool_calls=parallel_tool_calls)
+
+
+@udf
+def _extract_str_tool_arg(tool_calls: dict[str, Any], func_name: str, param_name: str) -> Optional[str]:
+    return _extract_arg(str, tool_calls, func_name, param_name)
+
+
+@udf
+def _extract_int_tool_arg(tool_calls: dict[str, Any], func_name: str, param_name: str) -> Optional[int]:
+    return _extract_arg(int, tool_calls, func_name, param_name)
+
+
+@udf
+def _extract_float_tool_arg(tool_calls: dict[str, Any], func_name: str, param_name: str) -> Optional[float]:
+    return _extract_arg(float, tool_calls, func_name, param_name)
+
+
+@udf
+def _extract_bool_tool_arg(tool_calls: dict[str, Any], func_name: str, param_name: str) -> Optional[bool]:
+    return _extract_arg(bool, tool_calls, func_name, param_name)
+
+
+T = TypeVar('T')
+
+
+def _extract_arg(eval_fn: Callable[[Any], T], tool_calls: dict[str, Any], func_name: str, param_name: str) -> Optional[T]:
+    if func_name in tool_calls:
+        arguments = tool_calls[func_name]['args']
+        if param_name in arguments:
+            return eval_fn(arguments[param_name])
+        return None
+    return None

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, 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,9 @@ from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Union
 import tenacity
 
 import pixeltable as pxt
-from pixeltable import env
+import pixeltable.exceptions as excs
+from pixeltable import env, exprs
+from pixeltable.func import Tools
 from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
@@ -46,8 +48,8 @@ def messages(
     stop_sequences: Optional[list[str]] = None,
     system: Optional[str] = None,
     temperature: Optional[float] = None,
-    tool_choice: Optional[list[dict]] = None,
-    tools: Optional[dict] = None,
+    tool_choice: Optional[dict] = None,
+    tools: Optional[list[dict]] = None,
     top_k: Optional[int] = None,
     top_p: Optional[float] = None,
 ) -> dict:
@@ -77,6 +79,33 @@ 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
+        ]
+
+    tool_choice_: Optional[dict] = None
+    if tool_choice is not None:
+        if tool_choice['auto']:
+            tool_choice_ = {'type': 'auto'}
+        elif tool_choice['required']:
+            tool_choice_ = {'type': 'any'}
+        else:
+            assert tool_choice['tool'] is not None
+            tool_choice_ = {'type': 'tool', 'name': tool_choice['tool']}
+        if not tool_choice['parallel_tool_calls']:
+            tool_choice_['disable_parallel_tool_use'] = True
+
     return _retry(_anthropic_client().messages.create)(
         messages=messages,
         model=model,
@@ -85,13 +114,31 @@ def messages(
         stop_sequences=_opt(stop_sequences),
         system=_opt(system),
         temperature=_opt(temperature),
-        tool_choice=_opt(tool_choice),
+        tool_choice=_opt(tool_choice_),
         tools=_opt(tools),
         top_k=_opt(top_k),
         top_p=_opt(top_p),
     ).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/gemini.py

@@ -0,0 +1,85 @@
+"""
+Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
+that wrap various endpoints from the Google Gemini API. In order to use them, you must
+first `pip install google-generativeai` and configure your Gemini credentials, as described in
+the [Working with Gemini](https://pixeltable.readme.io/docs/working-with-gemini) tutorial.
+"""
+
+from typing import Optional
+
+import pixeltable as pxt
+from pixeltable import env
+
+
+@env.register_client('gemini')
+def _(api_key: str) -> None:
+    import google.generativeai as genai  # type: ignore[import-untyped]
+    genai.configure(api_key=api_key)
+
+
+def _ensure_loaded() -> None:
+    env.Env.get().get_client('gemini')
+
+
+@pxt.udf
+def generate_content(
+    contents: str,
+    *,
+    model_name: str,
+    candidate_count: Optional[int] = None,
+    stop_sequences: Optional[list[str]] = None,
+    max_output_tokens: Optional[int] = None,
+    temperature: Optional[float] = None,
+    top_p: Optional[float] = None,
+    top_k: Optional[int] = None,
+    response_mime_type: Optional[str] = None,
+    response_schema: Optional[dict] = None,
+    presence_penalty: Optional[float] = None,
+    frequency_penalty: Optional[float] = None,
+    response_logprobs: Optional[bool] = None,
+    logprobs: Optional[int] = None,
+) -> dict:
+    """
+    Generate content from the specified model. For additional details, see:
+    <https://ai.google.dev/gemini-api/docs>
+
+    __Requirements:__
+
+    - `pip install google-generativeai`
+
+    Args:
+        contents: The input content to generate from.
+        model_name: The name of the model to use.
+
+    For details on the other parameters, see: <https://ai.google.dev/gemini-api/docs>
+
+    Returns:
+        A dictionary containing the response and other metadata.
+
+    Examples:
+        Add a computed column that applies the model `gemini-1.5-flash`
+        to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
+
+        >>> tbl['response'] = generate_content(tbl.prompt, model_name='gemini-1.5-flash')
+    """
+    env.Env.get().require_package('google.generativeai')
+    _ensure_loaded()
+    import google.generativeai as genai
+
+    model = genai.GenerativeModel(model_name=model_name)
+    gc = genai.GenerationConfig(
+        candidate_count=candidate_count,
+        stop_sequences=stop_sequences,
+        max_output_tokens=max_output_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        response_mime_type=response_mime_type,
+        response_schema=response_schema,
+        presence_penalty=presence_penalty,
+        frequency_penalty=frequency_penalty,
+        response_logprobs=response_logprobs,
+        logprobs=logprobs,
+    )
+    response = model.generate_content(contents, generation_config=gc)
+    return response.to_dict()

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