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

pixeltable/functions/document.py

@@ -0,0 +1,81 @@
+"""
+Pixeltable UDFs for `DocumentType`.
+"""
+
+from typing import Any, Literal
+
+import pixeltable as pxt
+
+
+def document_splitter(
+    document: Any,
+    separators: str,
+    *,
+    elements: list[Literal['text', 'image']] | None = None,
+    limit: int | None = None,
+    overlap: int | None = None,
+    metadata: str = '',
+    skip_tags: list[str] | None = None,
+    tiktoken_encoding: str | None = 'cl100k_base',
+    tiktoken_target_model: str | None = None,
+    image_dpi: int = 300,
+    image_format: str = 'png',
+) -> tuple[type[pxt.iterators.ComponentIterator], dict[str, Any]]:
+    """Iterator over chunks of a document. The document is chunked according to the specified `separators`.
+
+    The iterator yields a `text` field containing the text of the chunk, and it may also
+    include additional metadata fields if specified in the `metadata` parameter, as explained below.
+
+    Chunked text will be cleaned with `ftfy.fix_text` to fix up common problems with unicode sequences.
+
+    Args:
+        separators: separators to use to chunk the document. Options are:
+             `'heading'`, `'paragraph'`, `'sentence'`, `'token_limit'`, `'char_limit'`, `'page'`.
+             This may be a comma-separated string, e.g., `'heading,token_limit'`.
+        elements: list of elements to extract from the document. Options are:
+            `'text'`, `'image'`. Defaults to `['text']` if not specified. The `'image'` element is only supported
+            for the `'page'` separator on PDF documents.
+        limit: the maximum number of tokens or characters in each chunk, if `'token_limit'`
+             or `'char_limit'` is specified.
+        metadata: additional metadata fields to include in the output. Options are:
+             `'title'`, `'heading'` (HTML and Markdown), `'sourceline'` (HTML), `'page'` (PDF), `'bounding_box'`
+             (PDF). The input may be a comma-separated string, e.g., `'title,heading,sourceline'`.
+        image_dpi: DPI to use when extracting images from PDFs. Defaults to 300.
+        image_format: format to use when extracting images from PDFs. Defaults to 'png'.
+
+    Examples:
+        All these examples assume an existing table `tbl` with a column `doc` of type `pxt.Document`.
+
+        Create a view that splits all documents into chunks of up to 300 tokens:
+
+        >>> pxt.create_view('chunks', tbl, iterator=document_splitter(tbl.doc, separators='token_limit', limit=300))
+
+        Create a view that splits all documents along sentence boundaries, including title and heading metadata:
+
+        >>> pxt.create_view(
+        ...     'sentence_chunks',
+        ...     tbl,
+        ...     iterator=document_splitter(tbl.doc, separators='sentence', metadata='title,heading')
+        ... )
+    """
+
+    kwargs: dict[str, Any] = {}
+    if elements is not None:
+        kwargs['elements'] = elements
+    if limit is not None:
+        kwargs['limit'] = limit
+    if overlap is not None:
+        kwargs['overlap'] = overlap
+    if metadata != '':
+        kwargs['metadata'] = metadata
+    if skip_tags is not None:
+        kwargs['skip_tags'] = skip_tags
+    if tiktoken_encoding != 'cl100k_base':
+        kwargs['tiktoken_encoding'] = tiktoken_encoding
+    if tiktoken_target_model is not None:
+        kwargs['tiktoken_target_model'] = tiktoken_target_model
+    if image_dpi != 300:
+        kwargs['image_dpi'] = image_dpi
+    if image_format != 'png':
+        kwargs['image_format'] = image_format
+    return pxt.iterators.document.DocumentSplitter._create(document=document, separators=separators, **kwargs)

pixeltable/functions/fal.py

@@ -0,0 +1,76 @@
+"""
+Pixeltable UDFs
+that wrap various endpoints from the fal.ai API. In order to use them, you must
+first `pip install fal-client` and configure your fal.ai credentials, as described in
+the [Working with fal.ai](https://docs.pixeltable.com/notebooks/integrations/working-with-fal) tutorial.
+"""
+
+from typing import TYPE_CHECKING, Any
+
+import pixeltable as pxt
+from pixeltable.env import Env, register_client
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    import fal_client
+
+
+@register_client('fal')
+def _(api_key: str) -> 'fal_client.AsyncClient':
+    import fal_client
+
+    return fal_client.AsyncClient(key=api_key)
+
+
+def _fal_client() -> 'fal_client.AsyncClient':
+    return Env.get().get_client('fal')
+
+
+@pxt.udf(resource_pool='request-rate:fal')
+async def run(input: dict[str, Any], *, app: str) -> pxt.Json:
+    """
+    Run a model on fal.ai.
+
+    Uses fal's queue-based subscribe mechanism for reliable execution.
+    For additional details, see: <https://fal.ai/docs>
+
+    Request throttling:
+    Applies the rate limit set in the config (section `fal`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
+    __Requirements:__
+
+    - `pip install fal-client`
+
+    Args:
+        input: The input parameters for the model.
+        app: The name or ID of the fal.ai application to run (e.g., 'fal-ai/flux/schnell').
+
+    Returns:
+        The output of the model as a JSON object.
+
+    Examples:
+        Add a computed column that applies the model `fal-ai/flux/schnell`
+        to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
+
+        >>> input = {'prompt': tbl.prompt}
+        ... tbl.add_computed_column(response=run(input, app='fal-ai/flux/schnell'))
+
+        Add a computed column that uses the model `fal-ai/fast-sdxl`
+        to generate images from an existing Pixeltable column `tbl.prompt`:
+
+        >>> input = {'prompt': tbl.prompt, 'image_size': 'square', 'num_inference_steps': 25}
+        ... tbl.add_computed_column(response=run(input, app='fal-ai/fast-sdxl'))
+        ... tbl.add_computed_column(image=tbl.response['images'][0]['url'].astype(pxt.Image))
+    """
+    Env.get().require_package('fal_client')
+    client = _fal_client()
+    result = await client.subscribe(app, arguments=input)
+    return result
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/functions/fireworks.py

@@ -1,11 +1,11 @@
 """
-Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
+Pixeltable UDFs
 that wrap various endpoints from the Fireworks AI API. In order to use them, you must
 first `pip install fireworks-ai` and configure your Fireworks AI credentials, as described in
-the [Working with Fireworks](https://pixeltable.readme.io/docs/working-with-fireworks) tutorial.
+the [Working with Fireworks](https://docs.pixeltable.com/notebooks/integrations/working-with-fireworks) tutorial.
 """
 
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any
 
 import pixeltable as pxt
 from pixeltable import env
@@ -29,14 +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,
-    max_tokens: Optional[int] = None,
-    top_k: Optional[int] = None,
-    top_p: Optional[float] = None,
-    temperature: Optional[float] = None,
-    request_timeout: Optional[int] = 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.
@@ -55,8 +48,8 @@ async def chat_completions(
     Args:
         messages: A list of messages comprising the conversation so far.
         model: The name of the model to use.
-
-    For details on the other parameters, see: <https://docs.fireworks.ai/api-reference/post-chatcompletions>
+        model_kwargs: Additional keyword args for the Fireworks `chat_completions` API. For details on the available
+            parameters, see: <https://docs.fireworks.ai/api-reference/post-chatcompletions>
 
     Returns:
         A dictionary containing the response and other metadata.
@@ -70,20 +63,18 @@ async def chat_completions(
         ...     response=chat_completions(messages, model='accounts/fireworks/models/mixtral-8x22b-instruct')
         ... )
     """
-    kwargs = {'max_tokens': max_tokens, 'top_k': top_k, 'top_p': top_p, 'temperature': temperature}
-    kwargs_not_none = {k: v for k, v in kwargs.items() if v is not None}
+    if model_kwargs is None:
+        model_kwargs = {}
 
     # for debugging purposes:
     # res_sync = _fireworks_client().chat.completions.create(model=model, messages=messages, **kwargs_not_none)
     # res_sync_dict = res_sync.dict()
 
-    if request_timeout is None:
-        request_timeout = Config.get().get_int_value('timeout', section='fireworks') or 600
+    if 'request_timeout' not in model_kwargs:
+        model_kwargs['request_timeout'] = Config.get().get_int_value('timeout', section='fireworks') or 600
     # TODO: this timeout doesn't really work, I think it only applies to returning the stream, but not to the timing
     # of the chunks; addressing this would require a timeout for the task running this udf
-    stream = _fireworks_client().chat.completions.acreate(
-        model=model, messages=messages, request_timeout=request_timeout, **kwargs_not_none
-    )
+    stream = _fireworks_client().chat.completions.acreate(model=model, messages=messages, **model_kwargs)
     chunks = []
     async for chunk in stream:
         chunks.append(chunk)

pixeltable/functions/gemini.py

@@ -1,14 +1,21 @@
 """
-Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
+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://pixeltable.readme.io/docs/working-with-gemini) tutorial.
+the [Working with Gemini](https://docs.pixeltable.com/notebooks/integrations/working-with-gemini) tutorial.
 """
 
-from typing import TYPE_CHECKING, Optional
+import asyncio
+import io
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import PIL.Image
 
 import pixeltable as pxt
-from pixeltable import env
+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
@@ -27,26 +34,13 @@ def _genai_client() -> 'genai.client.Client':
 
 @pxt.udf(resource_pool='request-rate:gemini')
 async def generate_content(
-    contents: str,
-    *,
-    model: 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: pxt.Json, *, 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>
+    Generate content from the specified model.
 
     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:__
@@ -54,40 +48,202 @@ async def generate_content(
     - `pip install google-genai`
 
     Args:
-        contents: The input content to generate from.
+        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.
-
-    For details on the other parameters, see: <https://ai.google.dev/gemini-api/docs>
+        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-1.5-flash`
+        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-1.5-flash'))
+        >>> 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(
-        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 _genai_client().aio.models.generate_content(model=model, contents=contents, config=config)
+    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__