pixeltable 0.4.0rc1__py3-none-any.whl → 0.4.0rc3__py3-none-any.whl

pixeltable/exec/sql_node.py

@@ -14,6 +14,7 @@ from .exec_node import ExecNode
 
 if TYPE_CHECKING:
     import pixeltable.plan
+    from pixeltable.plan import SampleClause
 
 _logger = logging.getLogger('pixeltable')
 
@@ -64,8 +65,12 @@ def print_order_by_clause(clause: OrderByClause) -> str:
 
 class SqlNode(ExecNode):
     """
-    Materializes data from the store via a Select stmt.
+    Materializes data from the store via an SQL statement.
     This only provides the select list. The subclasses are responsible for the From clause and any additional clauses.
+    The pk columns are not included in the select list.
+    If set_pk is True, they are added to the end of the result set when creating the SQL statement
+    so they can always be referenced as cols[-num_pk_cols:] in the result set.
+    The pk_columns consist of the rowid columns of the target table followed by the version number.
     """
 
     tbl: Optional[catalog.TableVersionPath]
@@ -122,6 +127,7 @@ class SqlNode(ExecNode):
             # we also need to retrieve the pk columns
             assert tbl is not None
             self.num_pk_cols = len(tbl.tbl_version.get().store_tbl.pk_columns())
+            assert self.num_pk_cols > 1
 
         # additional state
         self.result_cursor = None
@@ -139,15 +145,20 @@ class SqlNode(ExecNode):
             if tv is not None:
                 assert tv.is_validated
 
+    def _create_pk_cols(self) -> list[sql.Column]:
+        """Create a list of pk columns"""
+        # we need to retrieve the pk columns
+        if self.set_pk:
+            assert self.tbl is not None
+            assert self.tbl.tbl_version.get().is_validated
+            return self.tbl.tbl_version.get().store_tbl.pk_columns()
+        return []
+
     def _create_stmt(self) -> sql.Select:
         """Create Select from local state"""
 
         assert self.sql_elements.contains_all(self.select_list)
-        sql_select_list = [self.sql_elements.get(e) for e in self.select_list]
-        if self.set_pk:
-            assert self.tbl is not None
-            assert self.tbl.tbl_version.get().is_validated
-            sql_select_list += self.tbl.tbl_version.get().store_tbl.pk_columns()
+        sql_select_list = [self.sql_elements.get(e) for e in self.select_list] + self._create_pk_cols()
         stmt = sql.select(*sql_select_list)
 
         where_clause_element = (
@@ -173,9 +184,10 @@ class SqlNode(ExecNode):
     def _ordering_tbl_ids(self) -> set[UUID]:
         return exprs.Expr.all_tbl_ids(e for e, _ in self.order_by_clause)
 
-    def to_cte(self) -> Optional[tuple[sql.CTE, exprs.ExprDict[sql.ColumnElement]]]:
+    def to_cte(self, keep_pk: bool = False) -> Optional[tuple[sql.CTE, exprs.ExprDict[sql.ColumnElement]]]:
         """
-        Returns a CTE that materializes the output of this node plus a mapping from select list expr to output column
+        Creates a CTE that materializes the output of this node plus a mapping from select list expr to output column.
+        keep_pk: if True, the PK columns are included in the CTE Select statement
 
         Returns:
             (CTE, dict from Expr to output column)
@@ -183,11 +195,13 @@ class SqlNode(ExecNode):
         if self.py_filter is not None:
             # the filter needs to run in Python
             return None
-        self.set_pk = False  # we don't need the PK if we use this SqlNode as a CTE
         if self.cte is None:
+            if not keep_pk:
+                self.set_pk = False  # we don't need the PK if we use this SqlNode as a CTE
             self.cte = self._create_stmt().cte()
-            assert len(self.cte.c) == len(self.select_list)
-        return self.cte, exprs.ExprDict(zip(self.select_list, self.cte.c))
+        pk_count = self.num_pk_cols if self.set_pk else 0
+        assert len(self.select_list) + pk_count == len(self.cte.c)
+        return self.cte, exprs.ExprDict(zip(self.select_list, self.cte.c))  # skip pk cols
 
     @classmethod
     def retarget_rowid_refs(cls, target: catalog.TableVersionPath, expr_seq: Iterable[exprs.Expr]) -> None:
@@ -293,7 +307,9 @@ class SqlNode(ExecNode):
                 stmt_str = str(stmt.compile(compile_kwargs={'literal_binds': True}))
                 _logger.debug(f'SqlLookupNode stmt:\n{stmt_str}')
             except Exception:
-                pass
+                # log something if we can't log the compiled stmt
+                stmt_str = repr(stmt)
+                _logger.debug(f'SqlLookupNode proto-stmt:\n{stmt_str}')
             self._log_explain(stmt)
 
             conn = Env.get().conn
@@ -510,3 +526,127 @@ class SqlJoinNode(SqlNode):
                 full=join_clause == plan.JoinType.FULL_OUTER,
             )
         return stmt
+
+
+class SqlSampleNode(SqlNode):
+    """
+    Returns rows from a stratified sample with N samples per strata.
+    """
+
+    stratify_exprs: Optional[list[exprs.Expr]]
+    n_samples: Optional[int]
+    fraction_samples: Optional[float]
+    seed: int
+    input_cte: Optional[sql.CTE]
+    pk_count: int
+
+    def __init__(
+        self,
+        row_builder: exprs.RowBuilder,
+        input: SqlNode,
+        select_list: Iterable[exprs.Expr],
+        stratify_exprs: Optional[list[exprs.Expr]] = None,
+        sample_clause: Optional['SampleClause'] = None,
+    ):
+        """
+        Args:
+            select_list: can contain calls to AggregateFunctions
+            stratify_exprs: list of expressions to group by
+            n: number of samples per strata
+        """
+        self.input_cte, input_col_map = input.to_cte(keep_pk=True)
+        self.pk_count = input.num_pk_cols
+        assert self.pk_count > 1
+        sql_elements = exprs.SqlElementCache(input_col_map)
+        super().__init__(input.tbl, row_builder, select_list, sql_elements, set_pk=True)
+        self.stratify_exprs = stratify_exprs
+        self.n_samples = sample_clause.n
+        self.n_per_stratum = sample_clause.n_per_stratum
+        self.fraction_samples = sample_clause.fraction
+        self.seed = sample_clause.seed if sample_clause.seed is not None else 0
+
+    @classmethod
+    def key_sql_expr(cls, seed: sql.ColumnElement, sql_cols: Iterable[sql.ColumnElement]) -> sql.ColumnElement:
+        """Construct expression which is the ordering key for rows to be sampled
+        General SQL form is:
+        - MD5(<seed::text> [ + '___' + <rowid_col_val>::text]+
+        """
+        sql_expr: sql.ColumnElement = sql.cast(seed, sql.Text)
+        for e in sql_cols:
+            sql_expr = sql_expr + sql.literal_column("'___'") + sql.cast(e, sql.Text)
+        sql_expr = sql.func.md5(sql_expr)
+        return sql_expr
+
+    def _create_order_by(self, cte: sql.CTE) -> sql.ColumnElement:
+        """Create an expression for randomly ordering rows with a given seed"""
+        rowid_cols = [*cte.c[-self.pk_count : -1]]  # exclude the version column
+        assert len(rowid_cols) > 0
+        return self.key_sql_expr(sql.literal_column(str(self.seed)), rowid_cols)
+
+    def _create_stmt(self) -> sql.Select:
+        if self.fraction_samples is not None:
+            return self._create_stmt_fraction(self.fraction_samples)
+        return self._create_stmt_n(self.n_samples, self.n_per_stratum)
+
+    def _create_stmt_n(self, n: Optional[int], n_per_stratum: Optional[int]) -> sql.Select:
+        """Create a Select stmt that returns n samples across all strata"""
+        sql_strata_exprs = [self.sql_elements.get(e) for e in self.stratify_exprs]
+        order_by = self._create_order_by(self.input_cte)
+
+        # Create a list of all columns plus the rank
+        # Get all columns from the input CTE dynamically
+        select_columns = [*self.input_cte.c]
+        select_columns.append(
+            sql.func.row_number().over(partition_by=sql_strata_exprs, order_by=order_by).label('rank')
+        )
+        row_rank_cte = sql.select(*select_columns).select_from(self.input_cte).cte('row_rank_cte')
+
+        final_columns = [*row_rank_cte.c[:-1]]  # exclude the rank column
+        if n_per_stratum is not None:
+            return sql.select(*final_columns).filter(row_rank_cte.c.rank <= n_per_stratum)
+        else:
+            secondary_order = self._create_order_by(row_rank_cte)
+            return sql.select(*final_columns).order_by(row_rank_cte.c.rank, secondary_order).limit(n)
+
+    def _create_stmt_fraction(self, fraction_samples: float) -> sql.Select:
+        """Create a Select stmt that returns a fraction of the rows per strata"""
+
+        # Build the strata count CTE
+        # Produces a table of the form:
+        #   ([stratify_exprs], s_s_size)
+        # where s_s_size is the number of samples to take from each stratum
+        sql_strata_exprs = [self.sql_elements.get(e) for e in self.stratify_exprs]
+        per_strata_count_cte = (
+            sql.select(
+                *sql_strata_exprs,
+                sql.func.ceil(fraction_samples * sql.func.count(1).cast(sql.Integer)).label('s_s_size'),
+            )
+            .select_from(self.input_cte)
+            .group_by(*sql_strata_exprs)
+            .cte('per_strata_count_cte')
+        )
+
+        # Build a CTE that ranks the rows within each stratum
+        # Include all columns from the input CTE dynamically
+        order_by = self._create_order_by(self.input_cte)
+        select_columns = [*self.input_cte.c]
+        select_columns.append(
+            sql.func.row_number().over(partition_by=sql_strata_exprs, order_by=order_by).label('rank')
+        )
+        row_rank_cte = sql.select(*select_columns).select_from(self.input_cte).cte('row_rank_cte')
+
+        # Build the join criterion dynamically to accommodate any number of group by columns
+        join_c = sql.true()
+        for col in per_strata_count_cte.c[:-1]:
+            join_c &= row_rank_cte.c[col.name].isnot_distinct_from(col)
+
+        # Join srcp with per_strata_count_cte to limit returns to the requested fraction of rows
+        final_columns = [*row_rank_cte.c[:-1]]  # exclude the rank column
+        stmt = (
+            sql.select(*final_columns)
+            .select_from(row_rank_cte)
+            .join(per_strata_count_cte, join_c)
+            .where(row_rank_cte.c.rank <= per_strata_count_cte.c.s_s_size)
+        )
+
+        return stmt

pixeltable/exprs/data_row.py

@@ -214,6 +214,7 @@ class DataRow:
         """Assign in-memory cell value
         This allows overwriting
         """
+        assert isinstance(idx, int)
         assert self.excs[idx] is None
 
         if (idx in self.img_slot_idxs or idx in self.media_slot_idxs) and isinstance(val, str):
@@ -253,14 +254,15 @@ class DataRow:
         assert self.excs[index] is None
         if self.file_paths[index] is None:
             if filepath is not None:
-                # we want to save this to a file
-                self.file_paths[index] = filepath
-                self.file_urls[index] = urllib.parse.urljoin('file:', urllib.request.pathname2url(filepath))
                 image = self.vals[index]
                 assert isinstance(image, PIL.Image.Image)
                 # Default to JPEG unless the image has a transparency layer (which isn't supported by JPEG).
                 # In that case, use WebP instead.
                 format = 'webp' if image.has_transparency_data else 'jpeg'
+                if not filepath.endswith(f'.{format}'):
+                    filepath += f'.{format}'
+                self.file_paths[index] = filepath
+                self.file_urls[index] = urllib.parse.urljoin('file:', urllib.request.pathname2url(filepath))
                 image.save(filepath, format=format)
             else:
                 # we discard the content of this cell

pixeltable/exprs/expr.py

@@ -276,6 +276,13 @@ class Expr(abc.ABC):
         tbl_versions = {tbl_version.id: tbl_version.get() for tbl_version in tbl.get_tbl_versions()}
         return self._retarget(tbl_versions)
 
+    @classmethod
+    def retarget_list(cls, expr_list: list[Expr], tbl: catalog.TableVersionPath) -> None:
+        """Retarget ColumnRefs in expr_list to the specific TableVersions in tbl."""
+        tbl_versions = {tbl_version.id: tbl_version.get() for tbl_version in tbl.get_tbl_versions()}
+        for i in range(len(expr_list)):
+            expr_list[i] = expr_list[i]._retarget(tbl_versions)
+
     def _retarget(self, tbl_versions: dict[UUID, catalog.TableVersion]) -> Self:
         for i in range(len(self.components)):
             self.components[i] = self.components[i]._retarget(tbl_versions)

pixeltable/exprs/literal.py

@@ -16,6 +16,8 @@ from .sql_element_cache import SqlElementCache
 
 
 class Literal(Expr):
+    val: Any
+
     def __init__(self, val: Any, col_type: Optional[ts.ColumnType] = None):
         if col_type is not None:
             val = col_type.create_literal(val)

pixeltable/func/tools.py

@@ -2,7 +2,7 @@ from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Union
 
 import pydantic
 
-import pixeltable.exceptions as excs
+from pixeltable import exceptions as excs
 
 from .function import Function
 from .signature import Parameter

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

@@ -53,8 +53,8 @@ async def generate_content(
         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.
+        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.
@@ -103,7 +103,6 @@ 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]:
-    print(response)
     pxt_tool_calls: dict[str, list[dict]] = {}
     for part in response['candidates'][0]['content']['parts']:
         tool_call = part.get('function_call')

pixeltable/functions/llama_cpp.py

@@ -17,7 +17,7 @@ def create_chat_completion(
     model_path: Optional[str] = None,
     repo_id: Optional[str] = None,
     repo_filename: Optional[str] = None,
-    args: Optional[dict[str, Any]] = None,
+    model_kwargs: Optional[dict[str, Any]] = None,
 ) -> dict:
     """
     Generate a chat completion from a list of messages.
@@ -35,14 +35,14 @@ def create_chat_completion(
         repo_id: The Hugging Face model repo id (if using a pretrained model).
         repo_filename: A filename or glob pattern to match the model file in the repo (optional, if using a
             pretrained model).
-        args: Additional arguments to pass to the `create_chat_completions` call, such as `max_tokens`, `temperature`,
-            `top_p`, and `top_k`. For details, see the
+        model_kwargs: Additional keyword args for the llama_cpp `create_chat_completions` API, such as `max_tokens`,
+            `temperature`, `top_p`, and `top_k`. For details, see the
             [llama_cpp create_chat_completions documentation](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.create_chat_completion).
     """
     Env.get().require_package('llama_cpp', min_version=[0, 3, 1])
 
-    if args is None:
-        args = {}
+    if model_kwargs is None:
+        model_kwargs = {}
 
     if (model_path is None) == (repo_id is None):
         raise excs.Error('Exactly one of `model_path` or `repo_id` must be provided.')
@@ -56,7 +56,7 @@ def create_chat_completion(
     else:
         Env.get().require_package('huggingface_hub')
         llm = _lookup_pretrained_model(repo_id, repo_filename, n_gpu_layers)
-    return llm.create_chat_completion(messages, **args)  # type: ignore
+    return llm.create_chat_completion(messages, **model_kwargs)  # type: ignore
 
 
 def _is_gpu_available() -> bool: