pixeltable 0.3.13__py3-none-any.whl → 0.3.15__py3-none-any.whl

pixeltable/exprs/literal.py

@@ -50,6 +50,9 @@ class Literal(Expr):
             assert isinstance(self.val, datetime.datetime)
             default_tz = Env.get().default_time_zone
             return f"'{self.val.astimezone(default_tz).isoformat()}'"
+        if self.col_type.is_date_type():
+            assert isinstance(self.val, datetime.date)
+            return f"'{self.val.isoformat()}'"
         if self.col_type.is_array_type():
             assert isinstance(self.val, np.ndarray)
             return str(self.val.tolist())
@@ -82,6 +85,10 @@ class Literal(Expr):
             # stored as UTC in the database)
             encoded_val = self.val.isoformat()
             return {'val': encoded_val, 'val_t': self.col_type._type.name, **super()._as_dict()}
+        elif self.col_type.is_date_type():
+            assert isinstance(self.val, datetime.date)
+            encoded_val = self.val.isoformat()
+            return {'val': encoded_val, 'val_t': self.col_type._type.name, **super()._as_dict()}
         elif self.col_type.is_array_type():
             assert isinstance(self.val, np.ndarray)
             return {'val': self.val.tolist(), 'val_t': self.col_type._type.name, **super()._as_dict()}
@@ -96,7 +103,10 @@ class Literal(Expr):
         assert 'val' in d
         if 'val_t' in d:
             val_t = d['val_t']
-            if val_t == ts.ColumnType.Type.TIMESTAMP.name:
+            if val_t == ts.ColumnType.Type.DATE.name:
+                dt = datetime.date.fromisoformat(d['val'])
+                return cls(dt)
+            elif val_t == ts.ColumnType.Type.TIMESTAMP.name:
                 dt = datetime.datetime.fromisoformat(d['val'])
                 assert dt.tzinfo == datetime.timezone.utc  # Must be UTC in the database
                 return cls(dt)

pixeltable/exprs/rowid_ref.py

@@ -30,7 +30,7 @@ class RowidRef(Expr):
 
     def __init__(
         self,
-        tbl: catalog.TableVersionHandle,
+        tbl: Optional[catalog.TableVersionHandle],
         idx: int,
         tbl_id: Optional[UUID] = None,
         normalized_base_id: Optional[UUID] = None,
@@ -98,6 +98,9 @@ class RowidRef(Expr):
     def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
         tbl = self.tbl.get() if self.tbl is not None else catalog.Catalog.get().get_tbl_version(self.tbl_id, None)
         rowid_cols = tbl.store_tbl.rowid_columns()
+        assert self.rowid_component_idx <= len(rowid_cols), (
+            f'{self.rowid_component_idx} not consistent with {rowid_cols}'
+        )
         return rowid_cols[self.rowid_component_idx]
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:

pixeltable/exprs/similarity_expr.py

@@ -26,7 +26,7 @@ class SimilarityExpr(Expr):
         from pixeltable import index
 
         # determine index to use
-        idx_dict = ColumnRef.find_embedding_index(col_ref.col, idx_name, 'similarity')
+        idx_dict = col_ref.find_embedding_index(idx_name, 'similarity')
         assert len(idx_dict) == 1
         self.idx_info = next(iter(idx_dict.values()))
         idx = self.idx_info.idx

pixeltable/func/function.py

@@ -514,7 +514,7 @@ class InvalidFunction(Function):
     def _as_dict(self) -> dict:
         """
         Here we write out (verbatim) the original metadata that failed to load (and that resulted in the
-        InvalidFunction). Note that the InvalidFunction itself is never serlialized, so there is no corresponding
+        InvalidFunction). Note that the InvalidFunction itself is never serialized, so there is no corresponding
         from_dict() method.
         """
         return self.fn_dict

pixeltable/functions/__init__.py

@@ -6,6 +6,7 @@ from . import (
     anthropic,
     audio,
     bedrock,
+    date,
     deepseek,
     fireworks,
     gemini,

pixeltable/functions/date.py

@@ -0,0 +1,185 @@
+"""
+Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs) for `DateType`.
+
+Usage example:
+```python
+import pixeltable as pxt
+
+t = pxt.get_table(...)
+t.select(t.date_col.year, t.date_col.weekday()).collect()
+```
+"""
+
+from datetime import date, timedelta
+
+import sqlalchemy as sql
+
+import pixeltable as pxt
+from pixeltable.utils.code import local_public_names
+
+_SQL_ZERO = sql.literal(0)
+
+# NOT YET SUPPORTED date +/- integer
+# NOT YET SUPPORTED date1 - date2 -> integer
+# NOT YET SUPPORTED timestamp(date)
+# NOT YET SUPPORTED date(timestamp)
+
+
+@pxt.udf(is_property=True)
+def year(self: date) -> int:
+    """
+    Between [`MINYEAR`](https://docs.python.org/3/library/datetime.html#datetime.MINYEAR) and
+    [`MAXYEAR`](https://docs.python.org/3/library/datetime.html#datetime.MAXYEAR) inclusive.
+
+    Equivalent to [`date.year`](https://docs.python.org/3/library/datetime.html#datetime.date.year).
+    """
+    return self.year
+
+
+@year.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.extract('year', self)
+
+
+@pxt.udf(is_property=True)
+def month(self: date) -> int:
+    """
+    Between 1 and 12 inclusive.
+
+    Equivalent to [`date.month`](https://docs.python.org/3/library/datetime.html#datetime.date.month).
+    """
+    return self.month
+
+
+@month.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.extract('month', self)
+
+
+@pxt.udf(is_property=True)
+def day(self: date) -> int:
+    """
+    Between 1 and the number of days in the given month of the given year.
+
+    Equivalent to [`date.day`](https://docs.python.org/3/library/datetime.html#datetime.date.day).
+    """
+    return self.day
+
+
+@day.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.extract('day', self)
+
+
+@pxt.udf(is_method=True)
+def make_date(year: int, month: int, day: int) -> date:
+    """
+    Create a date.
+
+    Equivalent to [`datetime()`](https://docs.python.org/3/library/datetime.html#datetime.date).
+    """
+    return date(year, month, day)
+
+
+@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))
+
+
+@pxt.udf(is_method=True)
+def weekday(self: date) -> int:
+    """
+    Between 0 (Monday) and 6 (Sunday) inclusive.
+
+    Equivalent to [`date.weekday()`](https://docs.python.org/3/library/datetime.html#datetime.date.weekday).
+    """
+    return self.weekday()
+
+
+@weekday.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.extract('isodow', self) - 1
+
+
+@pxt.udf(is_method=True)
+def isoweekday(self: date) -> int:
+    """
+    Return the day of the week as an integer, where Monday is 1 and Sunday is 7.
+
+    Equivalent to [`date.isoweekday()`](https://docs.python.org/3/library/datetime.html#datetime.date.isoweekday).
+    """
+    return self.isoweekday()
+
+
+@isoweekday.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.extract('isodow', self)
+
+
+@pxt.udf(is_method=True)
+def isocalendar(self: date) -> dict:
+    """
+    Return a dictionary with three entries: `'year'`, `'week'`, and `'weekday'`.
+
+    Equivalent to
+    [`date.isocalendar()`](https://docs.python.org/3/library/datetime.html#datetime.date.isocalendar).
+    """
+    iso_year, iso_week, iso_weekday = self.isocalendar()
+    return {'year': iso_year, 'week': iso_week, 'weekday': iso_weekday}
+
+
+@pxt.udf(is_method=True)
+def isoformat(self: date, sep: str = 'T', timespec: str = 'auto') -> str:
+    """
+    Return a string representing the date and time in ISO 8601 format.
+
+    Equivalent to [`date.isoformat()`](https://docs.python.org/3/library/datetime.html#datetime.date.isoformat).
+
+    Args:
+        sep: Separator between date and time.
+        timespec: The number of additional terms in the output. See the
+            [`date.isoformat()`](https://docs.python.org/3/library/datetime.html#datetime.date.isoformat)
+            documentation for more details.
+    """
+    return self.isoformat()
+
+
+@pxt.udf(is_method=True)
+def toordinal(self: date) -> int:
+    """
+    Return the proleptic Gregorian ordinal of the date, where January 1 of year 1 has ordinal 1.
+
+    Equivalent to [`date.toordinal()`](https://docs.python.org/3/library/datetime.html#datetime.date.toordinal).
+    """
+    return self.toordinal()
+
+
+@pxt.udf(is_method=True)
+def strftime(self: date, format: str) -> str:
+    """
+    Return a string representing the date and time, controlled by an explicit format string.
+
+    Equivalent to [`date.strftime()`](https://docs.python.org/3/library/datetime.html#datetime.date.strftime).
+
+    Args:
+        format: The format string to control the output. For a complete list of formatting directives, see
+            [`strftime()` and `strptime()` Behavior](https://docs.python.org/3/library/datetime.html#strftime-and-strptime-behavior).
+    """
+    return self.strftime(format)
+
+
+@pxt.udf(is_method=True)
+def add_days(self: date, n: int) -> date:
+    """
+    Add `n` days to the date.
+
+    Equivalent to [`date + timedelta(days=n)`](https://docs.python.org/3/library/datetime.html#datetime.timedelta).
+    """
+    return self + timedelta(days=n)
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/functions/gemini.py

@@ -1,46 +1,43 @@
 """
 Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
 that wrap various endpoints from the Google Gemini API. In order to use them, you must
-first `pip install google-generativeai` and configure your Gemini credentials, as described in
+first `pip install google-genai` and configure your Gemini credentials, as described in
 the [Working with Gemini](https://pixeltable.readme.io/docs/working-with-gemini) tutorial.
 """
 
-from typing import Optional
+import asyncio
+import io
+import tempfile
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional
+
+import PIL.Image
 
 import pixeltable as pxt
-from pixeltable import env
+from pixeltable import env, exceptions as excs, exprs
+
+if TYPE_CHECKING:
+    from google import genai
 
 
 @env.register_client('gemini')
-def _(api_key: str) -> None:
-    import google.generativeai as genai
+def _(api_key: str) -> 'genai.client.Client':
+    from google import genai
 
-    genai.configure(api_key=api_key)
+    return genai.client.Client(api_key=api_key)
 
 
-def _ensure_loaded() -> None:
-    env.Env.get().get_client('gemini')
+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: str,
-    *,
-    model_name: str,
-    candidate_count: Optional[int] = None,
-    stop_sequences: Optional[list[str]] = None,
-    max_output_tokens: Optional[int] = None,
-    temperature: Optional[float] = None,
-    top_p: Optional[float] = None,
-    top_k: Optional[int] = None,
-    response_mime_type: Optional[str] = None,
-    response_schema: Optional[dict] = None,
-    presence_penalty: Optional[float] = None,
-    frequency_penalty: Optional[float] = None,
+    contents: str, *, model: str, config: Optional[dict] = None, tools: Optional[list[dict]] = None
 ) -> dict:
     """
     Generate content from the specified model. For additional details, see:
-    <https://ai.google.dev/gemini-api/docs>
+    <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
@@ -48,44 +45,182 @@ async def generate_content(
 
     __Requirements:__
 
-    - `pip install google-generativeai`
+    - `pip install google-genai`
 
     Args:
         contents: The input content to generate from.
-        model_name: The name of the model to use.
-
-    For details on the other parameters, see: <https://ai.google.dev/gemini-api/docs>
+        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#module-genai.types>
+        tools: 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-1.5-flash`
+        Add a computed column that applies the model `gemini-2.0-flash`
         to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
 
-        >>> tbl.add_computed_column(response=generate_content(tbl.prompt, model_name='gemini-1.5-flash'))
+        >>> tbl.add_computed_column(response=generate_content(tbl.prompt, model='gemini-2.0-flash'))
     """
-    env.Env.get().require_package('google.generativeai')
-    _ensure_loaded()
-    import google.generativeai as genai
-
-    model = genai.GenerativeModel(model_name=model_name)
-    gc = genai.GenerationConfig(
-        candidate_count=candidate_count,
-        stop_sequences=stop_sequences,
-        max_output_tokens=max_output_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        top_k=top_k,
-        response_mime_type=response_mime_type,
-        response_schema=response_schema,
-        presence_penalty=presence_penalty,
-        frequency_penalty=frequency_penalty,
-    )
-    response = await model.generate_content_async(contents, generation_config=gc)
-    return response.to_dict()
+    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) -> Optional[dict]:
+    print(response)
+    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_name: str) -> str:
-    return f'request-rate:gemini:{model_name}'
+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: Optional[dict] = 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>
+
+    __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#module-genai.types>
+
+    Returns:
+        The generated image.
+
+    Examples:
+        Add a computed column that applies the model `imagen-3.0-generate-002`
+        to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
+
+        >>> tbl.add_computed_column(response=generate_images(tbl.prompt, model='imagen-3.0-generate-002'))
+    """
+    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: Optional[str] = None, image: Optional[PIL.Image.Image] = None, *, model: str, config: Optional[dict] = 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>
+
+    __Requirements:__
+
+    - `pip install google-genai`
+
+    Args:
+        prompt: A text description of the videos to generate.
+        image: An optional image to use as the first frame of the video. At least one of `prompt` or `image` must be
+            provided. (It is ok to specify both.)
+        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#module-genai.types>
+
+    Returns:
+        The generated video.
+
+    Examples:
+        Add a computed column that applies the model `veo-2.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-2.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_: Optional[types.Image] = None
+    if image is not None:
+        with io.BytesIO() as buffer:
+            image.save(buffer, format='jpeg')
+            image_ = types.Image(image_bytes=buffer.getvalue(), mime_type='image/jpeg')
+
+    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)
+
+    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
+
+    _, output_filename = tempfile.mkstemp(suffix='.mp4', dir=str(env.Env.get().tmp_dir))
+    Path(output_filename).write_bytes(video_bytes)
+    return output_filename
+
+
+@generate_videos.resource_pool
+def _(model: str) -> str:
+    return f'request-rate:veo:{model}'

pixeltable/functions/globals.py

@@ -49,22 +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]}  # type: ignore[misc]
-        for t in (
-            ts.String,
-            ts.Int,
-            ts.Float,
-            ts.Bool,
-            ts.Timestamp,
-            ts.Array,
-            ts.Json,
-            ts.Image,
-            ts.Video,
-            ts.Audio,
-            ts.Document,
-        )
-    ),
+    type_substitutions=tuple({T: Optional[t]} for t in ts.ALL_PIXELTABLE_TYPES),  # type: ignore[misc]
 )
 class count(func.Aggregator, typing.Generic[T]):
     def __init__(self) -> None:

pixeltable/functions/json.py

@@ -4,9 +4,10 @@ Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
 Example:
 ```python
 import pixeltable as pxt
+import pixeltable.functions as pxtf
 
 t = pxt.get_table(...)
-t.select(pxt.functions.json.make_list()).collect()
+t.select(pxtf.json.make_list(t.json_col)).collect()
 ```
 """