pixeltable 0.3.10__py3-none-any.whl → 0.3.12__py3-none-any.whl

pixeltable/func/query_template_function.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, Optional, overload
+from functools import reduce
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Optional, Union, overload
 
-from pixeltable import exprs, type_system as ts
+from pixeltable import catalog, exceptions as excs, exprs, func, type_system as ts
 
 from .function import Function
 from .signature import Signature
@@ -17,6 +18,7 @@ class QueryTemplateFunction(Function):
 
     template_df: Optional['DataFrame']
     self_name: Optional[str]
+    _comment: Optional[str]
 
     @classmethod
     def create(
@@ -34,15 +36,21 @@ class QueryTemplateFunction(Function):
         assert isinstance(template_df, DataFrame)
         # we take params and return json
         sig = Signature(return_type=ts.JsonType(), parameters=params)
-        return QueryTemplateFunction(template_df, sig, path=path, name=name)
+        return QueryTemplateFunction(template_df, sig, path=path, name=name, comment=inspect.getdoc(template_callable))
 
     def __init__(
-        self, template_df: Optional['DataFrame'], sig: Signature, path: Optional[str] = None, name: Optional[str] = None
+        self,
+        template_df: Optional['DataFrame'],
+        sig: Signature,
+        path: Optional[str] = None,
+        name: Optional[str] = None,
+        comment: Optional[str] = None,
     ):
         assert sig is not None
         super().__init__([sig], self_path=path)
         self.self_name = name
         self.template_df = template_df
+        self._comment = comment
 
     def _update_as_overload_resolution(self, signature_idx: int) -> None:
         pass  # only one signature supported for QueryTemplateFunction
@@ -74,6 +82,9 @@ class QueryTemplateFunction(Function):
     def name(self) -> str:
         return self.self_name
 
+    def comment(self) -> Optional[str]:
+        return self._comment
+
     def _as_dict(self) -> dict:
         return {'name': self.name, 'signature': self.signature.as_dict(), 'df': self.template_df.as_dict()}
 
@@ -112,3 +123,75 @@ def query(*args: Any, **kwargs: Any) -> Any:
     else:
         assert len(args) == 0 and len(kwargs) == 1 and 'param_types' in kwargs
         return lambda py_fn: make_query_template(py_fn, kwargs['param_types'])
+
+
+def retrieval_udf(
+    table: catalog.Table,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    parameters: Optional[Iterable[Union[str, exprs.ColumnRef]]] = None,
+    limit: Optional[int] = 10,
+) -> func.QueryTemplateFunction:
+    """
+    Constructs a retrieval UDF for the given table. The retrieval UDF is a UDF whose parameters are
+    columns of the table and whose return value is a list of rows from the table. The return value of
+    ```python
+    f(col1=x, col2=y, ...)
+    ```
+    will be a list of all rows from the table that match the specified arguments.
+
+    Args:
+        table: The table to use as the dataset for the retrieval tool.
+        name: The name of the tool. If not specified, then the name of the table will be used by default.
+        description: The description of the tool. If not specified, then a default description will be generated.
+        parameters: The columns of the table to use as parameters. If not specified, all data columns
+            (non-computed columns) will be used as parameters.
+
+            All of the specified parameters will be required parameters of the tool, regardless of their status
+            as columns.
+        limit: The maximum number of rows to return. If not specified, then all matching rows will be returned.
+
+    Returns:
+        A list of dictionaries containing data from the table, one per row that matches the input arguments.
+        If there are no matching rows, an empty list will be returned.
+    """
+    # Argument validation
+    col_refs: list[exprs.ColumnRef]
+    if parameters is None:
+        col_refs = [table[col_name] for col_name in table.columns if not table[col_name].col.is_computed]
+    else:
+        for param in parameters:
+            if isinstance(param, str) and param not in table.columns:
+                raise excs.Error(f'The specified parameter {param!r} is not a column of the table {table._path!r}')
+        col_refs = [table[param] if isinstance(param, str) else param for param in parameters]
+
+    if len(col_refs) == 0:
+        raise excs.Error('Parameter list cannot be empty.')
+
+    # Construct the dataframe
+    predicates = [col_ref == exprs.Variable(col_ref.col.name, col_ref.col.col_type) for col_ref in col_refs]
+    where_clause = reduce(lambda c1, c2: c1 & c2, predicates)
+    df = table.select().where(where_clause)
+    if limit is not None:
+        df = df.limit(limit)
+
+    # Construct the signature
+    query_params = [
+        func.Parameter(col_ref.col.name, col_ref.col.col_type, inspect.Parameter.POSITIONAL_OR_KEYWORD)
+        for col_ref in col_refs
+    ]
+    query_signature = func.Signature(return_type=ts.JsonType(), parameters=query_params)
+
+    # Construct a name and/or description if not provided
+    if name is None:
+        name = table._name
+    if description is None:
+        description = (
+            f'Retrieves an entry from the dataset {name!r} that matches the given parameters.\n\nParameters:\n'
+        )
+        description += '\n'.join(
+            [f'    {col_ref.col.name}: of type `{col_ref.col.col_type._to_base_str()}`' for col_ref in col_refs]
+        )
+
+    fn = func.QueryTemplateFunction(df, query_signature, name=name, comment=description)
+    return fn

pixeltable/func/signature.py

@@ -253,7 +253,7 @@ class Signature:
                 continue  # skip 'self' or 'cls' parameter
             if param.name in cls.SPECIAL_PARAM_NAMES:
                 raise excs.Error(f'{param.name!r} is a reserved parameter name')
-            if param.kind in {inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD}:
+            if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
                 parameters.append(Parameter(param.name, col_type=None, kind=param.kind))
                 continue
 

pixeltable/func/tools.py

@@ -39,7 +39,7 @@ class Tool(pydantic.BaseModel):
     def ser_model(self) -> dict[str, Any]:
         return {
             'name': self.name or self.fn.name,
-            'description': self.description or self.fn._docstring(),
+            '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()},

pixeltable/func/udf.py

@@ -43,7 +43,7 @@ def udf(
 ) -> ExprTemplateFunction: ...
 
 
-def udf(*args, **kwargs):
+def udf(*args, **kwargs):  # type: ignore[no-untyped-def]
     """A decorator to create a Function from a function definition.
 
     Examples:
@@ -79,7 +79,7 @@ def udf(*args, **kwargs):
         if len(args) > 0:
             raise excs.Error('Unexpected @udf decorator arguments.')
 
-        def decorator(decorated_fn: Callable):
+        def decorator(decorated_fn: Callable) -> CallableFunction:
             return make_function(
                 decorated_fn,
                 batch_size=batch_size,

pixeltable/functions/__init__.py

@@ -29,5 +29,5 @@ from .globals import count, map, max, mean, min, sum
 __all__ = local_public_names(__name__, exclude=['globals']) + local_public_names(globals.__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/anthropic.py

@@ -39,7 +39,7 @@ def _anthropic_client() -> 'anthropic.AsyncAnthropic':
 
 
 class AnthropicRateLimitsInfo(env.RateLimitsInfo):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(self._get_request_resources)
 
     def _get_request_resources(self, messages: dict, max_tokens: int) -> dict[str, int]:
@@ -236,5 +236,5 @@ def _opt(arg: _T) -> Union[_T, 'anthropic.NotGiven']:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/audio.py

@@ -26,5 +26,5 @@ def get_metadata(audio: pxt.Audio) -> dict:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/deepseek.py

@@ -117,5 +117,5 @@ async def chat_completions(
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/fireworks.py

@@ -131,5 +131,5 @@ async def chat_completions(
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/globals.py

@@ -23,7 +23,7 @@ T = typing.TypeVar('T')
 class sum(func.Aggregator, typing.Generic[T]):
     """Sums the selected integers or floats."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.sum: T = None
 
     def update(self, val: T) -> None:
@@ -67,7 +67,7 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     ),
 )
 class count(func.Aggregator, typing.Generic[T]):
-    def __init__(self):
+    def __init__(self) -> None:
         self.count = 0
 
     def update(self, val: T) -> None:
@@ -88,7 +88,7 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     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):
+    def __init__(self) -> None:
         self.val: T = None
 
     def update(self, val: T) -> None:
@@ -118,7 +118,7 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     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):
+    def __init__(self) -> None:
         self.val: T = None
 
     def update(self, val: T) -> None:
@@ -143,7 +143,7 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
 
 @func.uda(type_substitutions=({T: Optional[int]}, {T: Optional[float]}))  # type: ignore[misc]
 class mean(func.Aggregator, typing.Generic[T]):
-    def __init__(self):
+    def __init__(self) -> None:
         self.sum: T = None
         self.count = 0
 
@@ -182,5 +182,5 @@ def map(expr: exprs.Expr, fn: Callable[[exprs.Expr], Any]) -> exprs.Expr:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/huggingface.py

@@ -493,5 +493,5 @@ _processor_cache: dict[tuple[str, Callable], Any] = {}
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/image.py

@@ -426,5 +426,5 @@ def mode(self: PIL.Image.Image) -> str:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/json.py

@@ -37,5 +37,5 @@ class make_list(pxt.Aggregator):
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/llama_cpp.py

@@ -100,5 +100,5 @@ _IS_GPU_AVAILABLE: Optional[bool] = None
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/math.py

@@ -63,5 +63,5 @@ def _(self: sql.ColumnElement, digits: Optional[sql.ColumnElement] = None) -> sq
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/mistralai.py

@@ -193,5 +193,5 @@ def _opt(arg: Optional[_T]) -> Union[_T, 'mistralai.types.basemodel.Unset']:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/ollama.py

@@ -129,5 +129,5 @@ def embed(
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/openai.py

@@ -122,7 +122,7 @@ class OpenAIRateLimitsInfo(env.RateLimitsInfo):
 _header_duration_pattern = re.compile(r'(?:(\d+)d)?(?:(\d+)h)?(?:(\d+)ms)|(?:(\d+)m)?(?:([\d.]+)s)?')
 
 
-def _parse_header_duration(duration_str):
+def _parse_header_duration(duration_str: str) -> datetime.timedelta:
     match = _header_duration_pattern.match(duration_str)
     if not match:
         raise ValueError('Invalid duration format')
@@ -837,5 +837,5 @@ def _opt(arg: _T) -> Union[_T, 'openai.NotGiven']:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/replicate.py

@@ -69,5 +69,5 @@ async def run(input: dict[str, Any], *, ref: str) -> dict[str, Any]:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/string.py

@@ -675,5 +675,5 @@ def zfill(self: str, width: int) -> str:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/timestamp.py

@@ -312,5 +312,5 @@ def posix_timestamp(self: datetime) -> float:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/together.py

@@ -298,5 +298,5 @@ async def image_generations(
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/util.py

@@ -21,7 +21,7 @@ def normalize_image_mode(image: PIL.Image.Image) -> PIL.Image.Image:
     Converts grayscale images to 3-channel for compatibility with models that only work with
     multichannel input.
     """
-    if image.mode in {'1', 'L'}:
+    if image.mode in ('1', 'L'):
         return image.convert('RGB')
     if image.mode == 'LA':
         return image.convert('RGBA')

pixeltable/functions/video.py

@@ -143,7 +143,7 @@ def _get_metadata(path: str) -> dict:
 
 
 def __get_stream_metadata(stream: av.stream.Stream) -> dict:
-    if stream.type not in {'audio', 'video'}:
+    if stream.type not in ('audio', 'video'):
         return {'type': stream.type}  # Currently unsupported
 
     codec_context = stream.codec_context
@@ -190,5 +190,5 @@ def __get_stream_metadata(stream: av.stream.Stream) -> dict:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/vision.py

@@ -238,7 +238,7 @@ class mean_ap(pxt.Aggregator):
     - A `dict[int, float]` mapping each label class to an average precision (AP) value for that class.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.class_tpfp: dict[int, list[dict]] = defaultdict(list)
 
     def update(self, eval_dicts: list[dict]) -> None:
@@ -394,5 +394,5 @@ def draw_bounding_boxes(
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/globals.py

@@ -616,9 +616,14 @@ def _extract_paths(
         matches = [name for name, entry in dir_entries.items() if entry.dir is not None]
     else:
         matches = [name for name, entry in dir_entries.items() if entry.table is not None]
+
+    # Filter out system paths
+    matches = [name for name in matches if catalog.is_valid_identifier(name)]
     result = [parent.append(name) for name in matches]
-    for name, entry in [(name, entry) for name, entry in dir_entries.items() if len(entry.dir_entries) > 0]:
-        result.extend(_extract_paths(entry.dir_entries, parent=parent.append(name), entry_type=entry_type))
+
+    for name, entry in dir_entries.items():
+        if len(entry.dir_entries) > 0 and catalog.is_valid_identifier(name):
+            result.extend(_extract_paths(entry.dir_entries, parent=parent.append(name), entry_type=entry_type))
     return result
 
 

pixeltable/index/embedding_index.py

@@ -208,7 +208,18 @@ class EmbeddingIndex(IndexBase):
                 and len(sig.required_parameters) <= 1
                 and sig.parameters_by_pos[0].col_type.type_enum == expected_type
             ):
-                return resolved_fn
+                # We found a valid signature. Now, if it has more than one parameter, we need to transform it into a
+                # 1-ary function by fixing all the other parameters to their defaults. This is to ensure that
+                # conditional_return_type resolves correctly.
+                if len(sig.parameters) == 1:
+                    unary_fn = resolved_fn
+                else:
+                    assert all(sig.parameters_by_pos[i].has_default for i in range(1, len(sig.parameters)))
+                    defaults = {param.name: param.default for param in sig.parameters_by_pos[1:]}
+                    unary_fn = resolved_fn.using(**defaults)
+                assert not unary_fn.is_polymorphic
+                assert len(unary_fn.signature.parameters) == 1
+                return unary_fn
         return None
 
     @classmethod

pixeltable/io/__init__.py

@@ -1,3 +1,5 @@
+# ruff: noqa: F401
+
 from .datarows import import_json, import_rows
 from .external_store import ExternalStore, SyncStatus
 from .globals import create_label_studio_project, export_images_as_fo_dataset
@@ -5,10 +7,10 @@ from .hf_datasets import import_huggingface_dataset
 from .pandas import import_csv, import_excel, import_pandas
 from .parquet import export_parquet, import_parquet
 
-__default_dir = set(symbol for symbol in dir() if not symbol.startswith('_'))
+__default_dir = {symbol for symbol in dir() if not symbol.startswith('_')}
 __removed_symbols = {'globals', 'hf_datasets', 'pandas', 'parquet', 'datarows'}
-__all__ = sorted(list(__default_dir - __removed_symbols))
+__all__ = sorted(__default_dir - __removed_symbols)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/io/fiftyone.py

@@ -1,5 +1,5 @@
 import os
-from typing import Iterator, Optional, Union
+from typing import Any, Iterator, Optional, Union
 
 import fiftyone as fo  # type: ignore[import-untyped]
 import fiftyone.utils.data as foud  # type: ignore[import-untyped]
@@ -59,10 +59,9 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
         for exprs_, label_cls, default_name in label_categories:
             if exprs_ is None or isinstance(exprs_, dict):
                 continue
-            if isinstance(exprs_, exprs.Expr):
-                exprs_ = [exprs_]
-            assert isinstance(exprs_, list)
-            for expr in exprs_:
+            exprs_list = [exprs_] if isinstance(exprs_, exprs.Expr) else exprs_
+            assert isinstance(exprs_list, list)
+            for expr in exprs_list:
                 if default_name not in self.__labels:
                     name = default_name
                 else:
@@ -125,7 +124,7 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
             elif label_cls is fo.Detections:
                 label = fo.Detections(detections=self.__as_fo_detections(label_data))
             else:
-                assert False
+                raise AssertionError()
             labels[label_name] = label
 
         return file, metadata, labels
@@ -166,5 +165,5 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
     def get_dataset_info(self) -> dict:
         pass
 
-    def close(self, *args) -> None:
+    def close(self, *args: Any) -> None:
         pass