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

pixeltable/functions/bedrock.py

@@ -1,5 +1,12 @@
+"""
+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
+from typing import TYPE_CHECKING, Any
 
 import pixeltable as pxt
 from pixeltable import env, exprs
@@ -29,10 +36,10 @@ def converse(
     messages: list[dict[str, Any]],
     *,
     model_id: str,
-    system: Optional[list[dict[str, Any]]] = None,
-    inference_config: Optional[dict] = None,
-    additional_model_request_fields: Optional[dict] = None,
-    tool_config: Optional[list[dict]] = None,
+    system: list[dict[str, Any]] | None = None,
+    inference_config: dict | None = None,
+    additional_model_request_fields: dict | None = None,
+    tool_config: list[dict] | None = None,
 ) -> dict:
     """
     Generate a conversation response.
@@ -104,7 +111,7 @@ def invoke_tools(tools: Tools, response: exprs.Expr) -> exprs.InlineDict:
 
 
 @pxt.udf
-def _bedrock_response_to_pxt_tool_calls(response: dict) -> Optional[dict]:
+def _bedrock_response_to_pxt_tool_calls(response: dict) -> dict | None:
     if response.get('stopReason') != 'tool_use':
         return None
 

pixeltable/functions/date.py

@@ -83,7 +83,7 @@ def make_date(year: int, month: int, day: int) -> date:
 
 @make_date.to_sql
 def _(year: sql.ColumnElement, month: sql.ColumnElement, day: sql.ColumnElement) -> sql.ColumnElement:
-    return sql.func.make_date(sql.cast(year, sql.Integer), sql.cast(month, sql.Integer), sql.cast(day, sql.Integer))
+    return sql.func.make_date(year.cast(sql.Integer), month.cast(sql.Integer), day.cast(sql.Integer))
 
 
 @pxt.udf(is_method=True)

pixeltable/functions/deepseek.py

@@ -1,5 +1,12 @@
+"""
+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
+from typing import TYPE_CHECKING, Any
 
 import httpx
 
@@ -26,14 +33,14 @@ def _deepseek_client() -> 'openai.AsyncOpenAI':
     return env.Env.get().get_client('deepseek')
 
 
-@pxt.udf
+@pxt.udf(resource_pool='request-rate:deepseek')
 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,
+    model_kwargs: dict[str, Any] | None = None,
+    tools: list[dict[str, Any]] | None = None,
+    tool_choice: dict[str, Any] | None = None,
 ) -> dict:
     """
     Creates a model response for the given chat conversation.
@@ -43,6 +50,10 @@ async def chat_completions(
 
     Deepseek uses the OpenAI SDK, so you will need to install the `openai` package to use this UDF.
 
+    Request throttling:
+    Applies the rate limit set in the config (section `deepseek`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
     __Requirements:__
 
     - `pip install openai`
@@ -63,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/fireworks.py

@@ -5,7 +5,7 @@ first `pip install fireworks-ai` and configure your Fireworks AI credentials, as
 the [Working with Fireworks](https://pixeltable.readme.io/docs/working-with-fireworks) tutorial.
 """
 
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 import pixeltable as pxt
 from pixeltable import env
@@ -29,7 +29,7 @@ def _fireworks_client() -> 'fireworks.client.Fireworks':
 
 @pxt.udf(resource_pool='request-rate:fireworks')
 async def chat_completions(
-    messages: list[dict[str, str]], *, model: str, model_kwargs: Optional[dict[str, Any]] = None
+    messages: list[dict[str, str]], *, model: str, model_kwargs: dict[str, Any] | None = None
 ) -> dict:
     """
     Creates a model response for the given chat conversation.

pixeltable/functions/gemini.py

@@ -7,14 +7,15 @@ the [Working with Gemini](https://pixeltable.readme.io/docs/working-with-gemini)
 
 import asyncio
 import io
-import tempfile
 from pathlib import Path
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 import PIL.Image
 
 import pixeltable as pxt
 from pixeltable import env, exceptions as excs, exprs
+from pixeltable.utils.code import local_public_names
+from pixeltable.utils.local_store import TempStore
 
 if TYPE_CHECKING:
     from google import genai
@@ -33,14 +34,14 @@ def _genai_client() -> 'genai.client.Client':
 
 @pxt.udf(resource_pool='request-rate:gemini')
 async def generate_content(
-    contents: str, *, model: str, config: Optional[dict] = None, tools: Optional[list[dict]] = None
+    contents: str, *, model: str, config: dict | None = None, tools: list[dict] | None = None
 ) -> dict:
     """
     Generate content from the specified model. For additional details, see:
     <https://ai.google.dev/gemini-api/docs/text-generation>
 
     Request throttling:
-    Applies the rate limit set in the config (section `gemini`, key `rate_limit`). If no rate
+    Applies the rate limit set in the config (section `gemini.rate_limits`; use the model id as the key). If no rate
     limit is configured, uses a default of 600 RPM.
 
     __Requirements:__
@@ -102,7 +103,7 @@ def invoke_tools(tools: pxt.func.Tools, response: exprs.Expr) -> exprs.InlineDic
 
 
 @pxt.udf
-def _gemini_response_to_pxt_tool_calls(response: dict) -> Optional[dict]:
+def _gemini_response_to_pxt_tool_calls(response: dict) -> dict | None:
     pxt_tool_calls: dict[str, list[dict]] = {}
     for part in response['candidates'][0]['content']['parts']:
         tool_call = part.get('function_call')
@@ -122,11 +123,15 @@ def _(model: str) -> str:
 
 
 @pxt.udf(resource_pool='request-rate:imagen')
-async def generate_images(prompt: str, *, model: str, config: Optional[dict] = None) -> PIL.Image.Image:
+async def generate_images(prompt: str, *, model: str, config: dict | None = None) -> PIL.Image.Image:
     """
     Generates images based on a text description and configuration. For additional details, see:
     <https://ai.google.dev/gemini-api/docs/image-generation>
 
+    Request throttling:
+    Applies the rate limit set in the config (section `imagen.rate_limits`; use the model id as the key). If no rate
+    limit is configured, uses a default of 600 RPM.
+
     __Requirements:__
 
     - `pip install google-genai`
@@ -162,12 +167,16 @@ def _(model: str) -> str:
 
 @pxt.udf(resource_pool='request-rate:veo')
 async def generate_videos(
-    prompt: Optional[str] = None, image: Optional[PIL.Image.Image] = None, *, model: str, config: Optional[dict] = None
+    prompt: str | None = None, image: PIL.Image.Image | None = None, *, model: str, config: dict | None = None
 ) -> pxt.Video:
     """
     Generates videos based on a text description and configuration. For additional details, see:
     <https://ai.google.dev/gemini-api/docs/video-generation>
 
+    Request throttling:
+    Applies the rate limit set in the config (section `veo.rate_limits`; use the model id as the key). If no rate
+    limit is configured, uses a default of 600 RPM.
+
     __Requirements:__
 
     - `pip install google-genai`
@@ -196,7 +205,7 @@ async def generate_videos(
     if prompt is None and image is None:
         raise excs.Error('At least one of `prompt` or `image` must be provided.')
 
-    image_: Optional[types.Image] = None
+    image_: types.Image | None = None
     if image is not None:
         with io.BytesIO() as buffer:
             image.save(buffer, format='jpeg')
@@ -215,11 +224,19 @@ async def generate_videos(
     video_bytes = await _genai_client().aio.files.download(file=video.video)  # type: ignore[arg-type]
     assert video_bytes is not None
 
-    _, output_filename = tempfile.mkstemp(suffix='.mp4', dir=str(env.Env.get().tmp_dir))
-    Path(output_filename).write_bytes(video_bytes)
-    return output_filename
+    # Create a temporary file to store the video bytes
+    output_path = TempStore.create_path(extension='.mp4')
+    Path(output_path).write_bytes(video_bytes)
+    return str(output_path)
 
 
 @generate_videos.resource_pool
 def _(model: str) -> str:
     return f'request-rate:veo:{model}'
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/functions/globals.py

@@ -1,6 +1,6 @@
 import builtins
 import typing
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable
 
 import sqlalchemy as sql
 
@@ -11,7 +11,7 @@ from typing import _GenericAlias  # type: ignore[attr-defined]  # isort: skip
 
 
 # TODO: remove and replace calls with astype()
-def cast(expr: exprs.Expr, target_type: Union[ts.ColumnType, type, _GenericAlias]) -> exprs.Expr:
+def cast(expr: exprs.Expr, target_type: ts.ColumnType | type | _GenericAlias) -> exprs.Expr:
     expr.col_type = ts.ColumnType.normalize_type(target_type)
     return expr
 
@@ -19,7 +19,7 @@ def cast(expr: exprs.Expr, target_type: Union[ts.ColumnType, type, _GenericAlias
 T = typing.TypeVar('T')
 
 
-@func.uda(allows_window=True, type_substitutions=({T: Optional[int]}, {T: Optional[float]}))  # type: ignore[misc]
+@func.uda(allows_window=True, type_substitutions=({T: int | None}, {T: float | None}))  # type: ignore[misc]
 class sum(func.Aggregator, typing.Generic[T]):
     """Sums the selected integers or floats."""
 
@@ -39,7 +39,7 @@ class sum(func.Aggregator, typing.Generic[T]):
 
 
 @sum.to_sql
-def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
+def _(val: sql.ColumnElement) -> sql.ColumnElement | None:
     # This can produce a Decimal. We are deliberately avoiding an explicit cast to a Bigint here, because that can
     # cause overflows in Postgres. We're instead doing the conversion to the target type in SqlNode.__iter__().
     return sql.sql.func.sum(val)
@@ -49,7 +49,7 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     allows_window=True,
     # Allow counting non-null values of any type
     # TODO: should we have an "Any" type that can be used here?
-    type_substitutions=tuple({T: Optional[t]} for t in ts.ALL_PIXELTABLE_TYPES),  # type: ignore[misc]
+    type_substitutions=tuple({T: t | None} for t in ts.ALL_PIXELTABLE_TYPES),  # type: ignore[misc]
 )
 class count(func.Aggregator, typing.Generic[T]):
     def __init__(self) -> None:
@@ -64,13 +64,13 @@ class count(func.Aggregator, typing.Generic[T]):
 
 
 @count.to_sql
-def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
+def _(val: sql.ColumnElement) -> sql.ColumnElement | None:
     return sql.sql.func.count(val)
 
 
 @func.uda(
     allows_window=True,
-    type_substitutions=tuple({T: Optional[t]} for t in (str, int, float, bool, ts.Timestamp)),  # type: ignore[misc]
+    type_substitutions=tuple({T: t | None} for t in (str, int, float, bool, ts.Timestamp)),  # type: ignore[misc]
 )
 class min(func.Aggregator, typing.Generic[T]):
     def __init__(self) -> None:
@@ -89,7 +89,7 @@ class min(func.Aggregator, typing.Generic[T]):
 
 
 @min.to_sql
-def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
+def _(val: sql.ColumnElement) -> sql.ColumnElement | None:
     if val.type.python_type is bool:
         # TODO: min/max aggregation of booleans is not supported in Postgres (but it is in Python).
         # Right now we simply force the computation to be done in Python; we might consider implementing an alternate
@@ -100,7 +100,7 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
 
 @func.uda(
     allows_window=True,
-    type_substitutions=tuple({T: Optional[t]} for t in (str, int, float, bool, ts.Timestamp)),  # type: ignore[misc]
+    type_substitutions=tuple({T: t | None} for t in (str, int, float, bool, ts.Timestamp)),  # type: ignore[misc]
 )
 class max(func.Aggregator, typing.Generic[T]):
     def __init__(self) -> None:
@@ -119,14 +119,14 @@ class max(func.Aggregator, typing.Generic[T]):
 
 
 @max.to_sql
-def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
+def _(val: sql.ColumnElement) -> sql.ColumnElement | None:
     if val.type.python_type is bool:
         # TODO: see comment in @min.to_sql.
         return None
     return sql.sql.func.max(val)
 
 
-@func.uda(type_substitutions=({T: Optional[int]}, {T: Optional[float]}))  # type: ignore[misc]
+@func.uda(type_substitutions=({T: int | None}, {T: float | None}))  # type: ignore[misc]
 class mean(func.Aggregator, typing.Generic[T]):
     def __init__(self) -> None:
         self.sum: T = None
@@ -141,14 +141,14 @@ class mean(func.Aggregator, typing.Generic[T]):
             self.sum += val  # type: ignore[operator]
         self.count += 1
 
-    def value(self) -> Optional[float]:  # Always a float
+    def value(self) -> float | None:  # Always a float
         if self.count == 0:
             return None
         return self.sum / self.count  # type: ignore[operator]
 
 
 @mean.to_sql
-def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
+def _(val: sql.ColumnElement) -> sql.ColumnElement | None:
     return sql.sql.func.avg(val)
 
 

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
+
+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: dict[str, Any] | None = None,
+    tools: list[dict[str, Any]] | None = None,
+    tool_choice: dict[str, Any] | None = 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 `llama-3.1-8b-instant`
+        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='llama-3.1-8b-instant'))
+    """
+    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__