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

pixeltable/functions/globals.py

@@ -1,6 +1,6 @@
 import builtins
 import typing
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable
 
 import sqlalchemy as sql
 
@@ -11,7 +11,7 @@ from typing import _GenericAlias  # type: ignore[attr-defined]  # isort: skip
 
 
 # TODO: remove and replace calls with astype()
-def cast(expr: exprs.Expr, target_type: Union[ts.ColumnType, type, _GenericAlias]) -> exprs.Expr:
+def cast(expr: exprs.Expr, target_type: ts.ColumnType | type | _GenericAlias) -> exprs.Expr:
     expr.col_type = ts.ColumnType.normalize_type(target_type)
     return expr
 
@@ -19,9 +19,30 @@ def cast(expr: exprs.Expr, target_type: Union[ts.ColumnType, type, _GenericAlias
 T = typing.TypeVar('T')
 
 
-@func.uda(allows_window=True, type_substitutions=({T: Optional[int]}, {T: Optional[float]}))  # type: ignore[misc]
+@func.uda(allows_window=True, type_substitutions=({T: int | None}, {T: float | None}))  # type: ignore[misc]
 class sum(func.Aggregator, typing.Generic[T]):
-    """Sums the selected integers or floats."""
+    """
+    Aggregate function that computes the sum of non-null values of a numeric column or grouping.
+
+    Args:
+        val: The numeric value to add to the sum.
+
+    Returns:
+        The sum of the non-null values, or `None` if there are no non-null values.
+
+    Examples:
+        Sum the values in the `value` column of the table `tbl`:
+
+        >>> tbl.select(pxt.functions.sum(tbl.value)).collect()
+
+        Group by the `category` column and compute the sum of the `value` column for each category,
+        assigning the name `'category_total'` to the new column:
+
+        >>> tbl.group_by(tbl.category).select(
+        ...     tbl.category,
+        ...     category_total=pxt.functions.sum(tbl.value)
+        ... ).collect()
+    """
 
     def __init__(self) -> None:
         self.sum: T = None
@@ -39,7 +60,7 @@ class sum(func.Aggregator, typing.Generic[T]):
 
 
 @sum.to_sql
-def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
+def _(val: sql.ColumnElement) -> sql.ColumnElement | None:
     # This can produce a Decimal. We are deliberately avoiding an explicit cast to a Bigint here, because that can
     # cause overflows in Postgres. We're instead doing the conversion to the target type in SqlNode.__iter__().
     return sql.sql.func.sum(val)
@@ -49,9 +70,32 @@ 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]} for t in ts.ALL_PIXELTABLE_TYPES),  # type: ignore[misc]
+    type_substitutions=tuple({T: t | None} for t in ts.ALL_PIXELTABLE_TYPES),  # type: ignore[misc]
 )
 class count(func.Aggregator, typing.Generic[T]):
+    """
+    Aggregate function that counts the number of non-null values in a column or grouping.
+
+    Args:
+        val: The value to count.
+
+    Returns:
+        The count of non-null values.
+
+    Examples:
+        Count the number of non-null values in the `value` column of the table `tbl`:
+
+        >>> tbl.select(pxt.functions.count(tbl.value)).collect()
+
+        Group by the `category` column and compute the count of non-null values in the `value` column
+        for each category, assigning the name `'category_count'` to the new column:
+
+        >>> tbl.group_by(tbl.category).select(
+        ...     tbl.category,
+        ...     category_count=pxt.functions.count(tbl.value)
+        ... ).collect()
+    """
+
     def __init__(self) -> None:
         self.count = 0
 
@@ -64,15 +108,38 @@ class count(func.Aggregator, typing.Generic[T]):
 
 
 @count.to_sql
-def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
+def _(val: sql.ColumnElement) -> sql.ColumnElement | None:
     return sql.sql.func.count(val)
 
 
 @func.uda(
     allows_window=True,
-    type_substitutions=tuple({T: Optional[t]} for t in (str, int, float, bool, ts.Timestamp)),  # type: ignore[misc]
+    type_substitutions=tuple({T: t | None} for t in (str, int, float, bool, ts.Timestamp)),  # type: ignore[misc]
 )
 class min(func.Aggregator, typing.Generic[T]):
+    """
+    Aggregate function that computes the minimum value in a column or grouping.
+
+    Args:
+        val: The value to compare.
+
+    Returns:
+        The minimum value, or `None` if there are no non-null values.
+
+    Examples:
+        Compute the minimum value in the `value` column of the table `tbl`:
+
+        >>> tbl.select(pxt.functions.min(tbl.value)).collect()
+
+        Group by the `category` column and compute the minimum value in the `value` column for each category,
+        assigning the name `'category_min'` to the new column:
+
+        >>> tbl.group_by(tbl.category).select(
+        ...     tbl.category,
+        ...     category_min=pxt.functions.min(tbl.value)
+        ... ).collect()
+    """
+
     def __init__(self) -> None:
         self.val: T = None
 
@@ -89,7 +156,7 @@ class min(func.Aggregator, typing.Generic[T]):
 
 
 @min.to_sql
-def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
+def _(val: sql.ColumnElement) -> sql.ColumnElement | None:
     if val.type.python_type is bool:
         # TODO: min/max aggregation of booleans is not supported in Postgres (but it is in Python).
         # Right now we simply force the computation to be done in Python; we might consider implementing an alternate
@@ -100,9 +167,32 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
 
 @func.uda(
     allows_window=True,
-    type_substitutions=tuple({T: Optional[t]} for t in (str, int, float, bool, ts.Timestamp)),  # type: ignore[misc]
+    type_substitutions=tuple({T: t | None} for t in (str, int, float, bool, ts.Timestamp)),  # type: ignore[misc]
 )
 class max(func.Aggregator, typing.Generic[T]):
+    """
+    Aggregate function that computes the maximum value in a column or grouping.
+
+    Args:
+        val: The value to compare.
+
+    Returns:
+        The maximum value, or `None` if there are no non-null values.
+
+    Examples:
+        Compute the maximum value in the `value` column of the table `tbl`:
+
+        >>> tbl.select(pxt.functions.max(tbl.value)).collect()
+
+        Group by the `category` column and compute the maximum value in the `value` column for each category,
+        assigning the name `'category_max'` to the new column:
+
+        >>> tbl.group_by(tbl.category).select(
+        ...     tbl.category,
+        ...     category_max=pxt.functions.max(tbl.value)
+        ... ).collect()
+    """
+
     def __init__(self) -> None:
         self.val: T = None
 
@@ -119,15 +209,38 @@ class max(func.Aggregator, typing.Generic[T]):
 
 
 @max.to_sql
-def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
+def _(val: sql.ColumnElement) -> sql.ColumnElement | None:
     if val.type.python_type is bool:
         # TODO: see comment in @min.to_sql.
         return None
     return sql.sql.func.max(val)
 
 
-@func.uda(type_substitutions=({T: Optional[int]}, {T: Optional[float]}))  # type: ignore[misc]
+@func.uda(type_substitutions=({T: int | None}, {T: float | None}))  # type: ignore[misc]
 class mean(func.Aggregator, typing.Generic[T]):
+    """
+    Aggregate function that computes the mean (average) of non-null values of a numeric column or grouping.
+
+    Args:
+        val: The numeric value to include in the mean.
+
+    Returns:
+        The mean of the non-null values, or `None` if there are no non-null values.
+
+    Examples:
+        Compute the mean of the values in the `value` column of the table `tbl`:
+
+        >>> tbl.select(pxt.functions.mean(tbl.value)).collect()
+
+        Group by the `category` column and compute the mean of the `value` column for each category,
+        assigning the name `'category_mean'` to the new column:
+
+        >>> tbl.group_by(tbl.category).select(
+        ...     tbl.category,
+        ...     category_mean=pxt.functions.mean(tbl.value)
+        ... ).collect()
+    """
+
     def __init__(self) -> None:
         self.sum: T = None
         self.count = 0
@@ -141,18 +254,33 @@ class mean(func.Aggregator, typing.Generic[T]):
             self.sum += val  # type: ignore[operator]
         self.count += 1
 
-    def value(self) -> Optional[float]:  # Always a float
+    def value(self) -> float | None:  # Always a float
         if self.count == 0:
             return None
         return self.sum / self.count  # type: ignore[operator]
 
 
 @mean.to_sql
-def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
+def _(val: sql.ColumnElement) -> sql.ColumnElement | None:
     return sql.sql.func.avg(val)
 
 
 def map(expr: exprs.Expr, fn: Callable[[exprs.Expr], Any]) -> exprs.Expr:
+    """
+    Applies a mapping function to each element of a list.
+
+    Args:
+        expr: The list expression to map over; must be an expression of type `pxt.Json`.
+        fn: An operation on Pixeltable expressions that will be applied to each element of the JSON array.
+
+    Examples:
+        Given a table `tbl` with a column `data` of type `pxt.Json` containing lists of integers, add a computed
+        column that produces new lists with each integer doubled:
+
+        >>> tbl.add_computed_column(
+        ...     doubled=pxt.functions.map(t.data, lambda x: x * 2)
+        ... )
+    """
     target_expr: exprs.Expr
     try:
         target_expr = exprs.Expr.from_object(fn(exprs.json_path.RELATIVE_PATH_ROOT))

pixeltable/functions/groq.py

@@ -0,0 +1,108 @@
+"""
+Pixeltable UDFs
+that wrap various endpoints from the Groq API. In order to use them, you must
+first `pip install groq` and configure your Groq credentials, as described in
+the [Working with Groq](https://docs.pixeltable.com/notebooks/integrations/working-with-groq) tutorial.
+"""
+
+from typing import TYPE_CHECKING, Any
+
+import pixeltable as pxt
+from pixeltable import exprs
+from pixeltable.env import Env, register_client
+from pixeltable.utils.code import local_public_names
+
+from .openai import _openai_response_to_pxt_tool_calls
+
+if TYPE_CHECKING:
+    import groq
+
+
+@register_client('groq')
+def _(api_key: str) -> 'groq.AsyncGroq':
+    import groq
+
+    return groq.AsyncGroq(api_key=api_key)
+
+
+def _groq_client() -> 'groq.AsyncGroq':
+    return Env.get().get_client('groq')
+
+
+@pxt.udf(resource_pool='request-rate:groq')
+async def chat_completions(
+    messages: list[dict[str, str]],
+    *,
+    model: str,
+    model_kwargs: dict[str, Any] | None = None,
+    tools: list[dict[str, Any]] | None = None,
+    tool_choice: dict[str, Any] | None = None,
+) -> dict:
+    """
+    Chat Completion API.
+
+    Equivalent to the Groq `chat/completions` API endpoint.
+    For additional details, see: <https://console.groq.com/docs/api-reference#chat-create>
+
+    Request throttling:
+    Applies the rate limit set in the config (section `groq`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
+    __Requirements:__
+
+    - `pip install groq`
+
+    Args:
+        messages: A list of messages comprising the conversation so far.
+        model: ID of the model to use. (See overview here: <https://console.groq.com/docs/models>)
+        model_kwargs: Additional keyword args for the Groq `chat/completions` API.
+            For details on the available parameters, see: <https://console.groq.com/docs/api-reference#chat-create>
+
+    Returns:
+        A dictionary containing the response and other metadata.
+
+    Examples:
+        Add a computed column that applies the model `llama-3.1-8b-instant`
+        to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
+
+        >>> messages = [{'role': 'user', 'content': tbl.prompt}]
+        ... tbl.add_computed_column(response=chat_completions(messages, model='llama-3.1-8b-instant'))
+    """
+    if model_kwargs is None:
+        model_kwargs = {}
+
+    Env.get().require_package('groq')
+
+    if tools is not None:
+        model_kwargs['tools'] = [{'type': 'function', 'function': tool} for tool in tools]
+
+    if tool_choice is not None:
+        if tool_choice['auto']:
+            model_kwargs['tool_choice'] = 'auto'
+        elif tool_choice['required']:
+            model_kwargs['tool_choice'] = 'required'
+        else:
+            assert tool_choice['tool'] is not None
+            model_kwargs['tool_choice'] = {'type': 'function', 'function': {'name': tool_choice['tool']}}
+
+    if tool_choice is not None and not tool_choice['parallel_tool_calls']:
+        model_kwargs['parallel_tool_calls'] = False
+
+    result = await _groq_client().chat.completions.create(
+        messages=messages,  # type: ignore[arg-type]
+        model=model,
+        **model_kwargs,
+    )
+    return result.model_dump()
+
+
+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(_openai_response_to_pxt_tool_calls(response))
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__