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

pixeltable/functions/math.py

@@ -0,0 +1,169 @@
+"""
+Pixeltable UDFs for mathematical operations.
+
+Example:
+```python
+import pixeltable as pxt
+
+t = pxt.get_table(...)
+t.select(t.float_col.floor()).collect()
+```
+"""
+
+import builtins
+import math
+
+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 the absolute value of the given number.
+
+    Equivalent to Python [`builtins.abs()`](https://docs.python.org/3/library/functions.html#abs).
+    """
+    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:
+    """
+    Return the ceiling of the given number.
+
+    Equivalent to Python [`float(math.ceil(self))`](https://docs.python.org/3/library/math.html#math.ceil) if `self`
+    is finite, or `self` itself if `self` is infinite. (This is slightly different from the default behavior of
+    `math.ceil(self)`, which always returns an `int` and raises an error if `self` is infinite. The behavior in
+    Pixeltable generalizes the Python operator and is chosen to align with the SQL standard.)
+    """
+    # 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:
+    """
+    Return the ceiling of the given number.
+
+    Equivalent to Python [`float(math.floor(self))`](https://docs.python.org/3/library/math.html#math.ceil) if `self`
+    is finite, or `self` itself if `self` is infinite. (This is slightly different from the default behavior of
+    `math.floor(self)`, which always returns an `int` and raises an error if `self` is infinite. The behavior of
+    Pixeltable generalizes the Python operator and is chosen to align with the SQL standard.)
+    """
+    # 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: int | None = None) -> float:
+    """
+    Round a number to a given precision in decimal digits.
+
+    Equivalent to Python [`builtins.round(self, digits or 0)`](https://docs.python.org/3/library/functions.html#round).
+    Note that if `digits` is not specified, the behavior matches `builtins.round(self, 0)` rather than
+    `builtins.round(self)`; this ensures that the return type is always `float` (as in SQL) rather than `int`.
+    """
+    # 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: sql.ColumnElement | None = None) -> sql.ColumnElement:
+    if digits is None:
+        return sql.func.round(self)
+    else:
+        return sql.cast(sql.func.round(sql.cast(self, sql.Numeric), sql.cast(digits, sql.Integer)), sql.Float)
+
+
+@pxt.udf(is_method=True)
+def pow(self: int, other: int) -> float:
+    """
+    Raise `self` to the power of `other`.
+
+    Equivalent to Python [`self ** other`](https://docs.python.org/3/library/functions.html#pow).
+    """
+    return self**other
+
+
+@pow.to_sql
+def _(self: sql.ColumnElement, other: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.func.pow(self, other)
+
+
+@pxt.udf(is_method=True)
+def bitwise_and(self: int, other: int) -> int:
+    """
+    Bitwise AND of two integers.
+
+    Equivalent to Python
+    [`self & other`](https://docs.python.org/3/library/stdtypes.html#bitwise-operations-on-integer-types).
+    """
+    return self & other
+
+
+@bitwise_and.to_sql
+def _(self: sql.ColumnElement, other: sql.ColumnElement) -> sql.ColumnElement:
+    return self.bitwise_and(other)
+
+
+@pxt.udf(is_method=True)
+def bitwise_or(self: int, other: int) -> int:
+    """
+    Bitwise OR of two integers.
+
+    Equivalent to Python
+    [`self | other`](https://docs.python.org/3/library/stdtypes.html#bitwise-operations-on-integer-types).
+    """
+    return self | other
+
+
+@bitwise_or.to_sql
+def _(self: sql.ColumnElement, other: sql.ColumnElement) -> sql.ColumnElement:
+    return self.bitwise_or(other)
+
+
+@pxt.udf(is_method=True)
+def bitwise_xor(self: int, other: int) -> int:
+    """
+    Bitwise XOR of two integers.
+
+    Equivalent to Python
+    [`self ^ other`](https://docs.python.org/3/library/stdtypes.html#bitwise-operations-on-integer-types).
+    """
+    return self ^ other
+
+
+@bitwise_xor.to_sql
+def _(self: sql.ColumnElement, other: sql.ColumnElement) -> sql.ColumnElement:
+    return self.bitwise_xor(other)
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/functions/mistralai.py

@@ -1,26 +1,28 @@
 """
-Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
+Pixeltable UDFs
 that wrap various endpoints from the Mistral AI API. In order to use them, you must
 first `pip install mistralai` and configure your Mistral AI credentials, as described in
-the [Working with Mistral AI](https://pixeltable.readme.io/docs/working-with-mistralai) tutorial.
+the [Working with Mistral AI](https://docs.pixeltable.com/notebooks/integrations/working-with-mistralai) tutorial.
 """
 
-from typing import TYPE_CHECKING, Optional, TypeVar, Union
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
 import pixeltable as pxt
+import pixeltable.type_system as ts
 from pixeltable.env import Env, register_client
 from pixeltable.func.signature import Batch
 from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
-    import mistralai.types.basemodel
+    import mistralai
 
 
 @register_client('mistral')
 def _(api_key: str) -> 'mistralai.Mistral':
     import mistralai
+
     return mistralai.Mistral(api_key=api_key)
 
 
@@ -28,18 +30,9 @@ def _mistralai_client() -> 'mistralai.Mistral':
     return Env.get().get_client('mistral')
 
 
-@pxt.udf
-def chat_completions(
-    messages: list[dict[str, str]],
-    *,
-    model: str,
-    temperature: Optional[float] = 0.7,
-    top_p: Optional[float] = 1.0,
-    max_tokens: Optional[int] = None,
-    stop: Optional[list[str]] = None,
-    random_seed: Optional[int] = None,
-    response_format: Optional[dict] = None,
-    safe_prompt: Optional[bool] = False,
+@pxt.udf(resource_pool='request-rate:mistral')
+async def chat_completions(
+    messages: list[dict[str, str]], *, model: str, model_kwargs: dict[str, Any] | None = None
 ) -> dict:
     """
     Chat Completion API.
@@ -47,6 +40,10 @@ def chat_completions(
     Equivalent to the Mistral AI `chat/completions` API endpoint.
     For additional details, see: <https://docs.mistral.ai/api/#tag/chat>
 
+    Request throttling:
+    Applies the rate limit set in the config (section `mistral`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
     __Requirements:__
 
     - `pip install mistralai`
@@ -54,8 +51,8 @@ def chat_completions(
     Args:
         messages: The prompt(s) to generate completions for.
         model: ID of the model to use. (See overview here: <https://docs.mistral.ai/getting-started/models/>)
-
-    For details on the other parameters, see: <https://docs.mistral.ai/api/#tag/chat>
+        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>
 
     Returns:
         A dictionary containing the response and other metadata.
@@ -65,41 +62,32 @@ def chat_completions(
         to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
 
         >>> messages = [{'role': 'user', 'content': tbl.prompt}]
-        ... tbl['response'] = completions(messages, model='mistral-latest-small')
+        ... tbl.add_computed_column(response=completions(messages, model='mistral-latest-small'))
     """
+    if model_kwargs is None:
+        model_kwargs = {}
+
     Env.get().require_package('mistralai')
-    return _mistralai_client().chat.complete(
+    result = await _mistralai_client().chat.complete_async(
         messages=messages,  # type: ignore[arg-type]
         model=model,
-        temperature=temperature,
-        top_p=top_p,
-        max_tokens=_opt(max_tokens),
-        stop=stop,
-        random_seed=_opt(random_seed),
-        response_format=response_format,  # type: ignore[arg-type]
-        safe_prompt=safe_prompt,
-    ).dict()
-
-
-@pxt.udf
-def fim_completions(
-    prompt: str,
-    *,
-    model: str,
-    temperature: Optional[float] = 0.7,
-    top_p: Optional[float] = 1.0,
-    max_tokens: Optional[int] = None,
-    min_tokens: Optional[int] = None,
-    stop: Optional[list[str]] = None,
-    random_seed: Optional[int] = None,
-    suffix: Optional[str] = None,
-) -> dict:
+        **model_kwargs,
+    )
+    return result.dict()
+
+
+@pxt.udf(resource_pool='request-rate:mistral')
+async def fim_completions(prompt: str, *, model: str, model_kwargs: dict[str, Any] | None = None) -> dict:
     """
     Fill-in-the-middle Completion API.
 
     Equivalent to the Mistral AI `fim/completions` API endpoint.
     For additional details, see: <https://docs.mistral.ai/api/#tag/fim>
 
+    Request throttling:
+    Applies the rate limit set in the config (section `mistral`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
     __Requirements:__
 
     - `pip install mistralai`
@@ -107,6 +95,8 @@ def fim_completions(
     Args:
         prompt: The text/code to complete.
         model: ID of the model to use. (See overview here: <https://docs.mistral.ai/getting-started/models/>)
+        model_kwargs: Additional keyword args for the Mistral `fim/completions` API.
+            For details on the available parameters, see: <https://docs.mistral.ai/api/#tag/fim>
 
     For details on the other parameters, see: <https://docs.mistral.ai/api/#tag/fim>
 
@@ -117,35 +107,31 @@ def fim_completions(
         Add a computed column that applies the model `codestral-latest`
         to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
 
-        >>> tbl['response'] = completions(tbl.prompt, model='codestral-latest')
+        >>> tbl.add_computed_column(response=completions(tbl.prompt, model='codestral-latest'))
     """
+    if model_kwargs is None:
+        model_kwargs = {}
+
     Env.get().require_package('mistralai')
-    return _mistralai_client().fim.complete(
-        prompt=prompt,
-        model=model,
-        temperature=temperature,
-        top_p=top_p,
-        max_tokens=_opt(max_tokens),
-        min_tokens=_opt(min_tokens),
-        stop=stop,
-        random_seed=_opt(random_seed),
-        suffix=_opt(suffix)
-    ).dict()
+    result = await _mistralai_client().fim.complete_async(prompt=prompt, model=model, **model_kwargs)
+    return result.dict()
 
 
-_embedding_dimensions_cache: dict[str, int] = {
-    'mistral-embed': 1024
-}
+_embedding_dimensions_cache: dict[str, int] = {'mistral-embed': 1024}
 
 
-@pxt.udf(batch_size=16)
-def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,), pxt.Float]]:
+@pxt.udf(batch_size=16, resource_pool='request-rate:mistral')
+async def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,), pxt.Float]]:  # noqa: RUF029
     """
     Embeddings API.
 
     Equivalent to the Mistral AI `embeddings` API endpoint.
     For additional details, see: <https://docs.mistral.ai/api/#tag/embeddings>
 
+    Request throttling:
+    Applies the rate limit set in the config (section `mistral`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
     __Requirements:__
 
     - `pip install mistralai`
@@ -158,29 +144,18 @@ def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,), pxt
         An array representing the application of the given embedding to `input`.
     """
     Env.get().require_package('mistralai')
-    result = _mistralai_client().embeddings.create(
-        inputs=input,
-        model=model,
-    )
+    result = _mistralai_client().embeddings.create(inputs=input, model=model)
     return [np.array(data.embedding, dtype=np.float64) for data in result.data]
 
 
 @embeddings.conditional_return_type
-def _(model: str) -> pxt.ArrayType:
+def _(model: str) -> ts.ArrayType:
     dimensions = _embedding_dimensions_cache.get(model)  # `None` if unknown model
-    return pxt.ArrayType((dimensions,), dtype=pxt.FloatType())
-
-
-_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
+    return ts.ArrayType((dimensions,), dtype=ts.FloatType())
 
 
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/net.py

@@ -0,0 +1,70 @@
+"""
+Pixeltable UDF for converting media file URIs to presigned HTTP URLs.
+"""
+
+from pixeltable import exceptions as excs
+from pixeltable.func.udf import udf
+from pixeltable.utils.code import local_public_names
+from pixeltable.utils.object_stores import ObjectOps, ObjectPath, StorageTarget
+
+
+@udf
+def presigned_url(uri: str, expiration_seconds: int) -> str:
+    """
+    Convert a blob storage URI to a presigned HTTP URL for direct access.
+
+    Generates a time-limited, publicly accessible URL from cloud storage URIs
+    (S3, GCS, Azure, etc.) that can be used to serve media files over HTTP.
+
+    Note:
+        This function uses presigned URLs from storage providers. Provider-specific
+        limitations apply:
+
+        - Google Cloud Storage: maximum 7-day expiration
+        - AWS S3: requires proper region configuration
+        - Azure: subject to storage account access policies
+
+    Args:
+        uri: The media file URI (e.g., `s3://bucket/path`, `gs://bucket/path`, `azure://container/path`)
+        expiration_seconds: How long the URL remains valid
+
+    Returns:
+        A presigned HTTP URL for accessing the file
+
+    Raises:
+        Error: If the URI is a local file:// path
+
+    Examples:
+        Generate a presigned URL for a video column with 1-hour expiration:
+
+        >>> tbl.select(
+        ...     original_url=tbl.video.fileurl,
+        ...     presigned_url=pxtf.net.presigned_url(tbl.video.fileurl, 3600)
+        ... ).collect()
+    """
+    if not uri:
+        return uri
+
+    # Parse the object storage address from the URI
+    soa = ObjectPath.parse_object_storage_addr(uri, allow_obj_name=True)
+
+    # HTTP/HTTPS URLs are already publicly accessible
+    if soa.storage_target == StorageTarget.HTTP_STORE:
+        return uri
+
+    # For file:// URLs, we can't generate presigned URLs
+    if soa.storage_target == StorageTarget.LOCAL_STORE:
+        raise excs.Error(
+            'Cannot generate presigned URL for local file:// URLs. '
+            'Please use cloud storage (S3, GCS, Azure) for presigned URLs.'
+        )
+
+    store = ObjectOps.get_store(soa, allow_obj_name=True)
+    return store.create_presigned_url(soa, expiration_seconds)
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/functions/ollama.py

@@ -1,4 +1,11 @@
-from typing import TYPE_CHECKING, Optional
+"""
+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
 
 import numpy as np
 
@@ -14,10 +21,11 @@ if TYPE_CHECKING:
 @env.register_client('ollama')
 def _(host: str) -> 'ollama.Client':
     import ollama
+
     return ollama.Client(host=host)
 
 
-def _ollama_client() -> Optional['ollama.Client']:
+def _ollama_client() -> 'ollama.Client | None':
     try:
         return env.Env.get().get_client('ollama')
     except Exception:
@@ -32,10 +40,10 @@ def generate(
     suffix: str = '',
     system: str = '',
     template: str = '',
-    context: Optional[list[int]] = None,
+    context: list[int] | None = None,
     raw: bool = False,
-    format: str = '',
-    options: Optional[dict] = None,
+    format: str | None = None,
+    options: dict | None = None,
 ) -> dict:
     """
     Generate a response for a given prompt with a provided model.
@@ -44,14 +52,14 @@ def generate(
         prompt: The prompt to generate a response for.
         model: The model name.
         suffix: The text after the model response.
-        format: The format of the response; must be one of `'json'` or `''` (the empty string).
+        format: The format of the response; must be one of `'json'` or `None`.
         system: System message.
         template: Prompt template to use.
         context: The context parameter returned from a previous call to `generate()`.
         raw: If `True`, no formatting will be applied to the prompt.
-        options: Additional options to pass to the `chat` call, such as `max_tokens`, `temperature`, `top_p`, and `top_k`.
-            For details, see the
-            [Valid Parameters and Values](https://github.com/ollama/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values)
+        options: Additional options for the Ollama `chat` call, such as `max_tokens`, `temperature`, `top_p`, and
+            `top_k`. For details, see the
+            [Valid Parameters and Values](https://github.com/ollama/ollama/blob/main/docs/modelfile.mdx#valid-parameters-and-values)
             section of the Ollama documentation.
     """
     env.Env.get().require_package('ollama')
@@ -76,9 +84,9 @@ def chat(
     messages: list[dict],
     *,
     model: str,
-    tools: Optional[list[dict]] = None,
-    format: str = '',
-    options: Optional[dict] = None,
+    tools: list[dict] | None = None,
+    format: str | None = None,
+    options: dict | None = None,
 ) -> dict:
     """
     Generate the next message in a chat with a provided model.
@@ -87,32 +95,22 @@ def chat(
         messages: The messages of the chat.
         model: The model name.
         tools: Tools for the model to use.
-        format: The format of the response; must be one of `'json'` or `''` (the empty string).
-        options: Additional options to pass to the `chat` call, such as `max_tokens`, `temperature`, `top_p`, and `top_k`.
-            For details, see the
-            [Valid Parameters and Values](https://github.com/ollama/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values)
+        format: The format of the response; must be one of `'json'` or `None`.
+        options: Additional options to pass to the `chat` call, such as `max_tokens`, `temperature`, `top_p`, and
+            `top_k`. For details, see the
+            [Valid Parameters and Values](https://github.com/ollama/ollama/blob/main/docs/modelfile.mdx#valid-parameters-and-values)
             section of the Ollama documentation.
     """
     env.Env.get().require_package('ollama')
     import ollama
 
     client = _ollama_client() or ollama
-    return client.chat(
-        model=model,
-        messages=messages,
-        tools=tools,
-        format=format,
-        options=options,
-    ).dict()  # type: ignore[call-overload]
+    return client.chat(model=model, messages=messages, tools=tools, format=format, options=options).dict()  # type: ignore[call-overload]
 
 
 @pxt.udf(batch_size=16)
 def embed(
-    input: Batch[str],
-    *,
-    model: str,
-    truncate: bool = True,
-    options: Optional[dict] = None,
+    input: Batch[str], *, model: str, truncate: bool = True, options: dict | None = None
 ) -> Batch[pxt.Array[(None,), pxt.Float]]:
     """
     Generate embeddings from a model.
@@ -124,24 +122,19 @@ def embed(
             Returns error if false and context length is exceeded.
         options: Additional options to pass to the `embed` call.
             For details, see the
-            [Valid Parameters and Values](https://github.com/ollama/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values)
+            [Valid Parameters and Values](https://github.com/ollama/ollama/blob/main/docs/modelfile.mdx#valid-parameters-and-values)
             section of the Ollama documentation.
     """
     env.Env.get().require_package('ollama')
     import ollama
 
     client = _ollama_client() or ollama
-    results = client.embed(
-        model=model,
-        input=input,
-        truncate=truncate,
-        options=options,
-    ).dict()
+    results = client.embed(model=model, input=input, truncate=truncate, options=options).dict()
     return [np.array(data, dtype=np.float64) for data in results['embeddings']]
 
 
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__