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

pixeltable/functions/gemini.py

@@ -0,0 +1,249 @@
+"""
+Pixeltable UDFs
+that wrap various endpoints from the Google Gemini API. In order to use them, you must
+first `pip install google-genai` and configure your Gemini credentials, as described in
+the [Working with Gemini](https://docs.pixeltable.com/notebooks/integrations/working-with-gemini) tutorial.
+"""
+
+import asyncio
+import io
+from pathlib import Path
+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
+
+
+@env.register_client('gemini')
+def _(api_key: str) -> 'genai.client.Client':
+    from google import genai
+
+    return genai.client.Client(api_key=api_key)
+
+
+def _genai_client() -> 'genai.client.Client':
+    return env.Env.get().get_client('gemini')
+
+
+@pxt.udf(resource_pool='request-rate:gemini')
+async def generate_content(
+    contents: pxt.Json, *, model: str, config: dict | None = None, tools: list[dict] | None = None
+) -> dict:
+    """
+    Generate content from the specified model.
+
+    Request throttling:
+    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:__
+
+    - `pip install google-genai`
+
+    Args:
+        contents: The input content to generate from. Can be a prompt, or a list containing images and text
+            prompts, as described in: <https://ai.google.dev/gemini-api/docs/text-generation>
+        model: The name of the model to use.
+        config: Configuration for generation, corresponding to keyword arguments of
+            `genai.types.GenerateContentConfig`. For details on the parameters, see:
+            <https://googleapis.github.io/python-genai/genai.html#genai.types.GenerateContentConfig>
+        tools: An optional list of Pixeltable tools to use. It is also possible to specify tools manually via the
+            `config['tools']` parameter, but at most one of `config['tools']` or `tools` may be used.
+
+    Returns:
+        A dictionary containing the response and other metadata.
+
+    Examples:
+        Add a computed column that applies the model `gemini-2.5-flash`
+        to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
+
+        >>> tbl.add_computed_column(response=generate_content(tbl.prompt, model='gemini-2.5-flash'))
+
+        Add a computed column that applies the model `gemini-2.5-flash` for image understanding
+    """
+    env.Env.get().require_package('google.genai')
+    from google.genai import types
+
+    config_: types.GenerateContentConfig
+    if config is None and tools is None:
+        config_ = None
+    else:
+        if config is None:
+            config_ = types.GenerateContentConfig()
+        else:
+            config_ = types.GenerateContentConfig(**config)
+        if tools is not None:
+            gemini_tools = [__convert_pxt_tool(tool) for tool in tools]
+            config_.tools = [types.Tool(function_declarations=gemini_tools)]
+
+    response = await _genai_client().aio.models.generate_content(model=model, contents=contents, config=config_)
+    return response.model_dump()
+
+
+def __convert_pxt_tool(pxt_tool: dict) -> dict:
+    return {
+        'name': pxt_tool['name'],
+        'description': pxt_tool['description'],
+        'parameters': {
+            'type': 'object',
+            'properties': pxt_tool['parameters']['properties'],
+            'required': pxt_tool['required'],
+        },
+    }
+
+
+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(_gemini_response_to_pxt_tool_calls(response))
+
+
+@pxt.udf
+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')
+        if tool_call is not None:
+            tool_name = tool_call['name']
+            if tool_name not in pxt_tool_calls:
+                pxt_tool_calls[tool_name] = []
+            pxt_tool_calls[tool_name].append({'args': tool_call['args']})
+    if len(pxt_tool_calls) == 0:
+        return None
+    return pxt_tool_calls
+
+
+@generate_content.resource_pool
+def _(model: str) -> str:
+    return f'request-rate:gemini:{model}'
+
+
+@pxt.udf(resource_pool='request-rate:imagen')
+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`
+
+    Args:
+        prompt: A text description of the images to generate.
+        model: The model to use.
+        config: Configuration for generation, corresponding to keyword arguments of
+            `genai.types.GenerateImagesConfig`. For details on the parameters, see:
+            <https://googleapis.github.io/python-genai/genai.html#genai.types.GenerateImagesConfig>
+
+    Returns:
+        The generated image.
+
+    Examples:
+        Add a computed column that applies the model `imagen-4.0-generate-001`
+        to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
+
+        >>> tbl.add_computed_column(response=generate_images(tbl.prompt, model='imagen-4.0-generate-001'))
+    """
+    env.Env.get().require_package('google.genai')
+    from google.genai.types import GenerateImagesConfig
+
+    config_ = GenerateImagesConfig(**config) if config else None
+    response = await _genai_client().aio.models.generate_images(model=model, prompt=prompt, config=config_)
+    return response.generated_images[0].image._pil_image
+
+
+@generate_images.resource_pool
+def _(model: str) -> str:
+    return f'request-rate:imagen:{model}'
+
+
+@pxt.udf(resource_pool='request-rate:veo')
+async def generate_videos(
+    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>
+
+    At least one of `prompt` or `image` must be provided.
+
+    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`
+
+    Args:
+        prompt: A text description of the videos to generate.
+        image: An image to use as the first frame of the video.
+        model: The model to use.
+        config: Configuration for generation, corresponding to keyword arguments of
+            `genai.types.GenerateVideosConfig`. For details on the parameters, see:
+            <https://googleapis.github.io/python-genai/genai.html#genai.types.GenerateVideosConfig>
+
+    Returns:
+        The generated video.
+
+    Examples:
+        Add a computed column that applies the model `veo-3.0-generate-001`
+        to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
+
+        >>> tbl.add_computed_column(response=generate_videos(tbl.prompt, model='veo-3.0-generate-001'))
+    """
+    env.Env.get().require_package('google.genai')
+    from google.genai import types
+
+    if prompt is None and image is None:
+        raise excs.Error('At least one of `prompt` or `image` must be provided.')
+
+    image_: types.Image | None = None
+    if image is not None:
+        with io.BytesIO() as buffer:
+            image.save(buffer, format='webp')
+            image_ = types.Image(image_bytes=buffer.getvalue(), mime_type='image/webp')
+
+    config_ = types.GenerateVideosConfig(**config) if config else None
+
+    operation = await _genai_client().aio.models.generate_videos(
+        model=model, prompt=prompt, image=image_, config=config_
+    )
+    while not operation.done:
+        await asyncio.sleep(3)
+        operation = await _genai_client().aio.operations.get(operation)
+
+    if operation.error:
+        raise Exception(f'Video generation failed: {operation.error}')
+
+    video = operation.response.generated_videos[0]
+
+    video_bytes = await _genai_client().aio.files.download(file=video.video)  # type: ignore[arg-type]
+    assert video_bytes is not None
+
+    # 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,54 +1,105 @@
 import builtins
-from typing import _GenericAlias  # type: ignore[attr-defined]
-from typing import Optional, Union
+import typing
+from typing import Any, Callable
 
 import sqlalchemy as sql
 
-import pixeltable.func as func
-import pixeltable.type_system as ts
-from pixeltable import exprs
+from pixeltable import exceptions as excs, exprs, func, type_system as ts
 from pixeltable.utils.code import local_public_names
 
+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
 
 
-@func.uda(
-    update_types=[ts.IntType(nullable=True)], value_type=ts.IntType(nullable=False),
-    allows_window=True, requires_order_by=False)
-class sum(func.Aggregator):
-    """Sums the selected integers or floats."""
-    def __init__(self):
-        self.sum: Optional[int] = None
-
-    def update(self, val: Optional[int]) -> None:
+T = typing.TypeVar('T')
+
+
+@func.uda(allows_window=True, type_substitutions=({T: int | None}, {T: float | None}))  # type: ignore[misc]
+class sum(func.Aggregator, typing.Generic[T]):
+    """
+    Aggregate function that computes the sum of non-null values of a numeric column or grouping.
+
+    Args:
+        val: The numeric value to add to the sum.
+
+    Returns:
+        The sum of the non-null values, or `None` if there are no non-null values.
+
+    Examples:
+        Sum the values in the `value` column of the table `tbl`:
+
+        >>> tbl.select(pxt.functions.sum(tbl.value)).collect()
+
+        Group by the `category` column and compute the sum of the `value` column for each category,
+        assigning the name `'category_total'` to the new column:
+
+        >>> tbl.group_by(tbl.category).select(
+        ...     tbl.category,
+        ...     category_total=pxt.functions.sum(tbl.value)
+        ... ).collect()
+    """
+
+    def __init__(self) -> None:
+        self.sum: T = None
+
+    def update(self, val: T) -> None:
         if val is None:
             return
         if self.sum is None:
             self.sum = val
         else:
-            self.sum += val
+            self.sum += val  # type: ignore[operator]
 
-    def value(self) -> Union[int, float]:
+    def value(self) -> T:
         return self.sum
 
 
 @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)
 
 
-@func.uda(update_types=[ts.IntType(nullable=True)], value_type=ts.IntType(), allows_window=True, requires_order_by=False)
-class count(func.Aggregator):
-    def __init__(self):
+@func.uda(
+    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: t | None} for t in ts.ALL_PIXELTABLE_TYPES),  # type: ignore[misc]
+)
+class count(func.Aggregator, typing.Generic[T]):
+    """
+    Aggregate function that counts the number of non-null values in a column or grouping.
+
+    Args:
+        val: The value to count.
+
+    Returns:
+        The count of non-null values.
+
+    Examples:
+        Count the number of non-null values in the `value` column of the table `tbl`:
+
+        >>> tbl.select(pxt.functions.count(tbl.value)).collect()
+
+        Group by the `category` column and compute the count of non-null values in the `value` column
+        for each category, assigning the name `'category_count'` to the new column:
+
+        >>> tbl.group_by(tbl.category).select(
+        ...     tbl.category,
+        ...     category_count=pxt.functions.count(tbl.value)
+        ... ).collect()
+    """
+
+    def __init__(self) -> None:
         self.count = 0
 
-    def update(self, val: Optional[int]) -> None:
+    def update(self, val: T) -> None:
         if val is not None:
             self.count += 1
 
@@ -57,88 +108,192 @@ class count(func.Aggregator):
 
 
 @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(
-    update_types=[ts.IntType(nullable=True)], value_type=ts.IntType(nullable=True), allows_window=True,
-    requires_order_by=False)
-class min(func.Aggregator):
-    def __init__(self):
-        self.val: Optional[int] = None
+    allows_window=True,
+    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]):
+    """
+    Aggregate function that computes the minimum value in a column or grouping.
+
+    Args:
+        val: The value to compare.
+
+    Returns:
+        The minimum value, or `None` if there are no non-null values.
+
+    Examples:
+        Compute the minimum value in the `value` column of the table `tbl`:
+
+        >>> tbl.select(pxt.functions.min(tbl.value)).collect()
 
-    def update(self, val: Optional[int]) -> None:
+        Group by the `category` column and compute the minimum value in the `value` column for each category,
+        assigning the name `'category_min'` to the new column:
+
+        >>> tbl.group_by(tbl.category).select(
+        ...     tbl.category,
+        ...     category_min=pxt.functions.min(tbl.value)
+        ... ).collect()
+    """
+
+    def __init__(self) -> None:
+        self.val: T = None
+
+    def update(self, val: T) -> None:
         if val is None:
             return
         if self.val is None:
             self.val = val
         else:
-            self.val = builtins.min(self.val, val)
+            self.val = builtins.min(self.val, val)  # type: ignore[call-overload]
 
-    def value(self) -> Optional[int]:
+    def value(self) -> T:
         return self.val
 
 
 @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
+        # way of doing it in SQL. (min/max of booleans is simply logical and/or, respectively.)
+        return None
     return sql.sql.func.min(val)
 
 
 @func.uda(
-    update_types=[ts.IntType(nullable=True)], value_type=ts.IntType(nullable=True), allows_window=True,
-    requires_order_by=False)
-class max(func.Aggregator):
-    def __init__(self):
-        self.val: Optional[int] = None
+    allows_window=True,
+    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]):
+    """
+    Aggregate function that computes the maximum value in a column or grouping.
+
+    Args:
+        val: The value to compare.
+
+    Returns:
+        The maximum value, or `None` if there are no non-null values.
+
+    Examples:
+        Compute the maximum value in the `value` column of the table `tbl`:
 
-    def update(self, val: Optional[int]) -> None:
+        >>> tbl.select(pxt.functions.max(tbl.value)).collect()
+
+        Group by the `category` column and compute the maximum value in the `value` column for each category,
+        assigning the name `'category_max'` to the new column:
+
+        >>> tbl.group_by(tbl.category).select(
+        ...     tbl.category,
+        ...     category_max=pxt.functions.max(tbl.value)
+        ... ).collect()
+    """
+
+    def __init__(self) -> None:
+        self.val: T = None
+
+    def update(self, val: T) -> None:
         if val is None:
             return
         if self.val is None:
             self.val = val
         else:
-            self.val = builtins.max(self.val, val)
+            self.val = builtins.max(self.val, val)  # type: ignore[call-overload]
 
-    def value(self) -> Optional[int]:
+    def value(self) -> T:
         return self.val
 
 
 @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(
-    update_types=[ts.IntType(nullable=True)], value_type=ts.FloatType(nullable=True), allows_window=False,
-    requires_order_by=False)
-class mean(func.Aggregator):
-    def __init__(self):
-        self.sum: Optional[int] = None
+@func.uda(type_substitutions=({T: int | None}, {T: float | None}))  # type: ignore[misc]
+class mean(func.Aggregator, typing.Generic[T]):
+    """
+    Aggregate function that computes the mean (average) of non-null values of a numeric column or grouping.
+
+    Args:
+        val: The numeric value to include in the mean.
+
+    Returns:
+        The mean of the non-null values, or `None` if there are no non-null values.
+
+    Examples:
+        Compute the mean of the values in the `value` column of the table `tbl`:
+
+        >>> tbl.select(pxt.functions.mean(tbl.value)).collect()
+
+        Group by the `category` column and compute the mean of the `value` column for each category,
+        assigning the name `'category_mean'` to the new column:
+
+        >>> tbl.group_by(tbl.category).select(
+        ...     tbl.category,
+        ...     category_mean=pxt.functions.mean(tbl.value)
+        ... ).collect()
+    """
+
+    def __init__(self) -> None:
+        self.sum: T = None
         self.count = 0
 
-    def update(self, val: Optional[int]) -> None:
+    def update(self, val: T) -> None:
         if val is None:
             return
         if self.sum is None:
             self.sum = val
         else:
-            self.sum += val
+            self.sum += val  # type: ignore[operator]
         self.count += 1
 
-    def value(self) -> Optional[float]:
+    def value(self) -> float | None:  # Always a float
         if self.count == 0:
             return None
-        return self.sum / self.count
+        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)
 
 
+def map(expr: exprs.Expr, fn: Callable[[exprs.Expr], Any]) -> exprs.Expr:
+    """
+    Applies a mapping function to each element of a list.
+
+    Args:
+        expr: The list expression to map over; must be an expression of type `pxt.Json`.
+        fn: An operation on Pixeltable expressions that will be applied to each element of the JSON array.
+
+    Examples:
+        Given a table `tbl` with a column `data` of type `pxt.Json` containing lists of integers, add a computed
+        column that produces new lists with each integer doubled:
+
+        >>> tbl.add_computed_column(
+        ...     doubled=pxt.functions.map(t.data, lambda x: x * 2)
+        ... )
+    """
+    target_expr: exprs.Expr
+    try:
+        target_expr = exprs.Expr.from_object(fn(exprs.json_path.RELATIVE_PATH_ROOT))
+    except Exception as e:
+        raise excs.Error(
+            'Failed to evaluate map function. '
+            '(The `fn` argument to `map()` must produce a valid Pixeltable expression.)'
+        ) from e
+    return exprs.JsonMapper(expr, target_expr)
+
+
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__