pixeltable 0.4.14__py3-none-any.whl → 0.4.16__py3-none-any.whl

pixeltable/exprs/expr_set.py

@@ -9,26 +9,33 @@ T = TypeVar('T', bound='Expr')
 
 class ExprSet(Generic[T]):
     """
-    A set that also supports indexed lookup (by slot_idx and Expr.id). Exprs are uniquely identified by Expr.id.
+    An ordered set that also supports indexed lookup (by slot_idx and Expr.id). Exprs are uniquely identified by
+    Expr.id.
     """
 
     exprs: dict[int, T]  # key: Expr.id
+    expr_offsets: dict[int, int]  # key: Expr.id, value: offset into self.exprs.keys()
     exprs_by_idx: dict[int, T]  # key: slot_idx
 
     def __init__(self, elements: Optional[Iterable[T]] = None):
         self.exprs = {}
+        self.expr_offsets = {}
         self.exprs_by_idx = {}
         if elements is not None:
             for e in elements:
                 self.add(e)
 
-    def add(self, expr: T) -> None:
-        if expr.id in self.exprs:
-            return
+    def add(self, expr: T) -> int:
+        """Returns offset corresponding to iteration order"""
+        offset = self.expr_offsets.get(expr.id)
+        if offset is not None:
+            return offset
+        offset = len(self.exprs)
         self.exprs[expr.id] = expr
-        if expr.slot_idx is None:
-            return
-        self.exprs_by_idx[expr.slot_idx] = expr
+        self.expr_offsets[expr.id] = offset
+        if expr.slot_idx is not None:
+            self.exprs_by_idx[expr.slot_idx] = expr
+        return offset
 
     def update(self, *others: Iterable[T]) -> None:
         for other in others:

pixeltable/exprs/inline_expr.py

@@ -98,13 +98,7 @@ class InlineList(Expr):
     def __init__(self, elements: Iterable):
         exprs = [Expr.from_object(el) for el in elements]
 
-        json_schema = {
-            'type': 'array',
-            'prefixItems': [expr.col_type.to_json_schema() for expr in exprs],
-            'items': False,  # No additional items (fixed length)
-        }
-
-        super().__init__(ts.JsonType(json_schema))
+        super().__init__(ts.JsonType())
         self.components.extend(exprs)
         self.id = self._create_id()
 
@@ -150,18 +144,7 @@ class InlineDict(Expr):
             self.keys.append(key)
             exprs.append(Expr.from_object(val))
 
-        json_schema: Optional[dict[str, Any]]
-        try:
-            json_schema = {
-                'type': 'object',
-                'properties': {key: expr.col_type.to_json_schema() for key, expr in zip(self.keys, exprs)},
-            }
-        except excs.Error:
-            # InlineDicts are used to store iterator arguments, which are not required to be valid JSON types,
-            # so we can't always construct a valid schema.
-            json_schema = None
-
-        super().__init__(ts.JsonType(json_schema))
+        super().__init__(ts.JsonType())
         self.components.extend(exprs)
         self.id = self._create_id()
 

pixeltable/exprs/json_path.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import io
+from pathlib import Path
 from typing import Any, Optional
 
 import jmespath
@@ -7,6 +9,7 @@ import sqlalchemy as sql
 
 from pixeltable import catalog, exceptions as excs, type_system as ts
 
+from .column_ref import ColumnRef
 from .data_row import DataRow
 from .expr import Expr
 from .globals import print_slice
@@ -23,6 +26,11 @@ class JsonPath(Expr):
     (0: indicates the immediately preceding JsonMapper, -1: the parent of the immediately preceding mapper, ...)
     """
 
+    path_elements: list[str | int | slice]
+    compiled_path: jmespath.parser.ParsedResult | None
+    scope_idx: int
+    file_handles: dict[Path, io.BufferedReader]  # key: file path
+
     def __init__(
         self, anchor: Optional[Expr], path_elements: Optional[list[str | int | slice]] = None, scope_idx: int = 0
     ) -> None:
@@ -31,16 +39,22 @@ class JsonPath(Expr):
         super().__init__(ts.JsonType(nullable=True))  # JsonPath expressions are always nullable
         if anchor is not None:
             self.components = [anchor]
-        self.path_elements: list[str | int | slice] = path_elements
+        self.path_elements = path_elements
         self.compiled_path = jmespath.compile(self._json_path()) if len(path_elements) > 0 else None
         self.scope_idx = scope_idx
         # NOTE: the _create_id() result will change if set_anchor() gets called;
         # this is not a problem, because _create_id() shouldn't be called after init()
         self.id = self._create_id()
+        self.file_handles = {}
+
+    def release(self) -> None:
+        for fh in self.file_handles.values():
+            fh.close()
+        self.file_handles.clear()
 
     def __repr__(self) -> str:
         # else 'R': the anchor is RELATIVE_PATH_ROOT
-        anchor_str = str(self._anchor) if self._anchor is not None else 'R'
+        anchor_str = str(self.anchor) if self.anchor is not None else 'R'
         if len(self.path_elements) == 0:
             return anchor_str
         return f'{anchor_str}{"." if isinstance(self.path_elements[0], str) else ""}{self._json_path()}'
@@ -67,7 +81,7 @@ class JsonPath(Expr):
         return cls(anchor, path_elements, d['scope_idx'])
 
     @property
-    def _anchor(self) -> Optional[Expr]:
+    def anchor(self) -> Optional[Expr]:
         return None if len(self.components) == 0 else self.components[0]
 
     def set_anchor(self, anchor: Expr) -> None:
@@ -75,7 +89,7 @@ class JsonPath(Expr):
         self.components = [anchor]
 
     def is_relative_path(self) -> bool:
-        return self._anchor is None
+        return self.anchor is None
 
     def _has_relative_path(self) -> bool:
         return self.is_relative_path() or super()._has_relative_path()
@@ -85,7 +99,7 @@ class JsonPath(Expr):
             # TODO: take scope_idx into account
             self.set_anchor(mapper.scope_anchor)
         else:
-            self._anchor._bind_rel_paths(mapper)
+            self.anchor._bind_rel_paths(mapper)
 
     def __call__(self, *args: object, **kwargs: object) -> 'JsonPath':
         """
@@ -99,15 +113,15 @@ class JsonPath(Expr):
 
     def __getattr__(self, name: str) -> 'JsonPath':
         assert isinstance(name, str)
-        return JsonPath(self._anchor, [*self.path_elements, name])
+        return JsonPath(self.anchor, [*self.path_elements, name])
 
     def __getitem__(self, index: object) -> 'JsonPath':
         if isinstance(index, (int, slice, str)):
-            return JsonPath(self._anchor, [*self.path_elements, index])
+            return JsonPath(self.anchor, [*self.path_elements, index])
         raise excs.Error(f'Invalid json list index: {index}')
 
     def default_column_name(self) -> Optional[str]:
-        anchor_name = self._anchor.default_column_name() if self._anchor is not None else ''
+        anchor_name = self.anchor.default_column_name() if self.anchor is not None else ''
         ret_name = f'{anchor_name}.{self._json_path()}'
 
         def cleanup_char(s: str) -> str:
@@ -159,12 +173,31 @@ class JsonPath(Expr):
                 result.append(f'[{print_slice(element)}]')
         return ''.join(result)
 
-    def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
-        assert self._anchor is not None, self
-        val = data_row[self._anchor.slot_idx]
+    def eval(self, row: DataRow, row_builder: RowBuilder) -> None:
+        assert self.anchor is not None, self
+        val = row[self.anchor.slot_idx]
         if self.compiled_path is not None:
             val = self.compiled_path.search(val)
-        data_row[self.slot_idx] = val
+        row[self.slot_idx] = val
+        if val is None or self.anchor is None or not isinstance(self.anchor, ColumnRef):
+            return
+
+        # the origin of val is a json-typed column, which might stored inlined objects
+        if self.anchor.slot_idx not in row.slot_md:
+            # we can infer that there aren't any inlined objects because our execution plan doesn't include
+            # materializing the cellmd (eg, insert plans)
+            # TODO: have the planner pass that fact into ExprEvalNode explicitly to streamline this path a bit more
+            return
+
+        # defer import until it's needed
+        from pixeltable.exec.cell_reconstruction_node import json_has_inlined_objs, reconstruct_json
+
+        cell_md = row.slot_md[self.anchor.slot_idx]
+        if cell_md is None or cell_md.file_urls is None or not json_has_inlined_objs(val):
+            # val doesn't contain inlined objects
+            return
+
+        row.vals[self.slot_idx] = reconstruct_json(val, cell_md.file_urls, self.file_handles)
 
 
 RELATIVE_PATH_ROOT = JsonPath(None)

pixeltable/exprs/row_builder.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
+import dataclasses
 import sys
 import time
-from dataclasses import dataclass
 from typing import Any, Iterable, NamedTuple, Optional, Sequence
 from uuid import UUID
 
 import numpy as np
+import sqlalchemy as sql
 
 from pixeltable import catalog, exceptions as excs, exprs, utils
 from pixeltable.env import Env
+from pixeltable.utils.misc import non_none_dict_factory
 
 from .data_row import DataRow
 from .expr import Expr, ExprScope
@@ -68,7 +70,7 @@ class RowBuilder:
     input_exprs: ExprSet
 
     tbl: Optional[catalog.TableVersion]  # reference table of the RowBuilder; used to identify pk columns for writes
-    table_columns: list[ColumnSlotIdx]
+    table_columns: dict[catalog.Column, int | None]  # value: slot idx, if the result of an expr
     default_eval_ctx: EvalCtx
     unstored_iter_args: dict[UUID, Expr]
 
@@ -92,10 +94,9 @@ class RowBuilder:
     img_slot_idxs: list[int]  # Indices of image slots
     media_slot_idxs: list[int]  # Indices of non-image media slots
     array_slot_idxs: list[int]  # Indices of array slots
-    stored_img_cols: list[exprs.ColumnSlotIdx]
-    stored_media_cols: list[exprs.ColumnSlotIdx]
+    json_slot_idxs: list[int]  # Indices of json slots
 
-    @dataclass
+    @dataclasses.dataclass
     class EvalCtx:
         """Context for evaluating a set of target exprs"""
 
@@ -113,8 +114,6 @@ class RowBuilder:
     ):
         self.unique_exprs: ExprSet[Expr] = ExprSet()  # dependencies precede their dependents
         self.next_slot_idx = 0
-        self.stored_img_cols = []
-        self.stored_media_cols = []
 
         # record input and output exprs; make copies to avoid reusing execution state
         unique_input_exprs = [self._record_unique_expr(e.copy(), recursive=False) for e in input_exprs]
@@ -138,7 +137,7 @@ class RowBuilder:
         from .column_ref import ColumnRef
 
         self.tbl = tbl
-        self.table_columns: list[ColumnSlotIdx] = []
+        self.table_columns = {}
         self.input_exprs = ExprSet()
         validating_colrefs: dict[Expr, Expr] = {}  # key: non-validating colref, value: corresp. validating colref
         for col in columns:
@@ -245,17 +244,27 @@ class RowBuilder:
             e.slot_idx for e in self.unique_exprs if e.col_type.is_media_type() and not e.col_type.is_image_type()
         ]
         self.array_slot_idxs = [e.slot_idx for e in self.unique_exprs if e.col_type.is_array_type()]
+        self.json_slot_idxs = [e.slot_idx for e in self.unique_exprs if e.col_type.is_json_type()]
 
     def add_table_column(self, col: catalog.Column, slot_idx: int) -> None:
-        """Record a column that is part of the table row"""
+        """Record an output column for which the value is produced via expr evaluation"""
         assert self.tbl is not None
         assert col.is_stored
-        info = ColumnSlotIdx(col, slot_idx)
-        self.table_columns.append(info)
-        if col.col_type.is_media_type():
-            self.stored_media_cols.append(info)
-            if col.col_type.is_image_type():
-                self.stored_img_cols.append(info)
+        self.table_columns[col] = slot_idx
+
+    def add_table_columns(self, cols: list[catalog.Column]) -> None:
+        """Record output columns whose values are materialized into DataRow.cell_vals"""
+        for col in cols:
+            self.table_columns[col] = None
+
+    @property
+    def media_output_col_info(self) -> list[ColumnSlotIdx]:
+        """Return slot idxs for media output columns whose values are produced by expr evaluation"""
+        return [
+            ColumnSlotIdx(col, slot_idx)
+            for col, slot_idx in self.table_columns.items()
+            if col.col_type.is_media_type() and slot_idx is not None
+        ]
 
     @property
     def num_materialized(self) -> int:
@@ -462,13 +471,30 @@ class RowBuilder:
 
         num_excs = 0
         table_row: list[Any] = list(pk)
-        for col, slot_idx in self.table_columns:
+        # Nulls in JSONB columns need to be stored as sql.sql.null(), otherwise it stores a json 'null'
+        for col, slot_idx in self.table_columns.items():
+            if col.id in data_row.cell_vals:
+                table_row.append(data_row.cell_vals[col.id])
+                if col.stores_cellmd:
+                    if data_row.cell_md[col.id] is None:
+                        table_row.append(sql.sql.null())
+                    else:
+                        # we want to minimize the size of the stored dict and use dict_factory to remove Nones
+                        md = dataclasses.asdict(data_row.cell_md[col.id], dict_factory=non_none_dict_factory)
+                        assert len(md) > 0
+                        table_row.append(md)
+                if slot_idx is not None and data_row.has_exc(slot_idx):
+                    num_excs += 1
+                    if cols_with_excs is not None:
+                        cols_with_excs.add(col.id)
+                continue
+
             if data_row.has_exc(slot_idx):
                 exc = data_row.get_exc(slot_idx)
                 num_excs += 1
                 if cols_with_excs is not None:
                     cols_with_excs.add(col.id)
-                table_row.append(None)
+                table_row.append(sql.sql.null() if col.col_type.is_json_type() else None)
                 if col.stores_cellmd:
                     # exceptions get stored in the errortype/-msg properties of the cellmd column
                     table_row.append(ColumnPropertyRef.create_cellmd_exc(exc))
@@ -476,7 +502,7 @@ class RowBuilder:
                 val = data_row.get_stored_val(slot_idx, col.get_sa_col_type())
                 table_row.append(val)
                 if col.stores_cellmd:
-                    table_row.append(None)  # placeholder for cellmd column
+                    table_row.append(sql.sql.null())  # placeholder for cellmd column
 
         return table_row, num_excs
 
@@ -490,12 +516,18 @@ class RowBuilder:
         store_col_names: list[str] = [pk_col.name for pk_col in self.tbl.store_tbl.pk_columns()]
 
         for col in self.table_columns:
-            store_col_names.append(col.col.store_name())
-            if col.col.stores_cellmd:
-                store_col_names.append(col.col.cellmd_store_name())
+            store_col_names.append(col.store_name())
+            if col.stores_cellmd:
+                store_col_names.append(col.cellmd_store_name())
 
         return store_col_names
 
     def make_row(self) -> exprs.DataRow:
         """Creates a new DataRow with the current row_builder's configuration."""
-        return exprs.DataRow(self.num_materialized, self.img_slot_idxs, self.media_slot_idxs, self.array_slot_idxs)
+        return exprs.DataRow(
+            size=self.num_materialized,
+            img_slot_idxs=self.img_slot_idxs,
+            media_slot_idxs=self.media_slot_idxs,
+            array_slot_idxs=self.array_slot_idxs,
+            json_slot_idxs=self.json_slot_idxs,
+        )

pixeltable/functions/__init__.py

@@ -19,6 +19,7 @@ from . import (
     mistralai,
     ollama,
     openai,
+    openrouter,
     replicate,
     string,
     timestamp,

pixeltable/functions/bedrock.py

@@ -1,3 +1,10 @@
+"""
+Pixeltable UDFs for AWS Bedrock AI models.
+
+Provides integration with AWS Bedrock for accessing various foundation models
+including Anthropic Claude, Amazon Titan, and other providers.
+"""
+
 import logging
 from typing import TYPE_CHECKING, Any, Optional
 

pixeltable/functions/deepseek.py

@@ -1,3 +1,10 @@
+"""
+Pixeltable UDFs for Deepseek AI models.
+
+Provides integration with Deepseek's language models for chat completions
+and other AI capabilities.
+"""
+
 import json
 from typing import TYPE_CHECKING, Any, Optional
 
@@ -67,10 +74,10 @@ async def chat_completions(
         of the table `tbl`:
 
         >>> messages = [
-                {'role': 'system', 'content': 'You are a helpful assistant.'},
-                {'role': 'user', 'content': tbl.prompt}
-            ]
-            tbl.add_computed_column(response=chat_completions(messages, model='deepseek-chat'))
+        ...     {'role': 'system', 'content': 'You are a helpful assistant.'},
+        ...     {'role': 'user', 'content': tbl.prompt}
+        ... ]
+        >>> tbl.add_computed_column(response=chat_completions(messages, model='deepseek-chat'))
     """
     if model_kwargs is None:
         model_kwargs = {}

pixeltable/functions/llama_cpp.py

@@ -1,3 +1,10 @@
+"""
+Pixeltable UDFs for llama.cpp models.
+
+Provides integration with llama.cpp for running quantized language models locally,
+supporting chat completions and embeddings with GGUF format models.
+"""
+
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Optional
 

pixeltable/functions/math.py

@@ -97,7 +97,7 @@ def _(self: sql.ColumnElement, digits: Optional[sql.ColumnElement] = None) -> sq
     if digits is None:
         return sql.func.round(self)
     else:
-        return sql.func.round(self.cast(sql.Numeric), digits.cast(sql.Integer))
+        return sql.cast(sql.func.round(sql.cast(self, sql.Numeric), sql.cast(digits, sql.Integer)), sql.Float)
 
 
 @pxt.udf(is_method=True)

pixeltable/functions/ollama.py

@@ -1,3 +1,10 @@
+"""
+Pixeltable UDFs for Ollama local models.
+
+Provides integration with Ollama for running large language models locally,
+including chat completions and embeddings.
+"""
+
 from typing import TYPE_CHECKING, Optional
 
 import numpy as np

pixeltable/functions/openai.py

@@ -395,10 +395,10 @@ async def chat_completions(
         of the table `tbl`:
 
         >>> messages = [
-                {'role': 'system', 'content': 'You are a helpful assistant.'},
-                {'role': 'user', 'content': tbl.prompt}
-            ]
-            tbl.add_computed_column(response=chat_completions(messages, model='gpt-4o-mini'))
+        ...     {'role': 'system', 'content': 'You are a helpful assistant.'},
+        ...     {'role': 'user', 'content': tbl.prompt}
+        ... ]
+        >>> tbl.add_computed_column(response=chat_completions(messages, model='gpt-4o-mini'))
     """
     if model_kwargs is None:
         model_kwargs = {}

pixeltable/functions/openrouter.py

@@ -0,0 +1,143 @@
+"""
+Pixeltable UDFs that wrap the OpenRouter API.
+
+OpenRouter provides a unified interface to multiple LLM providers. In order to use it,
+you must first sign up at https://openrouter.ai, create an API key, and configure it
+as described in the Working with OpenRouter tutorial.
+"""
+
+from typing import TYPE_CHECKING, Any, Optional
+
+import pixeltable as pxt
+from pixeltable.env import Env, register_client
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    import openai
+
+
+@register_client('openrouter')
+def _(api_key: str, site_url: Optional[str] = None, app_name: Optional[str] = None) -> 'openai.AsyncOpenAI':
+    import openai
+
+    # Create default headers for OpenRouter
+    default_headers: dict[str, Any] = {}
+    if site_url:
+        default_headers['HTTP-Referer'] = site_url
+    if app_name:
+        default_headers['X-Title'] = app_name
+
+    return openai.AsyncOpenAI(base_url='https://openrouter.ai/api/v1', api_key=api_key, default_headers=default_headers)
+
+
+def _openrouter_client() -> 'openai.AsyncOpenAI':
+    return Env.get().get_client('openrouter')
+
+
+@pxt.udf(resource_pool='request-rate:openrouter')
+async def chat_completions(
+    messages: list,
+    *,
+    model: str,
+    model_kwargs: Optional[dict[str, Any]] = None,
+    tools: Optional[list[dict[str, Any]]] = None,
+    tool_choice: Optional[dict[str, Any]] = None,
+    provider: Optional[dict[str, Any]] = None,
+    transforms: Optional[list[str]] = None,
+) -> dict:
+    """
+    Chat Completion API via OpenRouter.
+
+    OpenRouter provides access to multiple LLM providers through a unified API.
+    For additional details, see: <https://openrouter.ai/docs>
+
+    Supported models can be found at: <https://openrouter.ai/models>
+
+    Request throttling:
+    Applies the rate limit set in the config (section `openrouter`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
+    __Requirements:__
+
+    - `pip install openai`
+
+    Args:
+        messages: A list of messages comprising the conversation so far.
+        model: ID of the model to use (e.g., 'anthropic/claude-3.5-sonnet', 'openai/gpt-4').
+        model_kwargs: Additional OpenAI-compatible parameters.
+        tools: List of tools available to the model.
+        tool_choice: Controls which (if any) tool is called by the model.
+        provider: OpenRouter-specific provider preferences (e.g., {'order': ['Anthropic', 'OpenAI']}).
+        transforms: List of message transforms to apply (e.g., ['middle-out']).
+
+    Returns:
+        A dictionary containing the response in OpenAI format.
+
+    Examples:
+        Basic chat completion:
+
+        >>> messages = [{'role': 'user', 'content': tbl.prompt}]
+        ... tbl.add_computed_column(
+        ...     response=chat_completions(
+        ...         messages,
+        ...         model='anthropic/claude-3.5-sonnet'
+        ...     )
+        ... )
+
+        With provider routing:
+
+        >>> tbl.add_computed_column(
+        ...     response=chat_completions(
+        ...         messages,
+        ...         model='anthropic/claude-3.5-sonnet',
+        ...         provider={'require_parameters': True, 'order': ['Anthropic']}
+        ...     )
+        ... )
+
+        With transforms:
+
+        >>> tbl.add_computed_column(
+        ...     response=chat_completions(
+        ...         messages,
+        ...         model='openai/gpt-4',
+        ...         transforms=['middle-out']  # Optimize for long contexts
+        ...     )
+        ... )
+    """
+    if model_kwargs is None:
+        model_kwargs = {}
+
+    Env.get().require_package('openai')
+
+    # Handle tools if provided
+    if tools is not None:
+        model_kwargs['tools'] = [{'type': 'function', 'function': tool} for tool in tools]
+
+    if tool_choice is not None:
+        if tool_choice['auto']:
+            model_kwargs['tool_choice'] = 'auto'
+        elif tool_choice['required']:
+            model_kwargs['tool_choice'] = 'required'
+        else:
+            assert tool_choice['tool'] is not None
+            model_kwargs['tool_choice'] = {'type': 'function', 'function': {'name': tool_choice['tool']}}
+
+    # Prepare OpenRouter-specific parameters for extra_body
+    extra_body: dict[str, Any] = {}
+    if provider is not None:
+        extra_body['provider'] = provider
+    if transforms is not None:
+        extra_body['transforms'] = transforms
+
+    # Make the API call
+    result = await _openrouter_client().chat.completions.create(
+        messages=messages, model=model, extra_body=extra_body if extra_body else None, **model_kwargs
+    )
+    return result.model_dump()
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__