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

pixeltable/functions/groq.py

@@ -0,0 +1,108 @@
+"""
+Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-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://pixeltable.readme.io/docs/working-with-groq) tutorial.
+"""
+
+from typing import TYPE_CHECKING, Any, Optional
+
+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: Optional[dict[str, Any]] = None,
+    tools: Optional[list[dict[str, Any]]] = None,
+    tool_choice: Optional[dict[str, Any]] = 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 `llama3-8b-8192`
+        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='llama3-8b-8192'))
+    """
+    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__

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:

pixeltable/functions/math.py

@@ -100,6 +100,69 @@ def _(self: sql.ColumnElement, digits: Optional[sql.ColumnElement] = None) -> sq
         return sql.func.round(sql.cast(self, sql.Numeric), sql.cast(digits, sql.Integer))
 
 
+@pxt.udf(is_method=True)
+def pow(self: int, other: int) -> float:
+    """
+    Raise `self` to the power of `other`.
+
+    Equivalent to Python [`self ** other`](https://docs.python.org/3/library/functions.html#pow).
+    """
+    return self**other
+
+
+@pow.to_sql
+def _(self: sql.ColumnElement, other: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.func.pow(self, other)
+
+
+@pxt.udf(is_method=True)
+def bitwise_and(self: int, other: int) -> int:
+    """
+    Bitwise AND of two integers.
+
+    Equivalent to Python
+    [`self & other`](https://docs.python.org/3/library/stdtypes.html#bitwise-operations-on-integer-types).
+    """
+    return self & other
+
+
+@bitwise_and.to_sql
+def _(self: sql.ColumnElement, other: sql.ColumnElement) -> sql.ColumnElement:
+    return self.bitwise_and(other)
+
+
+@pxt.udf(is_method=True)
+def bitwise_or(self: int, other: int) -> int:
+    """
+    Bitwise OR of two integers.
+
+    Equivalent to Python
+    [`self | other`](https://docs.python.org/3/library/stdtypes.html#bitwise-operations-on-integer-types).
+    """
+    return self | other
+
+
+@bitwise_or.to_sql
+def _(self: sql.ColumnElement, other: sql.ColumnElement) -> sql.ColumnElement:
+    return self.bitwise_or(other)
+
+
+@pxt.udf(is_method=True)
+def bitwise_xor(self: int, other: int) -> int:
+    """
+    Bitwise XOR of two integers.
+
+    Equivalent to Python
+    [`self ^ other`](https://docs.python.org/3/library/stdtypes.html#bitwise-operations-on-integer-types).
+    """
+    return self ^ other
+
+
+@bitwise_xor.to_sql
+def _(self: sql.ColumnElement, other: sql.ColumnElement) -> sql.ColumnElement:
+    return self.bitwise_xor(other)
+
+
 __all__ = local_public_names(__name__)
 
 

pixeltable/functions/mistralai.py

@@ -5,7 +5,7 @@ first `pip install mistralai` and configure your Mistral AI credentials, as desc
 the [Working with Mistral AI](https://pixeltable.readme.io/docs/working-with-mistralai) tutorial.
 """
 
-from typing import TYPE_CHECKING, Optional, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Optional
 
 import numpy as np
 
@@ -16,7 +16,7 @@ from pixeltable.func.signature import Batch
 from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
-    import mistralai.types.basemodel
+    import mistralai
 
 
 @register_client('mistral')
@@ -32,16 +32,7 @@ def _mistralai_client() -> 'mistralai.Mistral':
 
 @pxt.udf(resource_pool='request-rate:mistral')
 async def chat_completions(
-    messages: list[dict[str, str]],
-    *,
-    model: str,
-    temperature: Optional[float] = 0.7,
-    top_p: Optional[float] = 1.0,
-    max_tokens: Optional[int] = None,
-    stop: Optional[list[str]] = None,
-    random_seed: Optional[int] = None,
-    response_format: Optional[dict] = None,
-    safe_prompt: Optional[bool] = False,
+    messages: list[dict[str, str]], *, model: str, model_kwargs: Optional[dict[str, Any]] = None
 ) -> dict:
     """
     Chat Completion API.
@@ -60,8 +51,8 @@ async def chat_completions(
     Args:
         messages: The prompt(s) to generate completions for.
         model: ID of the model to use. (See overview here: <https://docs.mistral.ai/getting-started/models/>)
-
-    For details on the other parameters, see: <https://docs.mistral.ai/api/#tag/chat>
+        model_kwargs: Additional keyword args for the Mistral `chat/completions` API.
+            For details on the available parameters, see: <https://docs.mistral.ai/api/#tag/chat>
 
     Returns:
         A dictionary containing the response and other metadata.
@@ -73,34 +64,20 @@ async def chat_completions(
         >>> messages = [{'role': 'user', 'content': tbl.prompt}]
         ... tbl.add_computed_column(response=completions(messages, model='mistral-latest-small'))
     """
+    if model_kwargs is None:
+        model_kwargs = {}
+
     Env.get().require_package('mistralai')
     result = await _mistralai_client().chat.complete_async(
         messages=messages,  # type: ignore[arg-type]
         model=model,
-        temperature=temperature,
-        top_p=top_p,
-        max_tokens=_opt(max_tokens),
-        stop=stop,
-        random_seed=_opt(random_seed),
-        response_format=response_format,  # type: ignore[arg-type]
-        safe_prompt=safe_prompt,
+        **model_kwargs,
     )
     return result.dict()
 
 
 @pxt.udf(resource_pool='request-rate:mistral')
-async def fim_completions(
-    prompt: str,
-    *,
-    model: str,
-    temperature: Optional[float] = 0.7,
-    top_p: Optional[float] = 1.0,
-    max_tokens: Optional[int] = None,
-    min_tokens: Optional[int] = None,
-    stop: Optional[list[str]] = None,
-    random_seed: Optional[int] = None,
-    suffix: Optional[str] = None,
-) -> dict:
+async def fim_completions(prompt: str, *, model: str, model_kwargs: Optional[dict[str, Any]] = None) -> dict:
     """
     Fill-in-the-middle Completion API.
 
@@ -118,6 +95,8 @@ async def fim_completions(
     Args:
         prompt: The text/code to complete.
         model: ID of the model to use. (See overview here: <https://docs.mistral.ai/getting-started/models/>)
+        model_kwargs: Additional keyword args for the Mistral `fim/completions` API.
+            For details on the available parameters, see: <https://docs.mistral.ai/api/#tag/fim>
 
     For details on the other parameters, see: <https://docs.mistral.ai/api/#tag/fim>
 
@@ -130,18 +109,11 @@ async def fim_completions(
 
         >>> tbl.add_computed_column(response=completions(tbl.prompt, model='codestral-latest'))
     """
+    if model_kwargs is None:
+        model_kwargs = {}
+
     Env.get().require_package('mistralai')
-    result = await _mistralai_client().fim.complete_async(
-        prompt=prompt,
-        model=model,
-        temperature=temperature,
-        top_p=top_p,
-        max_tokens=_opt(max_tokens),
-        min_tokens=_opt(min_tokens),
-        stop=stop,
-        random_seed=_opt(random_seed),
-        suffix=_opt(suffix),
-    )
+    result = await _mistralai_client().fim.complete_async(prompt=prompt, model=model, **model_kwargs)
     return result.dict()
 
 
@@ -182,15 +154,6 @@ def _(model: str) -> ts.ArrayType:
     return ts.ArrayType((dimensions,), dtype=ts.FloatType())
 
 
-_T = TypeVar('_T')
-
-
-def _opt(arg: Optional[_T]) -> Union[_T, 'mistralai.types.basemodel.Unset']:
-    from mistralai.types import UNSET
-
-    return arg if arg is not None else UNSET
-
-
 __all__ = local_public_names(__name__)
 
 

pixeltable/functions/ollama.py

@@ -50,7 +50,7 @@ def generate(
         template: Prompt template to use.
         context: The context parameter returned from a previous call to `generate()`.
         raw: If `True`, no formatting will be applied to the prompt.
-        options: Additional options to pass to the `chat` call, such as `max_tokens`, `temperature`, `top_p`, and
+        options: Additional options for the Ollama `chat` call, such as `max_tokens`, `temperature`, `top_p`, and
             `top_k`. For details, see the
             [Valid Parameters and Values](https://github.com/ollama/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values)
             section of the Ollama documentation.