pixeltable 0.2.29__py3-none-any.whl → 0.2.30__py3-none-any.whl

pixeltable/__init__.py

@@ -4,7 +4,7 @@ from .exceptions import Error
 from .exprs import RELATIVE_PATH_ROOT
 from .func import Aggregator, Function, expr_udf, uda, udf
 from .globals import (array, configure_logging, create_dir, create_snapshot, create_table, create_view, drop_dir,
-                      drop_table, get_table, init, list_dirs, list_functions, list_tables, move)
+                      drop_table, get_table, init, list_dirs, list_functions, list_tables, move, 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)

pixeltable/__version__.py

@@ -1,3 +1,3 @@
 # These version placeholders will be replaced during build.
-__version__ = "0.2.29"
-__version_tuple__ = (0, 2, 29)
+__version__ = "0.2.30"
+__version_tuple__ = (0, 2, 30)

pixeltable/catalog/__init__.py

@@ -1,7 +1,7 @@
 from .catalog import Catalog
 from .column import Column
 from .dir import Dir
-from .globals import UpdateStatus, is_valid_identifier, is_valid_path, MediaValidation, IfExistsParam
+from .globals import UpdateStatus, is_valid_identifier, is_valid_path, MediaValidation, IfExistsParam, IfNotExistsParam
 from .insertable_table import InsertableTable
 from .named_function import NamedFunction
 from .path import Path

pixeltable/catalog/globals.py

@@ -65,6 +65,18 @@ class IfExistsParam(enum.Enum):
             val_strs = ', '.join(f'{s.lower()!r}' for s in cls.__members__.keys())
             raise excs.Error(f'{param_name} must be one of: [{val_strs}]')
 
+class IfNotExistsParam(enum.Enum):
+    ERROR = 0
+    IGNORE = 1
+
+    @classmethod
+    def validated(cls, param_val: str, param_name: str) -> IfNotExistsParam:
+        try:
+            return cls[param_val.upper()]
+        except KeyError:
+            val_strs = ', '.join(f'{s.lower()!r}' for s in cls.__members__.keys())
+            raise excs.Error(f'{param_name} must be one of: [{val_strs}]')
+
 def is_valid_identifier(name: str) -> bool:
     return name.isidentifier() and not name.startswith('_')
 

pixeltable/catalog/table.py

@@ -25,7 +25,7 @@ from ..exprs import ColumnRef
 from ..utils.description_helper import DescriptionHelper
 from ..utils.filecache import FileCache
 from .column import Column
-from .globals import _ROWID_COLUMN_NAME, MediaValidation, UpdateStatus, is_system_column_name, is_valid_identifier
+from .globals import _ROWID_COLUMN_NAME, MediaValidation, UpdateStatus, is_system_column_name, is_valid_identifier, IfNotExistsParam
 from .schema_object import SchemaObject
 from .table_version import TableVersion
 from .table_version_path import TableVersionPath
@@ -712,14 +712,19 @@ class Table(SchemaObject):
         if not exists:
             raise excs.Error(f'Unknown column: {col_ref.col.qualified_name}')
 
-    def drop_column(self, column: Union[str, ColumnRef]) -> None:
+    def drop_column(self, column: Union[str, ColumnRef], if_not_exists: Literal['error', 'ignore'] = 'error') -> None:
         """Drop a column from the table.
 
         Args:
             column: The name or reference of the column to drop.
+            if_not_exists: Directive for handling a non-existent column. Must be one of the following:
+
+                - `'error'`: raise an error if the column does not exist.
+                - `'ignore'`: do nothing if the column does not exist.
 
         Raises:
-            Error: If the column does not exist or if it is referenced by a dependent computed column.
+            Error: If the column does not exist and `if_exists='error'`,
+                or if it is referenced by a dependent computed column.
 
         Examples:
             Drop the column `col` from the table `my_table` by column name:
@@ -731,14 +736,32 @@ class Table(SchemaObject):
 
             >>> tbl = pxt.get_table('my_table')
             ... tbl.drop_column(tbl.col)
+
+            Drop the column `col` from the table `my_table` if it exists, otherwise do nothing:
+
+            >>> tbl = pxt.get_table('my_table')
+            ... tbl.drop_col(tbl.col, if_not_exists='ignore')
         """
         self._check_is_dropped()
+        if self._tbl_version_path.is_snapshot():
+            raise excs.Error('Cannot drop column from a snapshot.')
         col: Column = None
+        _if_not_exists = IfNotExistsParam.validated(if_not_exists, 'if_not_exists')
         if isinstance(column, str):
-            self.__check_column_name_exists(column)
+            col = self._tbl_version_path.get_column(column, include_bases=False)
+            if col is None:
+                if _if_not_exists == IfNotExistsParam.ERROR:
+                    raise excs.Error(f'Column {column!r} unknown')
+                assert _if_not_exists == IfNotExistsParam.IGNORE
+                return
             col = self._tbl_version.cols_by_name[column]
         else:
-            self.__check_column_ref_exists(column)
+            exists = self._tbl_version_path.has_column(column.col, include_bases=False)
+            if not exists:
+                if _if_not_exists == IfNotExistsParam.ERROR:
+                    raise excs.Error(f'Unknown column: {column.col.qualified_name}')
+                assert _if_not_exists == IfNotExistsParam.IGNORE
+                return
             col = column.col
 
         dependent_user_cols = [c for c in col.dependent_cols if c.name is not None]
@@ -866,7 +889,9 @@ class Table(SchemaObject):
     def drop_embedding_index(
             self, *,
             column: Union[str, ColumnRef, None] = None,
-            idx_name: Optional[str] = None) -> None:
+            idx_name: Optional[str] = None,
+            if_not_exists: Literal['error', 'ignore'] = 'error'
+    ) -> None:
         """
         Drop an embedding index from the table. Either a column name or an index name (but not both) must be
         specified. If a column name or reference is specified, it must be a column containing exactly one
@@ -876,11 +901,20 @@ class Table(SchemaObject):
             column: The name of, or reference to, the column from which to drop the index.
                     The column must have only one embedding index.
             idx_name: The name of the index to drop.
+            if_not_exists: Directive for handling a non-existent index. Must be one of the following:
+
+                - `'error'`: raise an error if the index does not exist.
+                - `'ignore'`: do nothing if the index does not exist.
+
+                Note that `if_not_exists` parameter is only applicable when an `idx_name` is specified
+                and it does not exist, or when `column` is specified and it has no index.
+                `if_not_exists` does not apply to non-exisitng column.
 
         Raises:
             Error: If `column` is specified, but the column does not exist, or it contains no embedding
-                indices or multiple embedding indices.
-            Error: If `idx_name` is specified, but the index does not exist or is not an embedding index.
+                indices and `if_not_exists='error'`, or the column has multiple embedding indices.
+            Error: If `idx_name` is specified, but the index is not an embedding index, or
+                the index does not exist and `if_not_exists='error'`.
 
         Examples:
             Drop the embedding index on the `img` column of the table `my_table` by column name:
@@ -897,6 +931,9 @@ class Table(SchemaObject):
             >>> tbl = pxt.get_table('my_table')
             ... tbl.drop_embedding_index(idx_name='idx1')
 
+            Drop the embedding index `idx1` of the table `my_table` by index name, if it exists, otherwise do nothing:
+            >>> tbl = pxt.get_table('my_table')
+            ... tbl.drop_embedding_index(idx_name='idx1', if_not_exists='ignore')
         """
         if (column is None) == (idx_name is None):
             raise excs.Error("Exactly one of 'column' or 'idx_name' must be provided")
@@ -910,12 +947,14 @@ class Table(SchemaObject):
                 self.__check_column_ref_exists(column, include_bases=True)
                 col = column.col
             assert col is not None
-        self._drop_index(col=col, idx_name=idx_name, _idx_class=index.EmbeddingIndex)
+        self._drop_index(col=col, idx_name=idx_name, _idx_class=index.EmbeddingIndex, if_not_exists=if_not_exists)
 
     def drop_index(
             self, *,
             column: Union[str, ColumnRef, None] = None,
-            idx_name: Optional[str] = None) -> None:
+            idx_name: Optional[str] = None,
+            if_not_exists: Literal['error', 'ignore'] = 'error'
+    ) -> None:
         """
         Drop an index from the table. Either a column name or an index name (but not both) must be
         specified. If a column name or reference is specified, it must be a column containing exactly one index;
@@ -925,6 +964,14 @@ class Table(SchemaObject):
             column: The name of, or reference to, the column from which to drop the index.
                     The column must have only one embedding index.
             idx_name: The name of the index to drop.
+            if_not_exists: Directive for handling a non-existent index. Must be one of the following:
+
+                - `'error'`: raise an error if the index does not exist.
+                - `'ignore'`: do nothing if the index does not exist.
+
+                Note that `if_not_exists` parameter is only applicable when an `idx_name` is specified
+                and it does not exist, or when `column` is specified and it has no index.
+                `if_not_exists` does not apply to non-exisitng column.
 
         Raises:
             Error: If `column` is specified, but the column does not exist, or it contains no
@@ -946,6 +993,10 @@ class Table(SchemaObject):
             >>> tbl = pxt.get_table('my_table')
             ... tbl.drop_index(idx_name='idx1')
 
+            Drop the index `idx1` of the table `my_table` by index name, if it exists, otherwise do nothing:
+            >>> tbl = pxt.get_table('my_table')
+            ... tbl.drop_index(idx_name='idx1', if_not_exists='ignore')
+
         """
         if (column is None) == (idx_name is None):
             raise excs.Error("Exactly one of 'column' or 'idx_name' must be provided")
@@ -959,20 +1010,25 @@ class Table(SchemaObject):
                 self.__check_column_ref_exists(column, include_bases=True)
                 col = column.col
             assert col is not None
-        self._drop_index(col=col, idx_name=idx_name)
+        self._drop_index(col=col, idx_name=idx_name, if_not_exists=if_not_exists)
 
     def _drop_index(
             self, *, col: Optional[Column] = None,
             idx_name: Optional[str] = None,
-            _idx_class: Optional[type[index.IndexBase]] = None
+            _idx_class: Optional[type[index.IndexBase]] = None,
+            if_not_exists: Literal['error', 'ignore'] = 'error'
     ) -> None:
         if self._tbl_version_path.is_snapshot():
             raise excs.Error('Cannot drop an index from a snapshot')
         assert (col is None) != (idx_name is None)
 
         if idx_name is not None:
+            _if_not_exists = IfNotExistsParam.validated(if_not_exists, 'if_not_exists')
             if idx_name not in self._tbl_version.idxs_by_name:
-                raise excs.Error(f'Index {idx_name!r} does not exist')
+                if _if_not_exists == IfNotExistsParam.ERROR:
+                    raise excs.Error(f'Index {idx_name!r} does not exist')
+                assert _if_not_exists == IfNotExistsParam.IGNORE
+                return
             idx_id = self._tbl_version.idxs_by_name[idx_name].id
         else:
             if col.tbl.id != self._tbl_version.id:
@@ -982,7 +1038,11 @@ class Table(SchemaObject):
             if _idx_class is not None:
                 idx_info = [info for info in idx_info if isinstance(info.idx, _idx_class)]
             if len(idx_info) == 0:
-                raise excs.Error(f'Column {col.name!r} does not have an index')
+                _if_not_exists = IfNotExistsParam.validated(if_not_exists, 'if_not_exists')
+                if _if_not_exists == IfNotExistsParam.ERROR:
+                    raise excs.Error(f'Column {col.name!r} does not have an index')
+                assert _if_not_exists == IfNotExistsParam.IGNORE
+                return
             if len(idx_info) > 1:
                 raise excs.Error(f"Column {col.name!r} has multiple indices; specify 'idx_name' instead")
             idx_id = idx_info[0].id

pixeltable/exprs/function_call.py

@@ -199,11 +199,6 @@ class FunctionCall(Expr):
                     pass
 
             if not isinstance(arg, Expr):
-                # make sure that non-Expr args are json-serializable and are literals of the correct type
-                try:
-                    _ = json.dumps(arg)
-                except TypeError:
-                    raise excs.Error(f'Argument for parameter {param_name!r} is not json-serializable: {arg} (of type {type(arg)})')
                 if arg is not None:
                     try:
                         param_type = param.col_type

pixeltable/func/__init__.py

@@ -5,4 +5,5 @@ from .function import Function
 from .function_registry import FunctionRegistry
 from .query_template_function import QueryTemplateFunction
 from .signature import Signature, Parameter, Batch
+from .tools import Tool, Tools
 from .udf import udf, make_function, expr_udf

pixeltable/func/aggregate_function.py

@@ -139,6 +139,9 @@ class AggregateFunction(Function):
         self.init_param_names.append(init_param_names)
         return self
 
+    def _docstring(self) -> Optional[str]:
+        return inspect.getdoc(self.agg_classes[0])
+
     def help_str(self) -> str:
         res = super().help_str()
         # We need to reference agg_classes[0] rather than agg_class here, because we want this to work even if the

pixeltable/func/callable_function.py

@@ -48,6 +48,9 @@ class CallableFunction(Function):
     def is_batched(self) -> bool:
         return self.batch_size is not None
 
+    def _docstring(self) -> Optional[str]:
+        return inspect.getdoc(self.py_fns[0])
+
     @property
     def py_fn(self) -> Callable:
         assert not self.is_polymorphic
@@ -109,11 +112,6 @@ class CallableFunction(Function):
         self.py_fns.append(fn)
         return self
 
-    def help_str(self) -> str:
-        res = super().help_str()
-        res += '\n\n' + inspect.getdoc(self.py_fns[0])
-        return res
-
     def _as_dict(self) -> dict:
         if self.self_path is None:
             # this is not a module function

pixeltable/func/expr_template_function.py

@@ -87,6 +87,13 @@ class ExprTemplateFunction(Function):
         assert not result._contains(exprs.Variable)
         return result
 
+    def _docstring(self) -> Optional[str]:
+        from pixeltable import exprs
+
+        if isinstance(self.templates[0].expr, exprs.FunctionCall):
+            return self.templates[0].expr.fn._docstring()
+        return None
+
     def exec(self, args: Sequence[Any], kwargs: dict[str, Any]) -> Any:
         from pixeltable import exec, exprs
 

pixeltable/func/function.py

@@ -92,6 +92,16 @@ class Function(abc.ABC):
         assert not self.is_polymorphic
         return len(self.signature.parameters)
 
+    def _docstring(self) -> Optional[str]:
+        return None
+
+    def help_str(self) -> str:
+        docstring = self._docstring()
+        display = self.display_name + str(self.signatures[0])
+        if docstring is None:
+            return display
+        return f'{display}\n\n{docstring}'
+
     @property
     def _resolved_fns(self) -> list[Self]:
         """
@@ -129,9 +139,6 @@ class Function(abc.ABC):
         """
         raise NotImplementedError()
 
-    def help_str(self) -> str:
-        return self.display_name + str(self.signatures[0])
-
     def __call__(self, *args: Any, **kwargs: Any) -> 'pxt.exprs.FunctionCall':
         from pixeltable import exprs
 
@@ -284,7 +291,7 @@ class Function(abc.ABC):
         """Print source code"""
         print('source not available')
 
-    def as_dict(self) -> dict:
+    def as_dict(self) -> dict[str, Any]:
         """
         Return a serialized reference to the instance that can be passed to json.dumps() and converted back
         to an instance with from_dict().

pixeltable/func/tools.py

@@ -0,0 +1,116 @@
+from dataclasses import dataclass
+import dataclasses
+import json
+from typing import TYPE_CHECKING, Any, Optional
+
+import pydantic
+
+from .function import Function
+from .signature import Parameter
+from .udf import udf
+
+if TYPE_CHECKING:
+    from pixeltable import exprs
+
+
+# The Tool and Tools classes are containers that hold Pixeltable UDFs and related metadata, so that they can be
+# realized as LLM tools. They are implemented as Pydantic models in order to provide a canonical way of converting
+# to JSON, via the Pydantic `model_serializer` interface. In this way, they can be passed directly as UDF
+# parameters as described in the `pixeltable.tools` and `pixeltable.tool` docstrings.
+#
+# (The dataclass dict serializer is insufficiently flexible for this purpose: `Tool` contains a member of type
+# `Function`, which is not natively JSON-serializable; Pydantic provides a way of customizing its default
+# serialization behavior, whereas dataclasses do not.)
+
+class Tool(pydantic.BaseModel):
+    # Allow arbitrary types so that we can include a Pixeltable function in the schema.
+    # We will implement a model_serializer to ensure the Tool model can be serialized.
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
+    fn: Function
+    name: Optional[str] = None
+    description: Optional[str] = None
+
+    @property
+    def parameters(self) -> dict[str, Parameter]:
+        return self.fn.signature.parameters
+
+    @pydantic.model_serializer
+    def ser_model(self) -> dict[str, Any]:
+        return {
+            'name': self.name or self.fn.name,
+            'description': self.description or self.fn._docstring(),
+            'parameters': {
+                'type': 'object',
+                'properties': {
+                    param.name: param.col_type._to_json_schema()
+                    for param in self.parameters.values()
+                }
+            },
+            'required': [
+                param.name for param in self.parameters.values() if not param.col_type.nullable
+            ],
+            'additionalProperties': False,  # TODO Handle kwargs?
+        }
+
+    # `tool_calls` must be in standardized tool invocation format:
+    # {tool_name: {'args': {name1: value1, name2: value2, ...}}, ...}
+    def invoke(self, tool_calls: 'exprs.Expr') -> 'exprs.FunctionCall':
+        kwargs = {
+            param.name: self.__extract_tool_arg(param, tool_calls)
+            for param in self.parameters.values()
+        }
+        return self.fn(**kwargs)
+
+    def __extract_tool_arg(self, param: Parameter, tool_calls: 'exprs.Expr') -> 'exprs.Expr':
+        func_name = self.name or self.fn.name
+        if param.col_type.is_string_type():
+            return _extract_str_tool_arg(tool_calls, func_name=func_name, param_name=param.name)
+        if param.col_type.is_int_type():
+            return _extract_int_tool_arg(tool_calls, func_name=func_name, param_name=param.name)
+        if param.col_type.is_float_type():
+            return _extract_float_tool_arg(tool_calls, func_name=func_name, param_name=param.name)
+        if param.col_type.is_bool_type():
+            return _extract_bool_tool_arg(tool_calls, func_name=func_name, param_name=param.name)
+        assert False
+
+
+class Tools(pydantic.BaseModel):
+    tools: list[Tool]
+
+    @pydantic.model_serializer
+    def ser_model(self) -> list[dict[str, Any]]:
+        return [tool.ser_model() for tool in self.tools]
+
+    # `tool_calls` must be in standardized tool invocation format:
+    # {tool_name: {'args': {name1: value1, name2: value2, ...}}, ...}
+    def _invoke(self, tool_calls: 'exprs.Expr') -> 'exprs.InlineDict':
+        from pixeltable import exprs
+
+        return exprs.InlineDict({
+            tool.name or tool.fn.name: tool.invoke(tool_calls)
+            for tool in self.tools
+        })
+
+
+@udf
+def _extract_str_tool_arg(tool_calls: dict, func_name: str, param_name: str) -> Optional[str]:
+    return str(_extract_arg(tool_calls, func_name, param_name))
+
+@udf
+def _extract_int_tool_arg(tool_calls: dict, func_name: str, param_name: str) -> Optional[int]:
+    return int(_extract_arg(tool_calls, func_name, param_name))
+
+@udf
+def _extract_float_tool_arg(tool_calls: dict, func_name: str, param_name: str) -> Optional[float]:
+    return float(_extract_arg(tool_calls, func_name, param_name))
+
+@udf
+def _extract_bool_tool_arg(tool_calls: dict, func_name: str, param_name: str) -> Optional[bool]:
+    return bool(_extract_arg(tool_calls, func_name, param_name))
+
+def _extract_arg(tool_calls: dict, func_name: str, param_name: str) -> Any:
+    if func_name in tool_calls:
+        arguments = tool_calls[func_name]['args']
+        return arguments.get(param_name)
+    return None

pixeltable/functions/__init__.py

@@ -1,7 +1,7 @@
 from pixeltable.utils.code import local_public_names
 
-from . import (anthropic, audio, fireworks, gemini, huggingface, image, json, llama_cpp, mistralai, ollama, openai,
-               string, timestamp, together, video, vision, whisper)
+from . import (anthropic, audio, fireworks, gemini, huggingface, image, json, llama_cpp, math, mistralai, ollama,
+               openai, string, timestamp, together, video, vision, whisper)
 from .globals import *
 
 __all__ = local_public_names(__name__, exclude=['globals']) + local_public_names(globals.__name__)

pixeltable/functions/anthropic.py

@@ -10,7 +10,8 @@ from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Union
 import tenacity
 
 import pixeltable as pxt
-from pixeltable import env
+from pixeltable import env, exprs
+from pixeltable.func import Tools
 from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
@@ -47,7 +48,7 @@ def messages(
     system: Optional[str] = None,
     temperature: Optional[float] = None,
     tool_choice: Optional[list[dict]] = None,
-    tools: Optional[dict] = None,
+    tools: Optional[list[dict]] = None,
     top_k: Optional[int] = None,
     top_p: Optional[float] = None,
 ) -> dict:
@@ -77,6 +78,21 @@ def messages(
         >>> msgs = [{'role': 'user', 'content': tbl.prompt}]
         ... tbl['response'] = messages(msgs, model='claude-3-haiku-20240307')
     """
+    if tools is not None:
+        # Reformat `tools` into Anthropic format
+        tools = [
+            {
+                'name': tool['name'],
+                'description': tool['description'],
+                'input_schema': {
+                    'type': 'object',
+                    'properties': tool['parameters']['properties'],
+                    'required': tool['required'],
+                },
+            }
+            for tool in tools
+        ]
+
     return _retry(_anthropic_client().messages.create)(
         messages=messages,
         model=model,
@@ -92,6 +108,24 @@ def messages(
     ).dict()
 
 
+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(_anthropic_response_to_pxt_tool_calls(response))
+
+
+@pxt.udf
+def _anthropic_response_to_pxt_tool_calls(response: dict) -> Optional[dict]:
+    anthropic_tool_calls = [r for r in response['content'] if r['type'] == 'tool_use']
+    if len(anthropic_tool_calls) > 0:
+        return {
+            tool_call['name']: {
+                'args': tool_call['input']
+            }
+            for tool_call in anthropic_tool_calls
+        }
+    return None
+
+
 _T = TypeVar('_T')
 
 

pixeltable/functions/math.py

@@ -0,0 +1,67 @@
+import builtins
+import math
+from typing import Optional
+
+import sqlalchemy as sql
+
+import pixeltable as pxt
+from pixeltable.utils.code import local_public_names
+
+
+@pxt.udf(is_method=True)
+def abs(self: float) -> float:
+    return builtins.abs(self)
+
+
+@abs.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.func.abs(self)
+
+
+@pxt.udf(is_method=True)
+def ceil(self: float) -> float:
+    # This ensures the same behavior as SQL
+    if math.isfinite(self):
+        return float(math.ceil(self))
+    else:
+        return self
+
+
+@ceil.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.func.ceiling(self)
+
+
+@pxt.udf(is_method=True)
+def floor(self: float) -> float:
+    # This ensures the same behavior as SQL
+    if math.isfinite(self):
+        return float(math.floor(self))
+    else:
+        return self
+
+
+@floor.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.func.floor(self)
+
+
+@pxt.udf(is_method=True)
+def round(self: float, digits: Optional[int] = None) -> float:
+    # Set digits explicitly to 0 to guarantee a return type of float; this ensures the same behavior as SQL
+    return builtins.round(self, digits or 0)
+
+
+@round.to_sql
+def _(self: sql.ColumnElement, digits: Optional[sql.ColumnElement] = None) -> sql.ColumnElement:
+    if digits is None:
+        return sql.func.round(self)
+    else:
+        return sql.func.round(sql.cast(self, sql.Numeric), sql.cast(digits, sql.Integer))
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/functions/openai.py

@@ -7,6 +7,7 @@ the [Working with OpenAI](https://pixeltable.readme.io/docs/working-with-openai)
 
 import base64
 import io
+import json
 import pathlib
 import uuid
 from typing import TYPE_CHECKING, Callable, Optional, TypeVar, Union
@@ -16,8 +17,8 @@ import PIL.Image
 import tenacity
 
 import pixeltable as pxt
-from pixeltable import env
-from pixeltable.func import Batch
+from pixeltable import env, exprs
+from pixeltable.func import Batch, Tools
 from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
@@ -225,6 +226,16 @@ def chat_completions(
             ]
             tbl['response'] = chat_completions(messages, model='gpt-4o-mini')
     """
+
+    if tools is not None:
+        tools = [
+            {
+                'type': 'function',
+                'function': tool
+            }
+            for tool in tools
+        ]
+
     result = _retry(_openai_client().chat.completions.create)(
         messages=messages,
         model=model,
@@ -453,6 +464,24 @@ def moderations(input: str, *, model: Optional[str] = None) -> dict:
     return result.dict()
 
 
+def invoke_tools(tools: 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))
+
+
+@pxt.udf
+def _openai_response_to_pxt_tool_calls(response: dict) -> Optional[dict]:
+    openai_tool_calls = response['choices'][0]['message']['tool_calls']
+    if openai_tool_calls is not None:
+        return {
+            tool_call['function']['name']: {
+                'args': json.loads(tool_call['function']['arguments'])
+            }
+            for tool_call in openai_tool_calls
+        }
+    return None
+
+
 _T = TypeVar('_T')