pixeltable 0.3.12__tar.gz → 0.3.13__tar.gz

{pixeltable-0.3.12 → pixeltable-0.3.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pixeltable
-Version: 0.3.12
+Version: 0.3.13
 Summary: AI Data Infrastructure: Declarative, Multimodal, and Incremental
 License: Apache-2.0
 Keywords: data-science,machine-learning,database,ai,computer-vision,chatbot,ml,artificial-intelligence,feature-engineering,multimodal,mlops,feature-store,vector-database,llm,genai

{pixeltable-0.3.12 → pixeltable-0.3.13}/pixeltable/__init__.py

@@ -24,32 +24,7 @@ from .globals import (
     tool,
     tools,
 )
-from .type_system import (
-    Array,
-    ArrayType,
-    Audio,
-    AudioType,
-    Bool,
-    BoolType,
-    ColumnType,
-    Document,
-    DocumentType,
-    Float,
-    FloatType,
-    Image,
-    ImageType,
-    Int,
-    IntType,
-    Json,
-    JsonType,
-    Required,
-    String,
-    StringType,
-    Timestamp,
-    TimestampType,
-    Video,
-    VideoType,
-)
+from .type_system import Array, Audio, Bool, Document, Float, Image, Int, Json, Required, String, Timestamp, Video
 
 # This import must go last to avoid circular imports.
 from . import ext, functions, io, iterators  # isort: skip

{pixeltable-0.3.12 → pixeltable-0.3.13}/pixeltable/__version__.py

@@ -1,3 +1,3 @@
 # These version placeholders will be replaced during build.
-__version__ = '0.3.12'
-__version_tuple__ = (0, 3, 12)
+__version__ = '0.3.13'
+__version_tuple__ = (0, 3, 13)

{pixeltable-0.3.12 → pixeltable-0.3.13}/pixeltable/catalog/insertable_table.py

@@ -228,3 +228,14 @@ class InsertableTable(Table):
         """
         with Env.get().begin_xact():
             return self._tbl_version.get().delete(where=where)
+
+    @property
+    def _base_table(self) -> Optional['Table']:
+        return None
+
+    @property
+    def _effective_base_versions(self) -> list[Optional[int]]:
+        return []
+
+    def _table_descriptor(self) -> str:
+        return f'Table {self._path()!r}'

{pixeltable-0.3.12 → pixeltable-0.3.13}/pixeltable/catalog/table.py

@@ -109,7 +109,7 @@ class Table(SchemaObject):
         self._check_is_dropped()
         with env.Env.get().begin_xact():
             md = super().get_metadata()
-            md['base'] = self._base._path() if self._base is not None else None
+            md['base'] = self._base_table._path() if self._base_table is not None else None
             md['schema'] = self._schema
             md['is_replica'] = self._tbl_version.get().is_replica
             md['version'] = self._version
@@ -255,28 +255,27 @@ class Table(SchemaObject):
         return {c.name: c.col_type for c in self._tbl_version_path.columns()}
 
     @property
-    def _base(self) -> Optional['Table']:
-        """
-        The base table of this `Table`. If this table is a view, returns the `Table`
-        from which it was derived. Otherwise, returns `None`.
-        """
-        if self._tbl_version_path.base is None:
-            return None
-        base_id = self._tbl_version_path.base.tbl_version.id
-        return catalog.Catalog.get().get_table_by_id(base_id)
+    @abc.abstractmethod
+    def _base_table(self) -> Optional['Table']:
+        """The base's Table instance"""
+        ...
 
     @property
-    def _bases(self) -> list['Table']:
-        """
-        The ancestor list of bases of this table, starting with its immediate base.
-        """
+    def _base_tables(self) -> list['Table']:
+        """The ancestor list of bases of this table, starting with its immediate base."""
         bases = []
-        base = self._base
+        base = self._base_table
         while base is not None:
             bases.append(base)
-            base = base._base
+            base = base._base_table
         return bases
 
+    @property
+    @abc.abstractmethod
+    def _effective_base_versions(self) -> list[Optional[int]]:
+        """The effective versions of the ancestor bases, starting with its immediate base."""
+        ...
+
     @property
     def _comment(self) -> str:
         return self._tbl_version.get().comment
@@ -300,7 +299,7 @@ class Table(SchemaObject):
         Constructs a list of descriptors for this table that can be pretty-printed.
         """
         helper = DescriptionHelper()
-        helper.append(self._title_descriptor())
+        helper.append(self._table_descriptor())
         helper.append(self._col_descriptor())
         idxs = self._index_descriptor()
         if not idxs.empty:
@@ -312,14 +311,8 @@ class Table(SchemaObject):
             helper.append(f'COMMENT: {self._comment}')
         return helper
 
-    def _title_descriptor(self) -> str:
-        title: str
-        if self._base is None:
-            title = f'Table\n{self._path()!r}'
-        else:
-            title = f'View\n{self._path()!r}'
-            title += f'\n(of {self.__bases_to_desc()})'
-        return title
+    @abc.abstractmethod
+    def _table_descriptor(self) -> str: ...
 
     def _col_descriptor(self, columns: Optional[list[str]] = None) -> pd.DataFrame:
         return pd.DataFrame(
@@ -332,14 +325,6 @@ class Table(SchemaObject):
             if columns is None or col.name in columns
         )
 
-    def __bases_to_desc(self) -> str:
-        bases = self._bases
-        assert len(bases) >= 1
-        if len(bases) <= 2:
-            return ', '.join(repr(b._path()) for b in bases)
-        else:
-            return f'{bases[0]._path()!r}, ..., {bases[-1]._path()!r}'
-
     def _index_descriptor(self, columns: Optional[list[str]] = None) -> pd.DataFrame:
         from pixeltable import index
 
@@ -373,9 +358,9 @@ class Table(SchemaObject):
         """
         self._check_is_dropped()
         if getattr(builtins, '__IPYTHON__', False):
-            from IPython.display import display
+            from IPython.display import Markdown, display
 
-            display(self._repr_html_())
+            display(Markdown(self._repr_html_()))
         else:
             print(repr(self))
 

{pixeltable-0.3.12 → pixeltable-0.3.13}/pixeltable/catalog/view.py

@@ -267,3 +267,34 @@ class View(Table):
 
     def delete(self, where: Optional[exprs.Expr] = None) -> UpdateStatus:
         raise excs.Error(f'{self._display_name()} {self._name!r}: cannot delete from view')
+
+    @property
+    def _base_table(self) -> Optional['Table']:
+        # if this is a pure snapshot, our tbl_version_path only reflects the base (there is no TableVersion instance
+        # for the snapshot itself)
+        base_id = self._tbl_version.id if self._snapshot_only else self._tbl_version_path.base.tbl_version.id
+        return catalog.Catalog.get().get_table_by_id(base_id)
+
+    @property
+    def _effective_base_versions(self) -> list[Optional[int]]:
+        effective_versions = [tv.effective_version for tv in self._tbl_version_path.get_tbl_versions()]
+        if self._snapshot_only:
+            return effective_versions
+        else:
+            return effective_versions[1:]
+
+    def _table_descriptor(self) -> str:
+        display_name = 'Snapshot' if self._snapshot_only else 'View'
+        result = [f'{display_name} {self._path()!r}']
+        bases_descrs: list[str] = []
+        for base, effective_version in zip(self._base_tables, self._effective_base_versions):
+            if effective_version is None:
+                bases_descrs.append(f'{base._path()!r}')
+            else:
+                base_descr = f'{base._path()}:{effective_version}'
+                bases_descrs.append(f'{base_descr!r}')
+        result.append(f' (of {", ".join(bases_descrs)})')
+
+        if self._tbl_version.get().predicate is not None:
+            result.append(f'\nWhere: {self._tbl_version.get().predicate!s}')
+        return ''.join(result)

{pixeltable-0.3.12 → pixeltable-0.3.13}/pixeltable/dataframe.py

@@ -513,9 +513,9 @@ class DataFrame:
                 (select list, where clause, ...) vertically.
         """
         if getattr(builtins, '__IPYTHON__', False):
-            from IPython.display import display
+            from IPython.display import Markdown, display
 
-            display(self._repr_html_())
+            display(Markdown(self._repr_html_()))
         else:
             print(repr(self))
 

pixeltable-0.3.13/pixeltable/exceptions.py

@@ -0,0 +1,42 @@
+from types import TracebackType
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from pixeltable import exprs
+
+
+class Error(Exception):
+    pass
+
+
+class ExprEvalError(Exception):
+    expr: 'exprs.Expr'
+    expr_msg: str
+    exc: Exception
+    exc_tb: TracebackType
+    input_vals: list[Any]
+    row_num: int
+
+    def __init__(
+        self,
+        expr: 'exprs.Expr',
+        expr_msg: str,
+        exc: Exception,
+        exc_tb: TracebackType,
+        input_vals: list[Any],
+        row_num: int,
+    ) -> None:
+        exct = type(exc)
+        super().__init__(
+            f'Expression evaluation failed with an error of type `{exct.__module__}.{exct.__qualname__}`:\n{expr}'
+        )
+        self.expr = expr
+        self.expr_msg = expr_msg
+        self.exc = exc
+        self.exc_tb = exc_tb
+        self.input_vals = input_vals
+        self.row_num = row_num
+
+
+class PixeltableWarning(Warning):
+    pass

{pixeltable-0.3.12 → pixeltable-0.3.13}/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-0.3.12 → pixeltable-0.3.13}/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-0.3.12 → pixeltable-0.3.13}/pixeltable/exec/sql_node.py

@@ -424,6 +424,7 @@ class SqlAggregationNode(SqlNode):
     """
 
     group_by_items: Optional[list[exprs.Expr]]
+    input_cte: Optional[sql.CTE]
 
     def __init__(
         self,
@@ -440,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-0.3.12 → pixeltable-0.3.13}/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-0.3.12 → pixeltable-0.3.13}/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-0.3.12 → pixeltable-0.3.13}/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-0.3.13/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-0.3.12 → pixeltable-0.3.13}/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-0.3.12 → pixeltable-0.3.13}/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-0.3.12 → pixeltable-0.3.13}/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-0.3.12 → pixeltable-0.3.13}/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))
 
 
 #####################################