pixeltable 0.4.0rc3__py3-none-any.whl → 0.4.2__py3-none-any.whl

pixeltable/exprs/row_builder.py

@@ -63,6 +63,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]
     default_eval_ctx: EvalCtx
     unstored_iter_args: dict[UUID, Expr]
@@ -93,7 +94,13 @@ class RowBuilder:
         target_slot_idxs: list[int]  # slot idxs of target exprs; might contain duplicates
         target_exprs: list[Expr]  # exprs corresponding to target_slot_idxs
 
-    def __init__(self, output_exprs: Sequence[Expr], columns: Sequence[catalog.Column], input_exprs: Iterable[Expr]):
+    def __init__(
+        self,
+        output_exprs: Sequence[Expr],
+        columns: Sequence[catalog.Column],
+        input_exprs: Iterable[Expr],
+        tbl: Optional[catalog.TableVersion] = None,
+    ):
         """
         Args:
             output_exprs: list of Exprs to be evaluated
@@ -125,6 +132,7 @@ class RowBuilder:
         #   * further references to that column (eg, computed cols) need to resolve to the validating ColumnRef
         from .column_ref import ColumnRef
 
+        self.tbl = tbl
         self.table_columns: list[ColumnSlotIdx] = []
         self.input_exprs = ExprSet()
         validating_colrefs: dict[Expr, Expr] = {}  # key: non-validating colref, value: corresp. validating colref
@@ -229,6 +237,7 @@ class RowBuilder:
 
     def add_table_column(self, col: catalog.Column, slot_idx: int) -> None:
         """Record a column that is part of the table row"""
+        assert self.tbl is not None
         self.table_columns.append(ColumnSlotIdx(col, slot_idx))
 
     def output_slot_idxs(self) -> list[ColumnSlotIdx]:
@@ -427,33 +436,55 @@ class RowBuilder:
                         expr, f'expression {expr}', data_row.get_exc(expr.slot_idx), exc_tb, input_vals, 0
                     ) from exc
 
-    def create_table_row(self, data_row: DataRow, exc_col_ids: set[int]) -> tuple[dict[str, Any], int]:
+    def create_table_row(
+        self, data_row: DataRow, cols_with_excs: Optional[set[int]], pk: tuple[int, ...]
+    ) -> tuple[list[Any], int]:
         """Create a table row from the slots that have an output column assigned
 
-        Return tuple[dict that represents a stored row (can be passed to sql.insert()), # of exceptions]
+        Return tuple[list of row values in `self.table_columns` order, # of exceptions]
             This excludes system columns.
         """
         num_excs = 0
-        table_row: dict[str, Any] = {}
+        table_row: list[Any] = list(pk)
         for info in self.table_columns:
             col, slot_idx = info.col, info.slot_idx
             if data_row.has_exc(slot_idx):
-                # exceptions get stored in the errortype/-msg columns
                 exc = data_row.get_exc(slot_idx)
                 num_excs += 1
-                exc_col_ids.add(col.id)
-                table_row[col.store_name()] = None
-                table_row[col.errortype_store_name()] = type(exc).__name__
-                table_row[col.errormsg_store_name()] = str(exc)
+                if cols_with_excs is not None:
+                    cols_with_excs.add(col.id)
+                table_row.append(None)
+                if col.records_errors:
+                    # exceptions get stored in the errortype/-msg columns
+                    table_row.extend((type(exc).__name__, str(exc)))
             else:
                 if col.col_type.is_image_type() and data_row.file_urls[slot_idx] is None:
                     # we have yet to store this image
                     filepath = str(MediaStore.prepare_media_path(col.tbl.id, col.id, col.tbl.version))
                     data_row.flush_img(slot_idx, filepath)
                 val = data_row.get_stored_val(slot_idx, col.get_sa_col_type())
-                table_row[col.store_name()] = val
-                # we unfortunately need to set these, even if there are no errors
-                table_row[col.errortype_store_name()] = None
-                table_row[col.errormsg_store_name()] = None
+                table_row.append(val)
+                if col.records_errors:
+                    table_row.extend((None, None))
 
         return table_row, num_excs
+
+    def store_column_names(self) -> tuple[list[str], dict[int, catalog.Column]]:
+        """
+        Returns the list of store column names corresponding to the table_columns of this RowBuilder.
+        The second tuple element of the return value is a dictionary containing all media columns in the
+        table; it's the mapping {list_index: column}.
+        """
+        assert self.tbl is not None, self.table_columns
+        store_col_names: list[str] = [pk_col.name for pk_col in self.tbl.store_tbl.pk_columns()]
+        media_cols: dict[int, catalog.Column] = {}
+
+        for col in self.table_columns:
+            if col.col.col_type.is_media_type():
+                media_cols[len(store_col_names)] = col.col
+            store_col_names.append(col.col.store_name())
+            if col.col.records_errors:
+                store_col_names.append(col.col.errortype_store_name())
+                store_col_names.append(col.col.errormsg_store_name())
+
+        return store_col_names, media_cols

pixeltable/func/__init__.py

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

pixeltable/func/mcp.py

@@ -0,0 +1,74 @@
+import asyncio
+import inspect
+from typing import TYPE_CHECKING, Any, Optional
+
+import pixeltable as pxt
+from pixeltable import exceptions as excs, type_system as ts
+from pixeltable.func.signature import Parameter
+
+if TYPE_CHECKING:
+    import mcp
+
+
+def mcp_udfs(url: str) -> list['pxt.func.Function']:
+    return asyncio.run(mcp_udfs_async(url))
+
+
+async def mcp_udfs_async(url: str) -> list['pxt.func.Function']:
+    import mcp
+    from mcp.client.streamable_http import streamablehttp_client
+
+    list_tools_result: Optional[mcp.types.ListToolsResult] = None
+    async with (
+        streamablehttp_client(url) as (read_stream, write_stream, _),
+        mcp.ClientSession(read_stream, write_stream) as session,
+    ):
+        await session.initialize()
+        list_tools_result = await session.list_tools()
+    assert list_tools_result is not None
+
+    return [mcp_tool_to_udf(url, tool) for tool in list_tools_result.tools]
+
+
+def mcp_tool_to_udf(url: str, mcp_tool: 'mcp.types.Tool') -> 'pxt.func.Function':
+    import mcp
+    from mcp.client.streamable_http import streamablehttp_client
+
+    async def invoke(**kwargs: Any) -> str:
+        # TODO: Cache session objects rather than creating a new one each time?
+        async with (
+            streamablehttp_client(url) as (read_stream, write_stream, _),
+            mcp.ClientSession(read_stream, write_stream) as session,
+        ):
+            await session.initialize()
+            res = await session.call_tool(name=mcp_tool.name, arguments=kwargs)
+            # TODO Handle image/audio responses?
+            return res.content[0].text  # type: ignore[union-attr]
+
+    if mcp_tool.description is not None:
+        invoke.__doc__ = mcp_tool.description
+
+    input_schema = mcp_tool.inputSchema
+    params = {
+        name: __mcp_param_to_pxt_type(mcp_tool.name, name, param) for name, param in input_schema['properties'].items()
+    }
+    required = input_schema.get('required', [])
+
+    # Ensure that any params not appearing in `required` are nullable.
+    # (A required param might or might not be nullable, since its type might be an 'anyOf' containing a null.)
+    for name in params.keys() - required:
+        params[name] = params[name].copy(nullable=True)
+
+    signature = pxt.func.Signature(
+        return_type=ts.StringType(),  # Return type is always string
+        parameters=[Parameter(name, col_type, inspect.Parameter.KEYWORD_ONLY) for name, col_type in params.items()],
+    )
+
+    return pxt.func.CallableFunction(signatures=[signature], py_fns=[invoke], self_name=mcp_tool.name)
+
+
+def __mcp_param_to_pxt_type(tool_name: str, name: str, param: dict[str, Any]) -> ts.ColumnType:
+    pxt_type = ts.ColumnType.from_json_schema(param)
+    if pxt_type is None:
+        raise excs.Error(f'Unknown type schema for MCP parameter {name!r} of tool {tool_name!r}: {param}')
+    return pxt_type

pixeltable/func/query_template_function.py

@@ -157,11 +157,13 @@ def retrieval_udf(
     """
     # Argument validation
     col_refs: list[exprs.ColumnRef]
+    # TODO: get rid of references to ColumnRef internals and replace instead with a public interface
+    col_names = table.columns()
     if parameters is None:
-        col_refs = [table[col_name] for col_name in table.columns if not table[col_name].col.is_computed]
+        col_refs = [table[col_name] for col_name in col_names if not table[col_name].col.is_computed]
     else:
         for param in parameters:
-            if isinstance(param, str) and param not in table.columns:
+            if isinstance(param, str) and param not in col_names:
                 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]
 

pixeltable/func/tools.py

@@ -1,8 +1,9 @@
+import json
 from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Union
 
 import pydantic
 
-from pixeltable import exceptions as excs
+from pixeltable import exceptions as excs, type_system as ts
 
 from .function import Function
 from .signature import Parameter
@@ -69,7 +70,9 @@ class Tool(pydantic.BaseModel):
             return _extract_float_tool_arg(kwargs, param_name=param.name)
         if param.col_type.is_bool_type():
             return _extract_bool_tool_arg(kwargs, param_name=param.name)
-        raise AssertionError()
+        if param.col_type.is_json_type():
+            return _extract_json_tool_arg(kwargs, param_name=param.name)
+        raise AssertionError(param.col_type)
 
 
 class ToolChoice(pydantic.BaseModel):
@@ -137,6 +140,13 @@ def _extract_bool_tool_arg(kwargs: dict[str, Any], param_name: str) -> Optional[
     return _extract_arg(bool, kwargs, param_name)
 
 
+@udf
+def _extract_json_tool_arg(kwargs: dict[str, Any], param_name: str) -> Optional[ts.Json]:
+    if param_name in kwargs:
+        return json.loads(kwargs[param_name])
+    return None
+
+
 T = TypeVar('T')
 
 

pixeltable/func/udf.py

@@ -262,7 +262,7 @@ def from_table(
     """
     from pixeltable import exprs
 
-    ancestors = [tbl, *tbl._base_tables]
+    ancestors = [tbl, *tbl._get_base_tables()]
     ancestors.reverse()  # We must traverse the ancestors in order from base to derived
 
     subst: dict[exprs.Expr, exprs.Expr] = {}
@@ -297,7 +297,7 @@ def from_table(
 
     if description is None:
         # Default description is the table comment
-        description = tbl._comment
+        description = tbl._get_comment()
         if len(description) == 0:
             description = f"UDF for table '{tbl._name}'"
 

pixeltable/functions/__init__.py

@@ -10,6 +10,7 @@ from . import (
     deepseek,
     fireworks,
     gemini,
+    groq,
     huggingface,
     image,
     json,

pixeltable/functions/groq.py

@@ -0,0 +1,108 @@
+"""
+Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
+that wrap various endpoints from the Groq API. In order to use them, you must
+first `pip install groq` and configure your Groq credentials, as described in
+the [Working with Groq](https://pixeltable.readme.io/docs/working-with-groq) tutorial.
+"""
+
+from typing import TYPE_CHECKING, Any, Optional
+
+import pixeltable as pxt
+from pixeltable import exprs
+from pixeltable.env import Env, register_client
+from pixeltable.utils.code import local_public_names
+
+from .openai import _openai_response_to_pxt_tool_calls
+
+if TYPE_CHECKING:
+    import groq
+
+
+@register_client('groq')
+def _(api_key: str) -> 'groq.AsyncGroq':
+    import groq
+
+    return groq.AsyncGroq(api_key=api_key)
+
+
+def _groq_client() -> 'groq.AsyncGroq':
+    return Env.get().get_client('groq')
+
+
+@pxt.udf(resource_pool='request-rate:groq')
+async def chat_completions(
+    messages: list[dict[str, str]],
+    *,
+    model: str,
+    model_kwargs: Optional[dict[str, Any]] = None,
+    tools: Optional[list[dict[str, Any]]] = None,
+    tool_choice: Optional[dict[str, Any]] = None,
+) -> dict:
+    """
+    Chat Completion API.
+
+    Equivalent to the Groq `chat/completions` API endpoint.
+    For additional details, see: <https://console.groq.com/docs/api-reference#chat-create>
+
+    Request throttling:
+    Applies the rate limit set in the config (section `groq`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
+    __Requirements:__
+
+    - `pip install groq`
+
+    Args:
+        messages: A list of messages comprising the conversation so far.
+        model: ID of the model to use. (See overview here: <https://console.groq.com/docs/models>)
+        model_kwargs: Additional keyword args for the Groq `chat/completions` API.
+            For details on the available parameters, see: <https://console.groq.com/docs/api-reference#chat-create>
+
+    Returns:
+        A dictionary containing the response and other metadata.
+
+    Examples:
+        Add a computed column that applies the model `llama3-8b-8192`
+        to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
+
+        >>> messages = [{'role': 'user', 'content': tbl.prompt}]
+        ... tbl.add_computed_column(response=chat_completions(messages, model='llama3-8b-8192'))
+    """
+    if model_kwargs is None:
+        model_kwargs = {}
+
+    Env.get().require_package('groq')
+
+    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']}}
+
+    if tool_choice is not None and not tool_choice['parallel_tool_calls']:
+        model_kwargs['parallel_tool_calls'] = False
+
+    result = await _groq_client().chat.completions.create(
+        messages=messages,  # type: ignore[arg-type]
+        model=model,
+        **model_kwargs,
+    )
+    return result.model_dump()
+
+
+def invoke_tools(tools: pxt.func.Tools, response: exprs.Expr) -> exprs.InlineDict:
+    """Converts an OpenAI response dict to Pixeltable tool invocation format and calls `tools._invoke()`."""
+    return tools._invoke(_openai_response_to_pxt_tool_calls(response))
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/functions/huggingface.py

@@ -51,7 +51,7 @@ def sentence_transformer(
     """
     env.Env.get().require_package('sentence_transformers')
     device = resolve_torch_device('auto')
-    from sentence_transformers import SentenceTransformer  # type: ignore
+    from sentence_transformers import SentenceTransformer
 
     # specifying the device, moves the model to device (gpu:cuda/mps, cpu)
     model = _lookup_model(model_id, SentenceTransformer, device=device, pass_device_to_create=True)
@@ -170,7 +170,7 @@ def clip(text: Batch[str], *, model_id: str) -> Batch[pxt.Array[(None,), pxt.Flo
     env.Env.get().require_package('transformers')
     device = resolve_torch_device('auto')
     import torch
-    from transformers import CLIPModel, CLIPProcessor  # type: ignore
+    from transformers import CLIPModel, CLIPProcessor
 
     model = _lookup_model(model_id, CLIPModel.from_pretrained, device=device)
     processor = _lookup_processor(model_id, CLIPProcessor.from_pretrained)
@@ -395,19 +395,21 @@ def speech2text_for_conditional_generation(audio: pxt.Audio, *, model_id: str, l
     device = resolve_torch_device('auto', allow_mps=False)  # Doesn't seem to work on 'mps'; use 'cpu' instead
     import torch
     import torchaudio  # type: ignore[import-untyped]
-    from transformers import Speech2TextForConditionalGeneration, Speech2TextProcessor
+    from transformers import Speech2TextForConditionalGeneration, Speech2TextProcessor, Speech2TextTokenizer
 
     model = _lookup_model(model_id, Speech2TextForConditionalGeneration.from_pretrained, device=device)
     processor = _lookup_processor(model_id, Speech2TextProcessor.from_pretrained)
+    tokenizer = processor.tokenizer
     assert isinstance(processor, Speech2TextProcessor)
+    assert isinstance(tokenizer, Speech2TextTokenizer)
 
-    if language is not None and language not in processor.tokenizer.lang_code_to_id:
+    if language is not None and language not in tokenizer.lang_code_to_id:
         raise excs.Error(
             f"Language code '{language}' is not supported by the model '{model_id}'. "
-            f'Supported languages are: {list(processor.tokenizer.lang_code_to_id.keys())}'
+            f'Supported languages are: {list(tokenizer.lang_code_to_id.keys())}'
         )
 
-    forced_bos_token_id: Optional[int] = None if language is None else processor.tokenizer.lang_code_to_id[language]
+    forced_bos_token_id: Optional[int] = None if language is None else tokenizer.lang_code_to_id[language]
 
     # Get the model's sampling rate. Default to 16 kHz (the standard) if not in config
     model_sampling_rate = getattr(model.config, 'sampling_rate', 16_000)

pixeltable/functions/mistralai.py

@@ -5,7 +5,7 @@ first `pip install mistralai` and configure your Mistral AI credentials, as desc
 the [Working with Mistral AI](https://pixeltable.readme.io/docs/working-with-mistralai) tutorial.
 """
 
-from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Optional
 
 import numpy as np
 
@@ -16,7 +16,7 @@ from pixeltable.func.signature import Batch
 from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
-    import mistralai.types.basemodel
+    import mistralai
 
 
 @register_client('mistral')
@@ -54,8 +54,6 @@ async def chat_completions(
         model_kwargs: Additional keyword args for the Mistral `chat/completions` API.
             For details on the available parameters, see: <https://docs.mistral.ai/api/#tag/chat>
 
-    For details on the other parameters, see: <https://docs.mistral.ai/api/#tag/chat>
-
     Returns:
         A dictionary containing the response and other metadata.
 
@@ -156,15 +154,6 @@ def _(model: str) -> ts.ArrayType:
     return ts.ArrayType((dimensions,), dtype=ts.FloatType())
 
 
-_T = TypeVar('_T')
-
-
-def _opt(arg: Optional[_T]) -> Union[_T, 'mistralai.types.basemodel.Unset']:
-    from mistralai.types import UNSET
-
-    return arg if arg is not None else UNSET
-
-
 __all__ = local_public_names(__name__)
 
 

pixeltable/functions/openai.py

@@ -205,12 +205,7 @@ async def speech(input: str, *, model: str, voice: str, model_kwargs: Optional[d
     if model_kwargs is None:
         model_kwargs = {}
 
-    content = await _openai_client().audio.speech.create(
-        input=input,
-        model=model,
-        voice=voice,  # type: ignore
-        **model_kwargs,
-    )
+    content = await _openai_client().audio.speech.create(input=input, model=model, voice=voice, **model_kwargs)
     ext = model_kwargs.get('response_format', 'mp3')
     output_filename = str(env.Env.get().tmp_dir / f'{uuid.uuid4()}.{ext}')
     content.write_to_file(output_filename)

pixeltable/functions/replicate.py

@@ -12,7 +12,7 @@ from pixeltable.env import Env, register_client
 from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
-    import replicate  # type: ignore[import-untyped]
+    import replicate
 
 
 @register_client('replicate')
@@ -27,7 +27,7 @@ def _replicate_client() -> 'replicate.Client':
 
 
 @pxt.udf(resource_pool='request-rate:replicate')
-async def run(input: dict[str, Any], *, ref: str) -> dict[str, Any]:
+async def run(input: dict[str, Any], *, ref: str) -> pxt.Json:
     """
     Run a model on Replicate.
 

pixeltable/functions/util.py

@@ -1,5 +1,6 @@
 import PIL.Image
 
+from pixeltable.config import Config
 from pixeltable.env import Env
 
 
@@ -7,10 +8,14 @@ def resolve_torch_device(device: str, allow_mps: bool = True) -> str:
     Env.get().require_package('torch')
     import torch
 
+    mps_enabled = Config.get().get_bool_value('enable_mps')
+    if mps_enabled is None:
+        mps_enabled = True  # Default to True if not set in config
+
     if device == 'auto':
         if torch.cuda.is_available():
             return 'cuda'
-        if allow_mps and torch.backends.mps.is_available():
+        if mps_enabled and allow_mps and torch.backends.mps.is_available():
             return 'mps'
         return 'cpu'
     return device

pixeltable/globals.py

@@ -428,8 +428,6 @@ def get_table(path: str) -> catalog.Table:
     """
     path_obj = catalog.Path(path)
     tbl = Catalog.get().get_table(path_obj)
-    tv = tbl._tbl_version.get()
-    _logger.debug(f'get_table(): tbl={tv.id}:{tv.effective_version} sa_tbl={id(tv.store_tbl.sa_tbl):x} tv={id(tv):x}')
     return tbl