pixeltable 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

pixeltable/exprs/inline_dict.py

@@ -1,15 +1,15 @@
 from __future__ import annotations
-from typing import Optional, List, Any, Dict, Tuple
+
 import copy
+from typing import Optional, List, Any, Dict, Tuple
 
 import sqlalchemy as sql
 
-from .expr import Expr
-from .data_row import DataRow
-from .row_builder import RowBuilder
 import pixeltable.exceptions as excs
-import pixeltable.catalog as catalog
 import pixeltable.type_system as ts
+from .data_row import DataRow
+from .expr import Expr
+from .row_builder import RowBuilder
 
 
 class InlineDict(Expr):
@@ -21,8 +21,8 @@ class InlineDict(Expr):
         super().__init__(ts.JsonType())  # we need to call this in order to populate self.components
         # dict_items contains
         # - for Expr fields: (key, index into components, None)
-        # - for non-Expr fields: (key, -1, value)
-        self.dict_items: List[Tuple[str, int, Any]] = []
+        # - for non-Expr fields: (key, None, value)
+        self.dict_items: List[Tuple[str, Optional[int], Any]] = []
         for key, val in d.items():
             if not isinstance(key, str):
                 raise excs.Error(f'Dictionary requires string keys, {key} has type {type(key)}')
@@ -35,11 +35,11 @@ class InlineDict(Expr):
                 self.dict_items.append((key, len(self.components), None))
                 self.components.append(val)
             else:
-                self.dict_items.append((key, -1, val))
+                self.dict_items.append((key, None, val))
 
         self.type_spec: Optional[Dict[str, ts.ColumnType]] = {}
         for key, idx, _ in self.dict_items:
-            if idx == -1:
+            if idx is None:
                 # TODO: implement type inference for values
                 self.type_spec = None
                 break
@@ -56,7 +56,7 @@ class InlineDict(Expr):
                 return f"'{val}'"
             return str(val)
         for key, idx, val in self.dict_items:
-            if idx != -1:
+            if idx is not None:
                 item_strs.append(f"'{key}': {str(self.components[i])}")
                 i += 1
             else:
@@ -71,7 +71,7 @@ class InlineDict(Expr):
 
     def to_dict(self) -> Dict[str, Any]:
         """Return the original dict used to construct this"""
-        return {key: val if idx == -1 else self.components[idx] for key, idx, val in self.dict_items}
+        return {key: val if idx is None else self.components[idx] for key, idx, val in self.dict_items}
 
     def sql_expr(self) -> Optional[sql.ClauseElement]:
         return None
@@ -80,7 +80,7 @@ class InlineDict(Expr):
         result = {}
         for key, idx, val in self.dict_items:
             assert isinstance(key, str)
-            if idx >= 0:
+            if idx is not None:
                 result[key] = data_row[self.components[idx].slot_idx]
             else:
                 result[key] = copy.deepcopy(val)
@@ -94,7 +94,9 @@ class InlineDict(Expr):
         assert 'dict_items' in d
         arg: Dict[str, Any] = {}
         for key, idx, val in d['dict_items']:
-            if idx >= 0:
+            # TODO Normalize idx -1 to None via schema migrations.
+            # Long-term we should not be allowing idx == -1.
+            if idx is not None and idx >= 0:  # Older schemas might have -1 instead of None
                 arg[key] = components[idx]
             else:
                 arg[key] = val

pixeltable/exprs/row_builder.py

@@ -60,6 +60,8 @@ class RowBuilder:
         Args:
             output_exprs: list of Exprs to be evaluated
             columns: list of columns to be materialized
+            input_exprs: list of Exprs that are excluded from evaluation (because they're already materialized)
+        TODO: enforce that output_exprs doesn't overlap with input_exprs?
         """
         self.unique_exprs = ExprSet()  # dependencies precede their dependents
         self.next_slot_idx = 0
@@ -179,12 +181,16 @@ class RowBuilder:
             for i, c in enumerate(expr.components):
                 # make sure we only refer to components that have themselves been recorded
                 expr.components[i] = self._record_unique_expr(c, True)
-        assert expr.slot_idx < 0
+        assert expr.slot_idx is None
         expr.slot_idx = self._next_slot_idx()
         self.unique_exprs.append(expr)
         return expr
 
     def _record_output_expr_id(self, e: Expr, output_expr_id: int) -> None:
+        assert e.slot_idx is not None
+        assert output_expr_id is not None
+        if e.slot_idx in self.input_expr_slot_idxs:
+            return
         self.output_expr_ids[e.slot_idx].add(output_expr_id)
         for d in e.dependencies():
             self._record_output_expr_id(d, output_expr_id)

pixeltable/exprs/similarity_expr.py

@@ -0,0 +1,67 @@
+from typing import Optional, List
+
+import sqlalchemy as sql
+import PIL.Image
+
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
+from .column_ref import ColumnRef
+from .data_row import DataRow
+from .expr import Expr
+from .literal import Literal
+from .row_builder import RowBuilder
+
+
+class SimilarityExpr(Expr):
+
+    def __init__(self, col_ref: ColumnRef, item: Expr):
+        super().__init__(ts.FloatType())
+        self.components = [col_ref, item]
+        self.id = self._create_id()
+        assert isinstance(item, Literal)
+        assert item.col_type.is_string_type() or item.col_type.is_image_type()
+
+        # determine index to use
+        idx_info = col_ref.col.get_idx_info()
+        import pixeltable.index as index
+        embedding_idx_info = [info for info in idx_info.values() if isinstance(info.idx, index.EmbeddingIndex)]
+        if len(embedding_idx_info) == 0:
+            raise excs.Error(f'No index found for column {col_ref.col}')
+        if len(embedding_idx_info) > 1:
+            raise excs.Error(
+                f'Column {col_ref.col.name} has multiple indices; use the index name to disambiguate, '
+                f'e.g., `{col_ref.col.name}.<index-name>.similarity(...)`')
+        self.idx_info = embedding_idx_info[0]
+        idx = self.idx_info.idx
+
+        if item.col_type.is_string_type() and idx.txt_embed is None:
+            raise excs.Error(
+                f'Embedding index {self.idx_info.name} on column {self.idx_info.col.name} was created without the '
+                f'text_embed parameter and does not support text queries')
+        if item.col_type.is_image_type() and idx.img_embed is None:
+            raise excs.Error(
+                f'Embedding index {self.idx_info.name} on column {self.idx_info.col.name} was created without the '
+                f'img_embed parameter and does not support image queries')
+
+    def __str__(self) -> str:
+        return f'{self.components[0]}.similarity({self.components[1]})'
+
+    def sql_expr(self) -> Optional[sql.ClauseElement]:
+        assert isinstance(self.components[1], Literal)
+        item = self.components[1].val
+        return self.idx_info.idx.similarity_clause(self.idx_info.val_col, item)
+
+    def as_order_by_clause(self, is_asc: bool) -> Optional[sql.ClauseElement]:
+        assert isinstance(self.components[1], Literal)
+        item = self.components[1].val
+        return self.idx_info.idx.order_by_clause(self.idx_info.val_col, item, is_asc)
+
+    def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
+        # this should never get called
+        assert False
+
+    @classmethod
+    def _from_dict(cls, d: dict, components: List[Expr]) -> Expr:
+        assert len(components) == 2
+        assert isinstance(components[0], ColumnRef)
+        return cls(components[0], components[1])

pixeltable/ext/functions/whisperx.py

@@ -0,0 +1,30 @@
+from typing import Optional
+
+import torch
+import whisperx
+from whisperx.asr import FasterWhisperPipeline
+
+import pixeltable as pxt
+
+
+@pxt.udf(param_types=[pxt.AudioType(), pxt.StringType(), pxt.StringType(), pxt.StringType(), pxt.IntType()])
+def transcribe(
+    audio: str, *, model: str, compute_type: Optional[str] = None, language: Optional[str] = None, chunk_size: int = 30
+) -> dict:
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    compute_type = compute_type or ('float16' if device == 'cuda' else 'int8')
+    model = _lookup_model(model, device, compute_type)
+    audio_array = whisperx.load_audio(audio)
+    result = model.transcribe(audio_array, batch_size=16, language=language, chunk_size=chunk_size)
+    return result
+
+
+def _lookup_model(model_id: str, device: str, compute_type: str) -> FasterWhisperPipeline:
+    key = (model_id, device, compute_type)
+    if key not in _model_cache:
+        model = whisperx.load_model(model_id, device, compute_type=compute_type)
+        _model_cache[key] = model
+    return _model_cache[key]
+
+
+_model_cache = {}

pixeltable/ext/functions/yolox.py

@@ -56,6 +56,22 @@ def yolox(images: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0
     return results
 
 
+@pxt.udf
+def yolo_to_coco(detections: dict) -> list:
+    bboxes, labels = detections['bboxes'], detections['labels']
+    num_annotations = len(detections['bboxes'])
+    assert num_annotations == len(detections['labels'])
+    result = []
+    for i in range(num_annotations):
+        bbox = bboxes[i]
+        ann = {
+            'bbox': [round(bbox[0]), round(bbox[1]), round(bbox[2] - bbox[0]), round(bbox[3] - bbox[1])],
+            'category': labels[i],
+        }
+        result.append(ann)
+    return result
+
+
 def _images_to_tensors(images: Iterable[PIL.Image.Image], exp: Exp) -> Iterator[torch.Tensor]:
     for image in images:
         image_transform, _ = _val_transform(np.array(image), None, exp.test_size)

pixeltable/func/__init__.py

@@ -1,9 +1,7 @@
 from .aggregate_function import Aggregator, AggregateFunction, uda
-from .batched_function import BatchedFunction, ExplicitBatchedFunction
 from .callable_function import CallableFunction
 from .expr_template_function import ExprTemplateFunction
 from .function import Function
 from .function_registry import FunctionRegistry
-from .nos_function import NOSFunction
 from .signature import Signature, Parameter, Batch
 from .udf import udf, make_function, expr_udf

pixeltable/func/aggregate_function.py

@@ -72,6 +72,9 @@ class AggregateFunction(Function):
             if param.lower() in self.RESERVED_PARAMS:
                 raise excs.Error(f'{self.name}(): parameter name {param} is reserved')
 
+    def exec(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError
+
     def help_str(self) -> str:
         res = super().help_str()
         res += '\n\n' + inspect.getdoc(self.agg_cls.update)
@@ -137,7 +140,7 @@ def uda(
         update_types: List[ts.ColumnType],
         init_types: Optional[List[ts.ColumnType]] = None,
         requires_order_by: bool = False, allows_std_agg: bool = True, allows_window: bool = False,
-) -> Callable:
+) -> Callable[[Type[Aggregator]], AggregateFunction]:
     """Decorator for user-defined aggregate functions.
 
     The decorated class must inherit from Aggregator and implement the following methods:
@@ -159,7 +162,7 @@ def uda(
     if init_types is None:
         init_types = []
 
-    def decorator(cls: Type[Aggregator]) -> Type[Function]:
+    def decorator(cls: Type[Aggregator]) -> AggregateFunction:
         # validate type parameters
         num_init_params = len(inspect.signature(cls.__init__).parameters) - 1
         if num_init_params > 0:

pixeltable/func/callable_function.py

@@ -1,16 +1,12 @@
 from __future__ import annotations
 
 import inspect
-import sys
-from typing import Optional, Dict, Callable, List, Tuple
+from typing import Optional, Callable, Tuple, Any
 from uuid import UUID
+
 import cloudpickle
 
-import pixeltable.type_system as ts
-import pixeltable.exceptions as excs
 from .function import Function
-from .function_registry import FunctionRegistry
-from .globals import get_caller_module_path
 from .signature import Signature
 
 
@@ -24,13 +20,48 @@ class CallableFunction(Function):
 
     def __init__(
             self, signature: Signature, py_fn: Callable, self_path: Optional[str] = None,
-            self_name: Optional[str] = None):
+            self_name: Optional[str] = None, batch_size: Optional[int] = None):
         assert py_fn is not None
         self.py_fn = py_fn
         self.self_name = self_name
+        self.batch_size = batch_size
         py_signature = inspect.signature(self.py_fn)
         super().__init__(signature, py_signature, self_path=self_path)
 
+    @property
+    def is_batched(self) -> bool:
+        return self.batch_size is not None
+
+    def exec(self, *args: Any, **kwargs: Any) -> Any:
+        if self.is_batched:
+            # Pack the batched parameters into singleton lists
+            constant_param_names = [p.name for p in self.signature.constant_parameters]
+            batched_args = [[arg] for arg in args]
+            constant_kwargs = {k: v for k, v in kwargs.items() if k in constant_param_names}
+            batched_kwargs = {k: [v] for k, v in kwargs.items() if k not in constant_param_names}
+            result = self.py_fn(*batched_args, **constant_kwargs, **batched_kwargs)
+            assert len(result) == 1
+            return result[0]
+        else:
+            return self.py_fn(*args, **kwargs)
+
+    def exec_batch(self, *args: Any, **kwargs: Any) -> list:
+        """Execute the function with the given arguments and return the result.
+        The arguments are expected to be batched: if the corresponding parameter has type T,
+        then the argument should have type T if it's a constant parameter, or list[T] if it's
+        a batched parameter.
+        """
+        assert self.is_batched
+        # Unpack the constant parameters
+        constant_param_names = [p.name for p in self.signature.constant_parameters]
+        constant_kwargs = {k: v[0] for k, v in kwargs.items() if k in constant_param_names}
+        batched_kwargs = {k: v for k, v in kwargs.items() if k not in constant_param_names}
+        return self.py_fn(*args, **constant_kwargs, **batched_kwargs)
+
+    # TODO(aaron-siegel): Implement conditional batch sizing
+    def get_batch_size(self, *args: Any, **kwargs: Any) -> Optional[int]:
+        return self.batch_size
+
     @property
     def display_name(self) -> str:
         return self.self_name
@@ -44,7 +75,7 @@ class CallableFunction(Function):
         res += '\n\n' + inspect.getdoc(self.py_fn)
         return res
 
-    def _as_dict(self) -> Dict:
+    def _as_dict(self) -> dict:
         if self.self_path is None:
             # this is not a module function
             from .function_registry import FunctionRegistry
@@ -53,17 +84,30 @@ class CallableFunction(Function):
         return super()._as_dict()
 
     @classmethod
-    def _from_dict(cls, d: Dict) -> Function:
+    def _from_dict(cls, d: dict) -> Function:
         if 'id' in d:
             from .function_registry import FunctionRegistry
             return FunctionRegistry.get().get_stored_function(UUID(hex=d['id']))
         return super()._from_dict(d)
 
-    def to_store(self) -> Tuple[Dict, bytes]:
-        return (self.signature.as_dict(), cloudpickle.dumps(self.py_fn))
+    def to_store(self) -> tuple[dict, bytes]:
+        md = self.signature.as_dict()
+        if self.batch_size is not None:
+            md['batch_size'] = self.batch_size
+        return md, cloudpickle.dumps(self.py_fn)
 
     @classmethod
-    def from_store(cls, name: Optional[str], md: Dict, binary_obj: bytes) -> Function:
+    def from_store(cls, name: Optional[str], md: dict, binary_obj: bytes) -> Function:
         py_fn = cloudpickle.loads(binary_obj)
         assert isinstance(py_fn, Callable)
-        return CallableFunction(Signature.from_dict(md), py_fn, self_name=name)
+        return CallableFunction(Signature.from_dict(md), py_fn, self_name=name, batch_size=md.get('batch_size'))
+
+    def validate_call(self, bound_args: dict[str, Any]) -> None:
+        import pixeltable.exprs as exprs
+        if self.is_batched:
+            for param in self.signature.constant_parameters:
+                if param.name in bound_args and isinstance(bound_args[param.name], exprs.Expr):
+                    raise ValueError(
+                        f'{self.display_name}(): '
+                        f'parameter {param.name} must be a constant value, not a Pixeltable expression'
+                    )

pixeltable/func/expr_template_function.py

@@ -1,9 +1,8 @@
 import inspect
-from typing import Dict, Optional, Callable, List
+from typing import Dict, Optional, Any
 
 import pixeltable
 import pixeltable.exceptions as excs
-import pixeltable.type_system as ts
 from .function import Function
 from .signature import Signature, Parameter
 
@@ -51,6 +50,7 @@ class ExprTemplateFunction(Function):
             {param_name: default for param_name, default in self.defaults.items() if param_name not in bound_args})
         result = self.expr.copy()
         import pixeltable.exprs as exprs
+        arg_exprs: dict[exprs.Expr, exprs.Expr] = {}
         for param_name, arg in bound_args.items():
             param_expr = self.param_exprs_by_name[param_name]
             if not isinstance(arg, exprs.Expr):
@@ -60,11 +60,22 @@ class ExprTemplateFunction(Function):
                     raise excs.Error(f'{self.self_name}(): cannot convert argument {arg} to a Pixeltable expression')
             else:
                 arg_expr = arg
-            result = result.substitute(param_expr, arg_expr)
+            arg_exprs[param_expr] = arg_expr
+        result = result.substitute(arg_exprs)
         import pixeltable.exprs as exprs
         assert not result.contains(exprs.Variable)
         return result
 
+    def exec(self, *args: Any, **kwargs: Any) -> Any:
+        expr = self.instantiate(*args, **kwargs)
+        import pixeltable.exprs as exprs
+        row_builder = exprs.RowBuilder(output_exprs=[expr], columns=[], input_exprs=[])
+        import pixeltable.exec as exec
+        row_batch = exec.DataRowBatch(tbl=None, row_builder=row_builder, len=1)
+        row = row_batch[0]
+        row_builder.eval(row, ctx=row_builder.default_eval_ctx)
+        return row[row_builder.get_output_exprs()[0].slot_idx]
+
     @property
     def display_name(self) -> str:
         return self.self_name

pixeltable/func/function.py

@@ -3,9 +3,10 @@ from __future__ import annotations
 import abc
 import importlib
 import inspect
-import pixeltable
-from typing import Optional, Dict, Any, Tuple
+from typing import Optional, Dict, Any, Tuple, Callable
 
+import pixeltable
+import pixeltable.type_system as ts
 from .globals import resolve_symbol
 from .signature import Signature
 
@@ -18,10 +19,13 @@ class Function(abc.ABC):
     via the member self_path.
     """
 
-    def __init__(self, signature: Signature, py_signature: inspect.Signature, self_path: Optional[str] = None):
+    def __init__(
+            self, signature: Signature, py_signature: inspect.Signature, self_path: Optional[str] = None
+    ):
         self.signature = signature
         self.py_signature = py_signature
         self.self_path = self_path  # fully-qualified path to self
+        self._conditional_return_type: Optional[Callable[..., ts.ColumnType]] = None
 
     @property
     def name(self) -> str:
@@ -40,7 +44,7 @@ class Function(abc.ABC):
     def help_str(self) -> str:
         return self.display_name + str(self.signature)
 
-    def __call__(self, *args: object, **kwargs: object) -> 'pixeltable.exprs.Expr':
+    def __call__(self, *args: Any, **kwargs: Any) -> 'pixeltable.exprs.Expr':
         from pixeltable import exprs
         bound_args = self.py_signature.bind(*args, **kwargs)
         self.validate_call(bound_args.arguments)
@@ -50,6 +54,33 @@ class Function(abc.ABC):
         """Override this to do custom validation of the arguments"""
         pass
 
+    def call_return_type(self, kwargs: dict[str, Any]) -> ts.ColumnType:
+        """Return the type of the value returned by calling this function with the given arguments"""
+        if self._conditional_return_type is None:
+            return self.signature.return_type
+        bound_args = self.py_signature.bind(**kwargs)
+        kw_args: dict[str, Any] = {}
+        sig = inspect.signature(self._conditional_return_type)
+        for param in sig.parameters.values():
+            if param.name in bound_args.arguments:
+                kw_args[param.name] = bound_args.arguments[param.name]
+        return self._conditional_return_type(**kw_args)
+
+    def conditional_return_type(self, fn: Callable[..., ts.ColumnType]) -> Callable[..., ts.ColumnType]:
+        """Instance decorator for specifying a conditional return type for this function"""
+        # verify that call_return_type only has parameters that are also present in the signature
+        sig = inspect.signature(fn)
+        for param in sig.parameters.values():
+            if param.name not in self.signature.parameters:
+                raise ValueError(f'`conditional_return_type` has parameter `{param.name}` that is not in the signature')
+        self._conditional_return_type = fn
+        return fn
+
+    @abc.abstractmethod
+    def exec(self, *args: Any, **kwargs: Any) -> Any:
+        """Execute the function with the given arguments and return the result."""
+        pass
+
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, self.__class__):
             return False

pixeltable/func/signature.py

@@ -29,21 +29,12 @@ class Signature:
     """
     Represents the signature of a Pixeltable function.
 
-    Regarding return type:
-    - most functions will have a fixed return type, which is specified directly
-    - some functions will have a return type that depends on the argument values;
-      ex.: PIL.Image.Image.resize() returns an image with dimensions specified as a parameter
-    - in the latter case, the 'return_type' field is a function that takes the bound arguments and returns the
-      return type; if no bound arguments are specified, a generic return type is returned (eg, ImageType() without a
-      size)
     - self.is_batched: return type is a Batch[...] type
     """
     SPECIAL_PARAM_NAMES = ['group_by', 'order_by']
 
-    def __init__(
-            self,
-            return_type: Union[ts.ColumnType, Callable[[Dict[str, Any]], ts.ColumnType]],
-            parameters: List[Parameter], is_batched: bool = False):
+    def __init__(self, return_type: ts.ColumnType, parameters: List[Parameter], is_batched: bool = False):
+        assert isinstance(return_type, ts.ColumnType)
         self.return_type = return_type
         self.is_batched = is_batched
         # we rely on the ordering guarantee of dicts in Python >=3.7
@@ -52,10 +43,9 @@ 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]
 
-    def get_return_type(self, bound_args: Optional[Dict[str, Any]] = None) -> ts.ColumnType:
-        if isinstance(self.return_type, ts.ColumnType):
-            return self.return_type
-        return self.return_type(bound_args)
+    def get_return_type(self) -> ts.ColumnType:
+        assert isinstance(self.return_type, ts.ColumnType)
+        return self.return_type
 
     def as_dict(self) -> Dict[str, Any]:
         result = {

pixeltable/func/udf.py

@@ -6,7 +6,6 @@ from typing import List, Callable, Optional, overload, Any
 import pixeltable as pxt
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
-from .batched_function import ExplicitBatchedFunction
 from .callable_function import CallableFunction
 from .expr_template_function import ExprTemplateFunction
 from .function import Function
@@ -29,7 +28,7 @@ def udf(
         batch_size: Optional[int] = None,
         substitute_fn: Optional[Callable] = None,
         _force_stored: bool = False
-) -> Callable: ...
+) -> Callable[[Callable], Function]: ...
 
 
 def udf(*args, **kwargs):
@@ -62,8 +61,8 @@ def udf(*args, **kwargs):
 
         def decorator(decorated_fn: Callable):
             return make_function(
-                decorated_fn, return_type, param_types, batch_size, substitute_fn=substitute_fn,
-                force_stored=force_stored)
+                decorated_fn, return_type, param_types, batch_size,
+                substitute_fn=substitute_fn, force_stored=force_stored)
 
         return decorator
 
@@ -78,8 +77,8 @@ def make_function(
     force_stored: bool = False
 ) -> Function:
     """
-    Constructs a `CallableFunction` or `BatchedFunction`, depending on the
-    supplied parameters. If `substitute_fn` is specified, then `decorated_fn`
+    Constructs a `CallableFunction` from the specified parameters.
+    If `substitute_fn` is specified, then `decorated_fn`
     will be used only for its signature, with execution delegated to
     `substitute_fn`.
     """
@@ -117,11 +116,8 @@ def make_function(
             raise excs.Error(f'{errmsg_name}(): @udf decorator with a `substitute_fn` can only be used in a module')
         py_fn = substitute_fn
 
-    if batch_size is None:
-        result = CallableFunction(signature=sig, py_fn=py_fn, self_path=function_path, self_name=function_name)
-    else:
-        result = ExplicitBatchedFunction(
-            signature=sig, batch_size=batch_size, invoker_fn=py_fn, self_path=function_path)
+    result = CallableFunction(
+        signature=sig, py_fn=py_fn, self_path=function_path, self_name=function_name, batch_size=batch_size)
 
     # If this function is part of a module, register it
     if function_path is not None:
@@ -135,7 +131,7 @@ def make_function(
 def expr_udf(py_fn: Callable) -> ExprTemplateFunction: ...
 
 @overload
-def expr_udf(*, param_types: Optional[List[ts.ColumnType]] = None) -> Callable: ...
+def expr_udf(*, param_types: Optional[List[ts.ColumnType]] = None) -> Callable[[Callable], ExprTemplateFunction]: ...
 
 def expr_udf(*args: Any, **kwargs: Any) -> Any:
     def decorator(py_fn: Callable, param_types: Optional[List[ts.ColumnType]]) -> ExprTemplateFunction:

pixeltable/functions/fireworks.py

@@ -6,8 +6,13 @@ import pixeltable as pxt
 from pixeltable import env
 
 
-def fireworks_client() -> fireworks.client.Fireworks:
-    return env.Env.get().get_client('fireworks', lambda api_key: fireworks.client.Fireworks(api_key=api_key))
+@env.register_client('fireworks')
+def _(api_key: str) -> fireworks.client.Fireworks:
+    return fireworks.client.Fireworks(api_key=api_key)
+
+
+def _fireworks_client() -> fireworks.client.Fireworks:
+    return env.Env.get().get_client('fireworks')
 
 
 @pxt.udf
@@ -26,8 +31,8 @@ def chat_completions(
         'top_p': top_p,
         'temperature': temperature
     }
-    kwargs_not_none = dict(filter(lambda x: x[1] is not None, kwargs.items()))
-    return fireworks_client().chat.completions.create(
+    kwargs_not_none = {k: v for k, v in kwargs.items() if v is not None}
+    return _fireworks_client().chat.completions.create(
         model=model,
         messages=messages,
         **kwargs_not_none