pixeltable 0.2.26__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/func/function_registry.py

@@ -4,25 +4,25 @@ import dataclasses
 import importlib
 import logging
 import sys
-from typing import Optional
 from uuid import UUID
 
 import sqlalchemy as sql
 
-import pixeltable.env as env
-import pixeltable.exceptions as excs
-import pixeltable.type_system as ts
+from pixeltable import env, exceptions as excs, type_system as ts
 from pixeltable.metadata import schema
+
 from .function import Function
 
 _logger = logging.getLogger('pixeltable')
 
+
 class FunctionRegistry:
     """
     A central registry for all Functions. Handles interactions with the backing store.
     Function are loaded from the store on demand.
     """
-    _instance: Optional[FunctionRegistry] = None
+
+    _instance: FunctionRegistry | None = None
 
     @classmethod
     def get(cls) -> FunctionRegistry:
@@ -30,7 +30,7 @@ class FunctionRegistry:
             cls._instance = FunctionRegistry()
         return cls._instance
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.stored_fns_by_id: dict[UUID, Function] = {}
         self.module_fns: dict[str, Function] = {}  # fqn -> Function
         self.type_methods: dict[ts.ColumnType.Type, dict[str, Function]] = {}
@@ -68,7 +68,7 @@ class FunctionRegistry:
             raise excs.Error(f'A UDF with that name already exists: {fqn}')
         self.module_fns[fqn] = fn
         if fn.is_method or fn.is_property:
-            base_type = fn.signature.parameters_by_pos[0].col_type.type_enum
+            base_type = fn.signatures[0].parameters_by_pos[0].col_type.type_enum
             if base_type not in self.type_methods:
                 self.type_methods[base_type] = {}
             if fn.name in self.type_methods[base_type]:
@@ -93,7 +93,7 @@ class FunctionRegistry:
         #         stored_fn_md.append(md)
         return list(self.module_fns.values())
 
-    # def get_function(self, *, id: Optional[UUID] = None, fqn: Optional[str] = None) -> Function:
+    # def get_function(self, *, id: UUID | None = None, fqn: str | None = None) -> Function:
     #     assert (id is not None) != (fqn is not None)
     #     if id is not None:
     #         if id not in self.stored_fns_by_id:
@@ -142,7 +142,7 @@ class FunctionRegistry:
             return list(self.type_methods[base_type].values())
         return []
 
-    def lookup_type_method(self, base_type: ts.ColumnType.Type, name: str) -> Optional[Function]:
+    def lookup_type_method(self, base_type: ts.ColumnType.Type, name: str) -> Function | None:
         """
         Look up a method (or property) by name for a given base type. If no such method is registered, return None.
         """
@@ -150,14 +150,16 @@ class FunctionRegistry:
             return self.type_methods[base_type][name]
         return None
 
-    #def create_function(self, md: schema.FunctionMd, binary_obj: bytes, dir_id: Optional[UUID] = None) -> UUID:
-    def create_stored_function(self, pxt_fn: Function, dir_id: Optional[UUID] = None) -> UUID:
+    # def create_function(self, md: schema.FunctionMd, binary_obj: bytes, dir_id: UUID | None = None) -> UUID:
+    def create_stored_function(self, pxt_fn: Function, dir_id: UUID | None = None) -> UUID:
         fn_md, binary_obj = pxt_fn.to_store()
         md = schema.FunctionMd(name=pxt_fn.name, md=fn_md, py_version=sys.version, class_name=pxt_fn.__class__.__name__)
         with env.Env.get().engine.begin() as conn:
             res = conn.execute(
-                sql.insert(schema.Function.__table__)
-                    .values(dir_id=dir_id, md=dataclasses.asdict(md), binary_obj=binary_obj))
+                sql.insert(schema.Function.__table__).values(
+                    dir_id=dir_id, md=dataclasses.asdict(md), binary_obj=binary_obj
+                )
+            )
             id = res.inserted_primary_key[0]
             _logger.info(f'Created function {pxt_fn.name} (id {id}) in store')
             self.stored_fns_by_id[id] = pxt_fn
@@ -166,8 +168,9 @@ class FunctionRegistry:
     def get_stored_function(self, id: UUID) -> Function:
         if id in self.stored_fns_by_id:
             return self.stored_fns_by_id[id]
-        stmt = sql.select(schema.Function.md, schema.Function.binary_obj, schema.Function.dir_id)\
-            .where(schema.Function.id == id)
+        stmt = sql.select(schema.Function.md, schema.Function.binary_obj, schema.Function.dir_id).where(
+            schema.Function.id == id
+        )
         with env.Env.get().engine.begin() as conn:
             row = conn.execute(stmt).fetchone()
             if row is None:
@@ -180,7 +183,7 @@ class FunctionRegistry:
             self.stored_fns_by_id[id] = instance
             return instance
 
-    # def create_function(self, fn: Function, dir_id: Optional[UUID] = None, name: Optional[str] = None) -> None:
+    # def create_function(self, fn: Function, dir_id: UUID | None = None, name: str | None = None) -> None:
     #     with env.Env.get().engine.begin() as conn:
     #         _logger.debug(f'Pickling function {name}')
     #         eval_fn_str = cloudpickle.dumps(fn.eval_fn) if fn.eval_fn is not None else None
@@ -237,7 +240,5 @@ class FunctionRegistry:
     def delete_function(self, id: UUID) -> None:
         assert id is not None
         with env.Env.get().engine.begin() as conn:
-            conn.execute(
-                sql.delete(schema.Function.__table__)
-                    .where(schema.Function.id == id))
+            conn.execute(sql.delete(schema.Function.__table__).where(schema.Function.id == id))
             _logger.info(f'Deleted function with id {id} from store')

pixeltable/func/globals.py

@@ -1,14 +1,13 @@
 import importlib
 import inspect
 from types import ModuleType
-from typing import Optional
 
 import pixeltable.exceptions as excs
 
 
-def resolve_symbol(symbol_path: str) -> Optional[object]:
+def resolve_symbol(symbol_path: str) -> object | None:
     path_elems = symbol_path.split('.')
-    module: Optional[ModuleType] = None
+    module: ModuleType | None = None
     i = len(path_elems) - 1
     while i > 0 and module is None:
         try:
@@ -28,10 +27,12 @@ def validate_symbol_path(fn_path: str) -> None:
     fn_name = path_elems[-1]
     if any(el == '<locals>' for el in path_elems):
         raise excs.Error(
-            f'{fn_name}(): nested functions are not supported. Move the function to the module level or into a class.')
+            f'{fn_name}(): nested functions are not supported. Move the function to the module level or into a class.'
+        )
     if any(not el.isidentifier() for el in path_elems):
         raise excs.Error(
-            f'{fn_name}(): cannot resolve symbol path {fn_path}. Move the function to the module level or into a class.')
+            f'{fn_name}(): cannot resolve symbol path {fn_path}. Move the function to the module level or into a class.'
+        )
 
 
 def get_caller_module_path() -> str:

pixeltable/func/mcp.py

@@ -0,0 +1,74 @@
+import asyncio
+import inspect
+from typing import TYPE_CHECKING, Any
+
+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: mcp.types.ListToolsResult | None = 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

@@ -1,69 +1,77 @@
 from __future__ import annotations
 
 import inspect
-from typing import Any, Callable, Optional
+from functools import reduce
+from typing import TYPE_CHECKING, Any, Callable, Iterable, overload
 
-import sqlalchemy as sql
-
-import pixeltable as pxt
-from pixeltable import exprs
+from pixeltable import catalog, exceptions as excs, exprs, func, type_system as ts
 
 from .function import Function
 from .signature import Signature
 
+if TYPE_CHECKING:
+    from pixeltable import Query
+
 
 class QueryTemplateFunction(Function):
-    """A parameterized query/DataFrame from which an executable DataFrame is created with a function call."""
+    """A parameterized query from which an executable Query is created with a function call."""
+
+    template_df: 'Query' | None
+    self_name: str | None
+    _comment: str | None
 
     @classmethod
     def create(
-        cls, template_callable: Callable, param_types: Optional[list[pxt.ColumnType]], path: str, name: str
+        cls, template_callable: Callable, param_types: list[ts.ColumnType] | None, path: str, name: str
     ) -> QueryTemplateFunction:
         # we need to construct a template df and a signature
         py_sig = inspect.signature(template_callable)
         py_params = list(py_sig.parameters.values())
         params = Signature.create_parameters(py_params=py_params, param_types=param_types)
-        # invoke template_callable with parameter expressions to construct a DataFrame with parameters
+        # invoke template_callable with parameter expressions to construct a Query with parameters
         var_exprs = [exprs.Variable(param.name, param.col_type) for param in params]
         template_df = template_callable(*var_exprs)
-        from pixeltable import DataFrame
-        assert isinstance(template_df, DataFrame)
+        from pixeltable import Query
+
+        assert isinstance(template_df, Query)
         # we take params and return json
-        sig = Signature(return_type=pxt.JsonType(), parameters=params)
-        return QueryTemplateFunction(template_df, sig, path=path, name=name)
+        sig = Signature(return_type=ts.JsonType(), parameters=params)
+        return QueryTemplateFunction(template_df, sig, path=path, name=name, comment=inspect.getdoc(template_callable))
 
     def __init__(
-            self, template_df: Optional['pxt.DataFrame'], sig: Optional[Signature], path: Optional[str] = None,
-            name: Optional[str] = None,
+        self,
+        template_df: 'Query' | None,
+        sig: Signature,
+        path: str | None = None,
+        name: str | None = None,
+        comment: str | None = None,
     ):
-        super().__init__(sig, self_path=path)
+        assert sig is not None
+        super().__init__([sig], self_path=path)
         self.self_name = name
         self.template_df = template_df
+        self._comment = comment
 
-        # if we're running as part of an ongoing update operation, we need to use the same connection, otherwise
-        # we end up with a deadlock
-        # TODO: figure out a more general way to make execution state available
-        self.conn: Optional[sql.engine.Connection] = None
+    def _update_as_overload_resolution(self, signature_idx: int) -> None:
+        pass  # only one signature supported for QueryTemplateFunction
 
-        # convert defaults to Literals
-        self.defaults: dict[str, exprs.Literal] = {}  # key: param name, value: default value converted to a Literal
-        param_types = self.template_df.parameters()
-        for param in [p for p in self.signature.parameters.values() if p.has_default()]:
-            assert param.name in param_types
-            param_type = param_types[param.name]
-            literal_default = exprs.Literal(param.default, col_type=param_type)
-            self.defaults[param.name] = literal_default
-
-    def set_conn(self, conn: Optional[sql.engine.Connection]) -> None:
-        self.conn = conn
+    @property
+    def is_async(self) -> bool:
+        return True
 
-    def exec(self, *args: Any, **kwargs: Any) -> Any:
+    async def aexec(self, *args: Any, **kwargs: Any) -> Any:
+        # assert not self.is_polymorphic
         bound_args = self.signature.py_signature.bind(*args, **kwargs).arguments
         # apply defaults, otherwise we might have Parameters left over
         bound_args.update(
-            {param_name: default for param_name, default in self.defaults.items() if param_name not in bound_args})
+            {
+                param.name: param.default
+                for param in self.signature.parameters.values()
+                if param.has_default() and param.name not in bound_args
+            }
+        )
         bound_df = self.template_df.bind(bound_args)
-        result = bound_df._collect(self.conn)
+        result = await bound_df._acollect()
         return list(result)
 
     @property
@@ -74,10 +82,118 @@ class QueryTemplateFunction(Function):
     def name(self) -> str:
         return self.self_name
 
+    def comment(self) -> str | None:
+        return self._comment
+
     def _as_dict(self) -> dict:
         return {'name': self.name, 'signature': self.signature.as_dict(), 'df': self.template_df.as_dict()}
 
     @classmethod
     def _from_dict(cls, d: dict) -> Function:
-        from pixeltable.dataframe import DataFrame
-        return cls(DataFrame.from_dict(d['df']), Signature.from_dict(d['signature']), name=d['name'])
+        from pixeltable._query import Query
+
+        return cls(Query.from_dict(d['df']), Signature.from_dict(d['signature']), name=d['name'])
+
+
+@overload
+def query(py_fn: Callable) -> QueryTemplateFunction: ...
+
+
+@overload
+def query(*, param_types: list[ts.ColumnType] | None = None) -> Callable[[Callable], QueryTemplateFunction]: ...
+
+
+def query(*args: Any, **kwargs: Any) -> Any:
+    def make_query_template(py_fn: Callable, param_types: list[ts.ColumnType] | None) -> QueryTemplateFunction:
+        if py_fn.__module__ != '__main__' and py_fn.__name__.isidentifier():
+            # this is a named function in a module
+            function_path = f'{py_fn.__module__}.{py_fn.__qualname__}'
+        else:
+            function_path = None
+        query_name = py_fn.__name__
+        query_fn = QueryTemplateFunction.create(py_fn, param_types=param_types, path=function_path, name=query_name)
+        return query_fn
+
+        # TODO: verify that the inferred return type matches that of the template
+        # TODO: verify that the signature doesn't contain batched parameters
+
+    if len(args) == 1:
+        assert len(kwargs) == 0 and callable(args[0])
+        return make_query_template(args[0], None)
+    else:
+        assert len(args) == 0 and len(kwargs) == 1 and 'param_types' in kwargs
+        return lambda py_fn: make_query_template(py_fn, kwargs['param_types'])
+
+
+def retrieval_udf(
+    table: catalog.Table,
+    name: str | None = None,
+    description: str | None = None,
+    parameters: Iterable[str | exprs.ColumnRef] | None = None,
+    limit: int | None = 10,
+) -> func.QueryTemplateFunction:
+    """
+    Constructs a retrieval UDF for the given table. The retrieval UDF is a UDF whose parameters are
+    columns of the table and whose return value is a list of rows from the table. The return value of
+    ```python
+    f(col1=x, col2=y, ...)
+    ```
+    will be a list of all rows from the table that match the specified arguments.
+
+    Args:
+        table: The table to use as the dataset for the retrieval tool.
+        name: The name of the tool. If not specified, then the name of the table will be used by default.
+        description: The description of the tool. If not specified, then a default description will be generated.
+        parameters: The columns of the table to use as parameters. If not specified, all data columns
+            (non-computed columns) will be used as parameters.
+
+            All of the specified parameters will be required parameters of the tool, regardless of their status
+            as columns.
+        limit: The maximum number of rows to return. If not specified, then all matching rows will be returned.
+
+    Returns:
+        A list of dictionaries containing data from the table, one per row that matches the input arguments.
+        If there are no matching rows, an empty list will be returned.
+    """
+    # Argument validation
+    col_refs: list[exprs.ColumnRef]
+    # 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 col_names if not table[col_name].col.is_computed]
+    else:
+        for param in parameters:
+            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]
+
+    if len(col_refs) == 0:
+        raise excs.Error('Parameter list cannot be empty.')
+
+    # Construct the Query
+    predicates = [col_ref == exprs.Variable(col_ref.col.name, col_ref.col.col_type) for col_ref in col_refs]
+    where_clause = reduce(lambda c1, c2: c1 & c2, predicates)
+    df = table.select().where(where_clause)
+    if limit is not None:
+        df = df.limit(limit)
+
+    # Construct the signature
+    query_params = [
+        func.Parameter(col_ref.col.name, col_ref.col.col_type, inspect.Parameter.POSITIONAL_OR_KEYWORD)
+        for col_ref in col_refs
+    ]
+    query_signature = func.Signature(return_type=ts.JsonType(), parameters=query_params)
+
+    # Construct a name and/or description if not provided
+    if name is None:
+        name = table._name
+    if description is None:
+        description = (
+            f'Retrieves an entry from the dataset {name!r} that matches the given parameters.\n\nParameters:\n'
+        )
+        description += '\n'.join(
+            [f'    {col_ref.col.name}: of type `{col_ref.col.col_type._to_base_str()}`' for col_ref in col_refs]
+        )
+
+    fn = func.QueryTemplateFunction(df, query_signature, name=name, comment=description)
+    return fn