pixeltable 0.3.11__py3-none-any.whl → 0.3.13__py3-none-any.whl

pixeltable/exec/expr_eval/evaluators.py

@@ -208,10 +208,6 @@ class FnCallEvaluator(Evaluator):
             _logger.debug(f'Evaluated slot {self.fn_call.slot_idx} in {end_ts - start_ts}')
             self.dispatcher.dispatch([call_args.row], self.exec_ctx)
         except Exception as exc:
-            import anthropic
-
-            if isinstance(exc, anthropic.RateLimitError):
-                _logger.debug(f'RateLimitError: {exc}')
             _, _, exc_tb = sys.exc_info()
             call_args.row.set_exc(self.fn_call.slot_idx, exc)
             self.dispatcher.dispatch_exc(call_args.rows, self.fn_call.slot_idx, exc_tb, self.exec_ctx)

pixeltable/exec/expr_eval/expr_eval_node.py

@@ -282,7 +282,6 @@ class ExprEvalNode(ExecNode):
 
                 if self.exc_event.is_set():
                     # we got an exception that we need to propagate through __iter__()
-                    _logger.debug(f'Propagating exception {self.error}')
                     if isinstance(self.error, excs.ExprEvalError):
                         raise self.error from self.error.exc
                     else:

pixeltable/exec/sql_node.py

@@ -103,7 +103,6 @@ class SqlNode(ExecNode):
         # create Select stmt
         self.sql_elements = sql_elements
         self.tbl = tbl
-        assert all(not isinstance(e, exprs.Literal) for e in select_list)  # we're never asked to materialize literals
         self.select_list = exprs.ExprSet(select_list)
         # unstored iter columns: we also need to retrieve whatever is needed to materialize the iter args
         for iter_arg in row_builder.unstored_iter_args.values():
@@ -425,6 +424,7 @@ class SqlAggregationNode(SqlNode):
     """
 
     group_by_items: Optional[list[exprs.Expr]]
+    input_cte: Optional[sql.CTE]
 
     def __init__(
         self,
@@ -441,13 +441,13 @@ class SqlAggregationNode(SqlNode):
             group_by_items: list of expressions to group by
             limit: max number of rows to return: None = no limit
         """
-        _, input_col_map = input.to_cte()
+        self.input_cte, input_col_map = input.to_cte()
         sql_elements = exprs.SqlElementCache(input_col_map)
         super().__init__(None, row_builder, select_list, sql_elements)
         self.group_by_items = group_by_items
 
     def _create_stmt(self) -> sql.Select:
-        stmt = super()._create_stmt()
+        stmt = super()._create_stmt().select_from(self.input_cte)
         if self.group_by_items is not None:
             sql_group_by_items = [self.sql_elements.get(e) for e in self.group_by_items]
             assert all(e is not None for e in sql_group_by_items)

pixeltable/exprs/json_path.py

@@ -5,7 +5,6 @@ from typing import Any, Optional, Union
 import jmespath
 import sqlalchemy as sql
 
-import pixeltable as pxt
 from pixeltable import catalog, exceptions as excs, type_system as ts
 
 from .data_row import DataRow
@@ -19,10 +18,7 @@ from .sql_element_cache import SqlElementCache
 
 class JsonPath(Expr):
     def __init__(
-        self,
-        anchor: Optional['pxt.exprs.Expr'],
-        path_elements: Optional[list[Union[str, int, slice]]] = None,
-        scope_idx: int = 0,
+        self, anchor: Optional[Expr], path_elements: Optional[list[Union[str, int, slice]]] = None, scope_idx: int = 0
     ) -> None:
         """
         anchor can be None, in which case this is a relative JsonPath and the anchor is set later via set_anchor().

pixeltable/func/__init__.py

@@ -5,7 +5,7 @@ from .callable_function import CallableFunction
 from .expr_template_function import ExprTemplateFunction
 from .function import Function, InvalidFunction
 from .function_registry import FunctionRegistry
-from .query_template_function import QueryTemplateFunction, query
+from .query_template_function import QueryTemplateFunction, query, retrieval_udf
 from .signature import Batch, Parameter, Signature
 from .tools import Tool, ToolChoice, Tools
 from .udf import expr_udf, make_function, udf

pixeltable/func/aggregate_function.py

@@ -159,7 +159,7 @@ class AggregateFunction(Function):
         self.init_param_names.append(init_param_names)
         return self
 
-    def _docstring(self) -> Optional[str]:
+    def comment(self) -> Optional[str]:
         return inspect.getdoc(self.agg_classes[0])
 
     def help_str(self) -> str:

pixeltable/func/callable_function.py

@@ -60,7 +60,7 @@ class CallableFunction(Function):
     def is_async(self) -> bool:
         return inspect.iscoroutinefunction(self.py_fn)
 
-    def _docstring(self) -> Optional[str]:
+    def comment(self) -> Optional[str]:
         return inspect.getdoc(self.py_fns[0])
 
     @property

pixeltable/func/expr_template_function.py

@@ -95,9 +95,9 @@ class ExprTemplateFunction(Function):
         )
         return substituted_expr.col_type
 
-    def _docstring(self) -> Optional[str]:
+    def comment(self) -> Optional[str]:
         if isinstance(self.templates[0].expr, exprs.FunctionCall):
-            return self.templates[0].expr.fn._docstring()
+            return self.templates[0].expr.fn.comment()
         return None
 
     def exec(self, args: Sequence[Any], kwargs: dict[str, Any]) -> Any:

pixeltable/func/function.py

@@ -10,8 +10,7 @@ from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence, cast
 import sqlalchemy as sql
 from typing_extensions import Self
 
-import pixeltable.exceptions as excs
-import pixeltable.type_system as ts
+from pixeltable import exceptions as excs, type_system as ts
 
 from .globals import resolve_symbol
 from .signature import Signature
@@ -106,11 +105,11 @@ class Function(ABC):
     @abstractmethod
     def is_async(self) -> bool: ...
 
-    def _docstring(self) -> Optional[str]:
+    def comment(self) -> Optional[str]:
         return None
 
     def help_str(self) -> str:
-        docstring = self._docstring()
+        docstring = self.comment()
         display = self.display_name + str(self.signatures[0])
         if docstring is None:
             return display

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/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

@@ -262,7 +262,7 @@ def from_table(
     """
     from pixeltable import exprs
 
-    ancestors = [tbl, *tbl._bases]
+    ancestors = [tbl, *tbl._base_tables]
     ancestors.reverse()  # We must traverse the ancestors in order from base to derived
 
     subst: dict[exprs.Expr, exprs.Expr] = {}

pixeltable/functions/__init__.py

@@ -5,6 +5,7 @@ from pixeltable.utils.code import local_public_names
 from . import (
     anthropic,
     audio,
+    bedrock,
     deepseek,
     fireworks,
     gemini,

pixeltable/functions/anthropic.py

@@ -112,7 +112,7 @@ async def messages(
         to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
 
         >>> msgs = [{'role': 'user', 'content': tbl.prompt}]
-        ... tbl.add_computed_column(response= messages(msgs, model='claude-3-haiku-20240307'))
+        ... tbl.add_computed_column(response=messages(msgs, model='claude-3-haiku-20240307'))
     """
 
     # it doesn't look like count_tokens() actually exists in the current version of the library

pixeltable/functions/bedrock.py

@@ -0,0 +1,130 @@
+import logging
+from typing import TYPE_CHECKING, Any, Optional
+
+import pixeltable as pxt
+from pixeltable import env, exprs
+from pixeltable.func import Tools
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    from botocore.client import BaseClient
+
+_logger = logging.getLogger('pixeltable')
+
+
+@env.register_client('bedrock')
+def _() -> 'BaseClient':
+    import boto3
+
+    return boto3.client(service_name='bedrock-runtime')
+
+
+# boto3 typing is weird; type information is dynamically defined, so the best we can do for the static checker is `Any`
+def _bedrock_client() -> Any:
+    return env.Env.get().get_client('bedrock')
+
+
+@pxt.udf
+def converse(
+    messages: list[dict[str, Any]],
+    *,
+    model_id: str,
+    system: Optional[list[dict[str, Any]]] = None,
+    inference_config: Optional[dict] = None,
+    additional_model_request_fields: Optional[dict] = None,
+    tool_config: Optional[list[dict]] = None,
+) -> dict:
+    """
+    Generate a conversation response.
+
+    Equivalent to the AWS Bedrock `converse` API endpoint.
+    For additional details, see: <https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html>
+
+    __Requirements:__
+
+    - `pip install boto3`
+
+    Args:
+        messages: Input messages.
+        model_id: The model that will complete your prompt.
+        system: An optional system prompt.
+        inference_config: Base inference parameters to use.
+        additional_model_request_fields: Additional inference parameters to use.
+
+    For details on the optional parameters, see:
+    <https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html>
+
+    Returns:
+        A dictionary containing the response and other metadata.
+
+    Examples:
+        Add a computed column that applies the model `anthropic.claude-3-haiku-20240307-v1:0`
+        to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
+
+        >>> msgs = [{'role': 'user', 'content': [{'text': tbl.prompt}]}]
+        ... tbl.add_computed_column(response=messages(msgs, model_id='anthropic.claude-3-haiku-20240307-v1:0'))
+    """
+
+    kwargs: dict[str, Any] = {'messages': messages, 'modelId': model_id}
+
+    if system is not None:
+        kwargs['system'] = system
+    if inference_config is not None:
+        kwargs['inferenceConfig'] = inference_config
+    if additional_model_request_fields is not None:
+        kwargs['additionalModelRequestFields'] = additional_model_request_fields
+
+    if tool_config is not None:
+        tool_config_ = {
+            'tools': [
+                {
+                    'toolSpec': {
+                        'name': tool['name'],
+                        'description': tool['description'],
+                        'inputSchema': {
+                            'json': {
+                                'type': 'object',
+                                'properties': tool['parameters']['properties'],
+                                'required': tool['required'],
+                            }
+                        },
+                    }
+                }
+                for tool in tool_config
+            ]
+        }
+        kwargs['toolConfig'] = tool_config_
+
+    return _bedrock_client().converse(**kwargs)
+
+
+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(_bedrock_response_to_pxt_tool_calls(response))
+
+
+@pxt.udf
+def _bedrock_response_to_pxt_tool_calls(response: dict) -> Optional[dict]:
+    if response.get('stopReason') != 'tool_use':
+        return None
+
+    pxt_tool_calls: dict[str, list[dict[str, Any]]] = {}
+    for message in response['output']['message']['content']:
+        if 'toolUse' in message:
+            tool_call = message['toolUse']
+            tool_name = tool_call['name']
+            if tool_name not in pxt_tool_calls:
+                pxt_tool_calls[tool_name] = []
+            pxt_tool_calls[tool_name].append({'args': tool_call['input']})
+
+    if len(pxt_tool_calls) == 0:
+        return None
+
+    return pxt_tool_calls
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/functions/huggingface.py

@@ -13,6 +13,7 @@ import PIL.Image
 
 import pixeltable as pxt
 import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
 from pixeltable import env
 from pixeltable.func import Batch
 from pixeltable.functions.util import normalize_image_mode, resolve_torch_device
@@ -61,14 +62,14 @@ def sentence_transformer(
 
 
 @sentence_transformer.conditional_return_type
-def _(model_id: str) -> pxt.ArrayType:
+def _(model_id: str) -> ts.ArrayType:
     try:
         from sentence_transformers import SentenceTransformer
 
         model = _lookup_model(model_id, SentenceTransformer)
-        return pxt.ArrayType((model.get_sentence_embedding_dimension(),), dtype=pxt.FloatType(), nullable=False)
+        return ts.ArrayType((model.get_sentence_embedding_dimension(),), dtype=ts.FloatType(), nullable=False)
     except ImportError:
-        return pxt.ArrayType((None,), dtype=pxt.FloatType(), nullable=False)
+        return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
 
 
 @pxt.udf
@@ -199,14 +200,14 @@ def _(image: Batch[PIL.Image.Image], *, model_id: str) -> Batch[pxt.Array[(None,
 
 
 @clip.conditional_return_type
-def _(model_id: str) -> pxt.ArrayType:
+def _(model_id: str) -> ts.ArrayType:
     try:
         from transformers import CLIPModel
 
         model = _lookup_model(model_id, CLIPModel.from_pretrained)
-        return pxt.ArrayType((model.config.projection_dim,), dtype=pxt.FloatType(), nullable=False)
+        return ts.ArrayType((model.config.projection_dim,), dtype=ts.FloatType(), nullable=False)
     except ImportError:
-        return pxt.ArrayType((None,), dtype=pxt.FloatType(), nullable=False)
+        return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
 
 
 @pxt.udf(batch_size=4)

pixeltable/functions/image.py

@@ -16,6 +16,7 @@ from typing import Optional
 import PIL.Image
 
 import pixeltable as pxt
+import pixeltable.type_system as ts
 from pixeltable.exprs import Expr
 from pixeltable.utils.code import local_public_names
 
@@ -88,10 +89,10 @@ def convert(self: PIL.Image.Image, mode: str) -> PIL.Image.Image:
 
 
 @convert.conditional_return_type
-def _(self: Expr, mode: str) -> pxt.ColumnType:
+def _(self: Expr, mode: str) -> ts.ColumnType:
     input_type = self.col_type
-    assert isinstance(input_type, pxt.ImageType)
-    return pxt.ImageType(size=input_type.size, mode=mode, nullable=input_type.nullable)
+    assert isinstance(input_type, ts.ImageType)
+    return ts.ImageType(size=input_type.size, mode=mode, nullable=input_type.nullable)
 
 
 # Image.crop()
@@ -108,14 +109,12 @@ def crop(self: PIL.Image.Image, box: tuple[int, int, int, int]) -> PIL.Image.Ima
 
 
 @crop.conditional_return_type
-def _(self: Expr, box: tuple[int, int, int, int]) -> pxt.ColumnType:
+def _(self: Expr, box: tuple[int, int, int, int]) -> ts.ColumnType:
     input_type = self.col_type
-    assert isinstance(input_type, pxt.ImageType)
+    assert isinstance(input_type, ts.ImageType)
     if (isinstance(box, (list, tuple))) and len(box) == 4 and all(isinstance(x, int) for x in box):
-        return pxt.ImageType(
-            size=(box[2] - box[0], box[3] - box[1]), mode=input_type.mode, nullable=input_type.nullable
-        )
-    return pxt.ImageType(mode=input_type.mode, nullable=input_type.nullable)  # we can't compute the size statically
+        return ts.ImageType(size=(box[2] - box[0], box[3] - box[1]), mode=input_type.mode, nullable=input_type.nullable)
+    return ts.ImageType(mode=input_type.mode, nullable=input_type.nullable)  # we can't compute the size statically
 
 
 # Image.getchannel()
@@ -134,10 +133,10 @@ def getchannel(self: PIL.Image.Image, channel: int) -> PIL.Image.Image:
 
 
 @getchannel.conditional_return_type
-def _(self: Expr) -> pxt.ColumnType:
+def _(self: Expr) -> ts.ColumnType:
     input_type = self.col_type
-    assert isinstance(input_type, pxt.ImageType)
-    return pxt.ImageType(size=input_type.size, mode='L', nullable=input_type.nullable)
+    assert isinstance(input_type, ts.ImageType)
+    return ts.ImageType(size=input_type.size, mode='L', nullable=input_type.nullable)
 
 
 @pxt.udf(is_method=True)
@@ -183,10 +182,10 @@ def resize(self: PIL.Image.Image, size: tuple[int, int]) -> PIL.Image.Image:
 
 
 @resize.conditional_return_type
-def _(self: Expr, size: tuple[int, int]) -> pxt.ColumnType:
+def _(self: Expr, size: tuple[int, int]) -> ts.ColumnType:
     input_type = self.col_type
-    assert isinstance(input_type, pxt.ImageType)
-    return pxt.ImageType(size=size, mode=input_type.mode, nullable=input_type.nullable)
+    assert isinstance(input_type, ts.ImageType)
+    return ts.ImageType(size=size, mode=input_type.mode, nullable=input_type.nullable)
 
 
 # Image.rotate()
@@ -237,7 +236,7 @@ def transpose(self: PIL.Image.Image, method: int) -> PIL.Image.Image:
 @rotate.conditional_return_type
 @effect_spread.conditional_return_type
 @transpose.conditional_return_type
-def _(self: Expr) -> pxt.ColumnType:
+def _(self: Expr) -> ts.ColumnType:
     return self.col_type
 
 

pixeltable/functions/mistralai.py

@@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Optional, TypeVar, Union
 import numpy as np
 
 import pixeltable as pxt
+import pixeltable.type_system as ts
 from pixeltable.env import Env, register_client
 from pixeltable.func.signature import Batch
 from pixeltable.utils.code import local_public_names
@@ -176,9 +177,9 @@ async def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,
 
 
 @embeddings.conditional_return_type
-def _(model: str) -> pxt.ArrayType:
+def _(model: str) -> ts.ArrayType:
     dimensions = _embedding_dimensions_cache.get(model)  # `None` if unknown model
-    return pxt.ArrayType((dimensions,), dtype=pxt.FloatType())
+    return ts.ArrayType((dimensions,), dtype=ts.FloatType())
 
 
 _T = TypeVar('_T')

pixeltable/functions/openai.py

@@ -21,6 +21,7 @@ import numpy as np
 import PIL
 
 import pixeltable as pxt
+import pixeltable.type_system as ts
 from pixeltable import env, exprs
 from pixeltable.func import Batch, Tools
 from pixeltable.utils.code import local_public_names
@@ -666,13 +667,13 @@ async def embeddings(
 
 
 @embeddings.conditional_return_type
-def _(model: str, dimensions: Optional[int] = None) -> pxt.ArrayType:
+def _(model: str, dimensions: Optional[int] = None) -> ts.ArrayType:
     if dimensions is None:
         if model not in _embedding_dimensions_cache:
             # TODO: find some other way to retrieve a sample
-            return pxt.ArrayType((None,), dtype=pxt.FloatType(), nullable=False)
+            return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
         dimensions = _embedding_dimensions_cache.get(model)
-    return pxt.ArrayType((dimensions,), dtype=pxt.FloatType(), nullable=False)
+    return ts.ArrayType((dimensions,), dtype=ts.FloatType(), nullable=False)
 
 
 #####################################
@@ -738,17 +739,17 @@ async def image_generations(
 
 
 @image_generations.conditional_return_type
-def _(size: Optional[str] = None) -> pxt.ImageType:
+def _(size: Optional[str] = None) -> ts.ImageType:
     if size is None:
-        return pxt.ImageType(size=(1024, 1024))
+        return ts.ImageType(size=(1024, 1024))
     x_pos = size.find('x')
     if x_pos == -1:
-        return pxt.ImageType()
+        return ts.ImageType()
     try:
         width, height = int(size[:x_pos]), int(size[x_pos + 1 :])
     except ValueError:
-        return pxt.ImageType()
-    return pxt.ImageType(size=(width, height))
+        return ts.ImageType()
+    return ts.ImageType(size=(width, height))
 
 
 #####################################

pixeltable/functions/together.py

@@ -16,6 +16,7 @@ import tenacity
 
 import pixeltable as pxt
 import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
 from pixeltable import env
 from pixeltable.func import Batch
 from pixeltable.utils.code import local_public_names
@@ -225,12 +226,12 @@ async def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,
 
 
 @embeddings.conditional_return_type
-def _(model: str) -> pxt.ArrayType:
+def _(model: str) -> ts.ArrayType:
     if model not in _embedding_dimensions_cache:
         # TODO: find some other way to retrieve a sample
-        return pxt.ArrayType((None,), dtype=pxt.FloatType())
+        return ts.ArrayType((None,), dtype=ts.FloatType())
     dimensions = _embedding_dimensions_cache[model]
-    return pxt.ArrayType((dimensions,), dtype=pxt.FloatType())
+    return ts.ArrayType((dimensions,), dtype=ts.FloatType())
 
 
 @pxt.udf(resource_pool='request-rate:together:images')