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

pixeltable/func/query_template_function.py

@@ -157,12 +157,14 @@ def retrieval_udf(
     """
     # Argument validation
     col_refs: list[exprs.ColumnRef]
+    # TODO: get rid of references to ColumnRef internals and replace instead with a public interface
+    col_names = table.columns()
     if parameters is None:
-        col_refs = [table[col_name] for col_name in table.columns if not table[col_name].col.is_computed]
+        col_refs = [table[col_name] for col_name in col_names if not table[col_name].col.is_computed]
     else:
         for param in parameters:
-            if isinstance(param, str) and param not in table.columns:
-                raise excs.Error(f'The specified parameter {param!r} is not a column of the table {table._path!r}')
+            if isinstance(param, str) and param not in col_names:
+                raise excs.Error(f'The specified parameter {param!r} is not a column of the table {table._path()!r}')
         col_refs = [table[param] if isinstance(param, str) else param for param in parameters]
 
     if len(col_refs) == 0:

pixeltable/func/tools.py

@@ -1,8 +1,9 @@
+import json
 from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Union
 
 import pydantic
 
-import pixeltable.exceptions as excs
+from pixeltable import exceptions as excs, type_system as ts
 
 from .function import Function
 from .signature import Parameter
@@ -69,7 +70,9 @@ class Tool(pydantic.BaseModel):
             return _extract_float_tool_arg(kwargs, param_name=param.name)
         if param.col_type.is_bool_type():
             return _extract_bool_tool_arg(kwargs, param_name=param.name)
-        raise AssertionError()
+        if param.col_type.is_json_type():
+            return _extract_json_tool_arg(kwargs, param_name=param.name)
+        raise AssertionError(param.col_type)
 
 
 class ToolChoice(pydantic.BaseModel):
@@ -137,6 +140,13 @@ def _extract_bool_tool_arg(kwargs: dict[str, Any], param_name: str) -> Optional[
     return _extract_arg(bool, kwargs, param_name)
 
 
+@udf
+def _extract_json_tool_arg(kwargs: dict[str, Any], param_name: str) -> Optional[ts.Json]:
+    if param_name in kwargs:
+        return json.loads(kwargs[param_name])
+    return None
+
+
 T = TypeVar('T')
 
 

pixeltable/func/udf.py

@@ -262,7 +262,7 @@ def from_table(
     """
     from pixeltable import exprs
 
-    ancestors = [tbl, *tbl._base_tables]
+    ancestors = [tbl, *tbl._get_base_tables()]
     ancestors.reverse()  # We must traverse the ancestors in order from base to derived
 
     subst: dict[exprs.Expr, exprs.Expr] = {}
@@ -297,7 +297,7 @@ def from_table(
 
     if description is None:
         # Default description is the table comment
-        description = tbl._comment
+        description = tbl._get_comment()
         if len(description) == 0:
             description = f"UDF for table '{tbl._name}'"
 

pixeltable/functions/__init__.py

@@ -10,6 +10,7 @@ from . import (
     deepseek,
     fireworks,
     gemini,
+    groq,
     huggingface,
     image,
     json,

pixeltable/functions/anthropic.py

@@ -8,7 +8,7 @@ the [Working with Anthropic](https://pixeltable.readme.io/docs/working-with-anth
 import datetime
 import json
 import logging
-from typing import TYPE_CHECKING, Any, Iterable, Optional, TypeVar, Union, cast
+from typing import TYPE_CHECKING, Any, Iterable, Optional, cast
 
 import httpx
 
@@ -73,16 +73,10 @@ async def messages(
     messages: list[dict[str, str]],
     *,
     model: str,
-    max_tokens: int = 1024,
-    metadata: Optional[dict[str, Any]] = None,
-    stop_sequences: Optional[list[str]] = None,
-    system: Optional[str] = None,
-    temperature: Optional[float] = None,
-    tool_choice: Optional[dict] = None,
-    tools: Optional[list[dict]] = None,
-    top_k: Optional[int] = None,
-    top_p: Optional[float] = None,
-    timeout: Optional[float] = None,
+    max_tokens: int,
+    model_kwargs: Optional[dict[str, Any]] = None,
+    tools: Optional[list[dict[str, Any]]] = None,
+    tool_choice: Optional[dict[str, Any]] = None,
 ) -> dict:
     """
     Create a Message.
@@ -101,25 +95,27 @@ async def messages(
     Args:
         messages: Input messages.
         model: The model that will complete your prompt.
-
-    For details on the other parameters, see: <https://docs.anthropic.com/en/api/messages>
+        model_kwargs: Additional keyword args for the Anthropic `messages` API.
+            For details on the available parameters, see: <https://docs.anthropic.com/en/api/messages>
+        tools: An optional list of Pixeltable tools to use for the request.
+        tool_choice: An optional tool choice configuration.
 
     Returns:
         A dictionary containing the response and other metadata.
 
     Examples:
-        Add a computed column that applies the model `claude-3-haiku-20240307`
+        Add a computed column that applies the model `claude-3-5-sonnet-20241022`
         to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
 
         >>> msgs = [{'role': 'user', 'content': tbl.prompt}]
-        ... tbl.add_computed_column(response=messages(msgs, model='claude-3-haiku-20240307'))
+        ... tbl.add_computed_column(response=messages(msgs, model='claude-3-5-sonnet-20241022'))
     """
-
-    # it doesn't look like count_tokens() actually exists in the current version of the library
+    if model_kwargs is None:
+        model_kwargs = {}
 
     if tools is not None:
         # Reformat `tools` into Anthropic format
-        tools = [
+        model_kwargs['tools'] = [
             {
                 'name': tool['name'],
                 'description': tool['description'],
@@ -132,17 +128,16 @@ async def messages(
             for tool in tools
         ]
 
-    tool_choice_: Optional[dict] = None
     if tool_choice is not None:
         if tool_choice['auto']:
-            tool_choice_ = {'type': 'auto'}
+            model_kwargs['tool_choice'] = {'type': 'auto'}
         elif tool_choice['required']:
-            tool_choice_ = {'type': 'any'}
+            model_kwargs['tool_choice'] = {'type': 'any'}
         else:
             assert tool_choice['tool'] is not None
-            tool_choice_ = {'type': 'tool', 'name': tool_choice['tool']}
+            model_kwargs['tool_choice'] = {'type': 'tool', 'name': tool_choice['tool']}
         if not tool_choice['parallel_tool_calls']:
-            tool_choice_['disable_parallel_tool_use'] = True
+            model_kwargs['tool_choice']['disable_parallel_tool_use'] = True
 
     # make sure the pool info exists prior to making the request
     resource_pool_id = f'rate-limits:anthropic:{model}'
@@ -152,20 +147,8 @@ async def messages(
     # TODO: timeouts should be set system-wide and be user-configurable
     from anthropic.types import MessageParam
 
-    # cast(Any, ...): avoid mypy errors
     result = await _anthropic_client().messages.with_raw_response.create(
-        messages=cast(Iterable[MessageParam], messages),
-        model=model,
-        max_tokens=max_tokens,
-        metadata=_opt(cast(Any, metadata)),
-        stop_sequences=_opt(stop_sequences),
-        system=_opt(system),
-        temperature=_opt(cast(Any, temperature)),
-        tools=_opt(cast(Any, tools)),
-        tool_choice=_opt(cast(Any, tool_choice_)),
-        top_k=_opt(top_k),
-        top_p=_opt(top_p),
-        timeout=_opt(timeout),
+        messages=cast(Iterable[MessageParam], messages), model=model, max_tokens=max_tokens, **model_kwargs
     )
 
     requests_limit_str = result.headers.get('anthropic-ratelimit-requests-limit')
@@ -224,15 +207,6 @@ def _anthropic_response_to_pxt_tool_calls(response: dict) -> Optional[dict]:
     return pxt_tool_calls
 
 
-_T = TypeVar('_T')
-
-
-def _opt(arg: _T) -> Union[_T, 'anthropic.NotGiven']:
-    import anthropic
-
-    return arg if arg is not None else anthropic.NOT_GIVEN
-
-
 __all__ = local_public_names(__name__)
 
 

pixeltable/functions/deepseek.py

@@ -1,5 +1,5 @@
 import json
-from typing import TYPE_CHECKING, Any, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Optional
 
 import httpx
 
@@ -7,8 +7,6 @@ import pixeltable as pxt
 from pixeltable import env
 from pixeltable.utils.code import local_public_names
 
-from .openai import _opt
-
 if TYPE_CHECKING:
     import openai
 
@@ -33,17 +31,9 @@ async def chat_completions(
     messages: list,
     *,
     model: str,
-    frequency_penalty: Optional[float] = None,
-    logprobs: Optional[bool] = None,
-    top_logprobs: Optional[int] = None,
-    max_tokens: Optional[int] = None,
-    presence_penalty: Optional[float] = None,
-    response_format: Optional[dict] = None,
-    stop: Optional[list[str]] = None,
-    temperature: Optional[float] = None,
-    tools: Optional[list[dict]] = None,
-    tool_choice: Optional[dict] = None,
-    top_p: Optional[float] = None,
+    model_kwargs: Optional[dict[str, Any]] = None,
+    tools: Optional[list[dict[str, Any]]] = None,
+    tool_choice: Optional[dict[str, Any]] = None,
 ) -> dict:
     """
     Creates a model response for the given chat conversation.
@@ -60,8 +50,10 @@ async def chat_completions(
     Args:
         messages: A list of messages to use for chat completion, as described in the Deepseek API documentation.
         model: The model to use for chat completion.
-
-    For details on the other parameters, see: <https://api-docs.deepseek.com/api/create-chat-completion>
+        model_kwargs: Additional keyword args for the Deepseek `chat/completions` API.
+            For details on the available parameters, see: <https://api-docs.deepseek.com/api/create-chat-completion>
+        tools: An optional list of Pixeltable tools to use for the request.
+        tool_choice: An optional tool choice configuration.
 
     Returns:
         A dictionary containing the response and other metadata.
@@ -76,39 +68,28 @@ async def chat_completions(
             ]
             tbl.add_computed_column(response=chat_completions(messages, model='deepseek-chat'))
     """
+    if model_kwargs is None:
+        model_kwargs = {}
+
     if tools is not None:
-        tools = [{'type': 'function', 'function': tool} for tool in tools]
+        model_kwargs['tools'] = [{'type': 'function', 'function': tool} for tool in tools]
 
-    tool_choice_: Union[str, dict, None] = None
     if tool_choice is not None:
         if tool_choice['auto']:
-            tool_choice_ = 'auto'
+            model_kwargs['tool_choice'] = 'auto'
         elif tool_choice['required']:
-            tool_choice_ = 'required'
+            model_kwargs['tool_choice'] = 'required'
         else:
             assert tool_choice['tool'] is not None
-            tool_choice_ = {'type': 'function', 'function': {'name': tool_choice['tool']}}
+            model_kwargs['tool_choice'] = {'type': 'function', 'function': {'name': tool_choice['tool']}}
 
-    extra_body: Optional[dict[str, Any]] = None
     if tool_choice is not None and not tool_choice['parallel_tool_calls']:
-        extra_body = {'parallel_tool_calls': False}
+        if 'extra_body' not in model_kwargs:
+            model_kwargs['extra_body'] = {}
+        model_kwargs['extra_body']['parallel_tool_calls'] = False
 
-    # cast(Any, ...): avoid mypy errors
     result = await _deepseek_client().chat.completions.with_raw_response.create(
-        messages=messages,
-        model=model,
-        frequency_penalty=_opt(frequency_penalty),
-        logprobs=_opt(logprobs),
-        top_logprobs=_opt(top_logprobs),
-        max_tokens=_opt(max_tokens),
-        presence_penalty=_opt(presence_penalty),
-        response_format=_opt(cast(Any, response_format)),
-        stop=_opt(stop),
-        temperature=_opt(temperature),
-        tools=_opt(cast(Any, tools)),
-        tool_choice=_opt(cast(Any, tool_choice_)),
-        top_p=_opt(top_p),
-        extra_body=extra_body,
+        messages=messages, model=model, **model_kwargs
     )
 
     return json.loads(result.text)

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, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 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: Optional[dict[str, Any]] = 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

@@ -5,10 +5,16 @@ first `pip install google-genai` and configure your Gemini credentials, as descr
 the [Working with Gemini](https://pixeltable.readme.io/docs/working-with-gemini) tutorial.
 """
 
+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
@@ -27,23 +33,11 @@ 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: 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
@@ -56,38 +50,176 @@ async def generate_content(
     Args:
         contents: The input content to generate from.
         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#module-genai.types>
+        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.0-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.0-flash'))
     """
     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) -> Optional[dict]:
+    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: 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}'