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

pixeltable/functions/anthropic.py

@@ -1,55 +1,167 @@
 """
-Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
+Pixeltable UDFs
 that wrap various endpoints from the Anthropic API. In order to use them, you must
 first `pip install anthropic` and configure your Anthropic credentials, as described in
-the [Working with Anthropic](https://pixeltable.readme.io/docs/working-with-anthropic) tutorial.
+the [Working with Anthropic](https://docs.pixeltable.com/notebooks/integrations/working-with-anthropic) tutorial.
 """
 
-from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Union
+import datetime
+import json
+import logging
+from typing import TYPE_CHECKING, Any, Iterable, cast
 
-import tenacity
+import httpx
 
 import pixeltable as pxt
-from pixeltable import env
+from pixeltable import env, exprs
+from pixeltable.func import Tools
 from pixeltable.utils.code import local_public_names
+from pixeltable.utils.http import exponential_backoff
 
 if TYPE_CHECKING:
     import anthropic
 
+_logger = logging.getLogger('pixeltable')
+
 
 @env.register_client('anthropic')
-def _(api_key: str) -> 'anthropic.Anthropic':
+def _(api_key: str) -> 'anthropic.AsyncAnthropic':
     import anthropic
-    return anthropic.Anthropic(api_key=api_key)
+
+    return anthropic.AsyncAnthropic(
+        api_key=api_key,
+        # recommended to increase limits for async client to avoid connection errors
+        http_client=httpx.AsyncClient(limits=httpx.Limits(max_keepalive_connections=100, max_connections=500)),
+    )
 
 
-def _anthropic_client() -> 'anthropic.Anthropic':
+def _anthropic_client() -> 'anthropic.AsyncAnthropic':
     return env.Env.get().get_client('anthropic')
 
 
-def _retry(fn: Callable) -> Callable:
-    import anthropic
-    return tenacity.retry(
-        retry=tenacity.retry_if_exception_type(anthropic.RateLimitError),
-        wait=tenacity.wait_random_exponential(multiplier=1, max=60),
-        stop=tenacity.stop_after_attempt(20),
-    )(fn)
+def _get_header_info(
+    headers: httpx.Headers,
+) -> tuple[
+    tuple[int, int, datetime.datetime] | None,
+    tuple[int, int, datetime.datetime] | None,
+    tuple[int, int, datetime.datetime] | None,
+]:
+    """Extract rate limit info from Anthropic API response headers."""
+    requests_limit_str = headers.get('anthropic-ratelimit-requests-limit')
+    requests_limit = int(requests_limit_str) if requests_limit_str is not None else None
+    requests_remaining_str = headers.get('anthropic-ratelimit-requests-remaining')
+    requests_remaining = int(requests_remaining_str) if requests_remaining_str is not None else None
+    requests_reset_str = headers.get('anthropic-ratelimit-requests-reset')
+    requests_reset = (
+        datetime.datetime.fromisoformat(requests_reset_str.replace('Z', '+00:00')) if requests_reset_str else None
+    )
+    requests_info = (
+        (requests_limit, requests_remaining, requests_reset) if requests_reset and requests_remaining else None
+    )
+
+    input_tokens_limit_str = headers.get('anthropic-ratelimit-input-tokens-limit')
+    input_tokens_limit = int(input_tokens_limit_str) if input_tokens_limit_str is not None else None
+    input_tokens_remaining_str = headers.get('anthropic-ratelimit-input-tokens-remaining')
+    input_tokens_remaining = int(input_tokens_remaining_str) if input_tokens_remaining_str is not None else None
+    input_tokens_reset_str = headers.get('anthropic-ratelimit-input-tokens-reset')
+    input_tokens_reset = (
+        datetime.datetime.fromisoformat(input_tokens_reset_str.replace('Z', '+00:00'))
+        if input_tokens_reset_str
+        else None
+    )
+    input_tokens_info = (
+        (input_tokens_limit, input_tokens_remaining, input_tokens_reset)
+        if input_tokens_reset and input_tokens_remaining
+        else None
+    )
+
+    output_tokens_limit_str = headers.get('anthropic-ratelimit-output-tokens-limit')
+    output_tokens_limit = int(output_tokens_limit_str) if output_tokens_limit_str is not None else None
+    output_tokens_remaining_str = headers.get('anthropic-ratelimit-output-tokens-remaining')
+    output_tokens_remaining = int(output_tokens_remaining_str) if output_tokens_remaining_str is not None else None
+    output_tokens_reset_str = headers.get('anthropic-ratelimit-output-tokens-reset')
+    output_tokens_reset = (
+        datetime.datetime.fromisoformat(output_tokens_reset_str.replace('Z', '+00:00'))
+        if output_tokens_reset_str
+        else None
+    )
+    output_tokens_info = (
+        (output_tokens_limit, output_tokens_remaining, output_tokens_reset)
+        if output_tokens_reset and output_tokens_remaining
+        else None
+    )
+
+    if requests_info is None or input_tokens_info is None or output_tokens_info is None:
+        _logger.debug(f'get_header_info(): incomplete rate limit info: {headers}')
+
+    return requests_info, input_tokens_info, output_tokens_info
+
+
+class AnthropicRateLimitsInfo(env.RateLimitsInfo):
+    def __init__(self) -> None:
+        super().__init__(self._get_request_resources)
+
+    def _get_request_resources(self, messages: dict, max_tokens: int) -> dict[str, int]:
+        input_len = 0
+        for message in messages:
+            if 'role' in message:
+                input_len += len(message['role'])
+            if 'content' in message:
+                input_len += len(message['content'])
+        return {'requests': 1, 'input_tokens': int(input_len / 4), 'output_tokens': max_tokens}
+
+    def record_exc(self, request_ts: datetime.datetime, exc: Exception) -> None:
+        import anthropic
+
+        if (
+            not isinstance(exc, anthropic.APIError)
+            or not hasattr(exc, 'response')
+            or not hasattr(exc.response, 'headers')
+        ):
+            return
+        requests_info, input_tokens_info, output_tokens_info = _get_header_info(exc.response.headers)
+        _logger.debug(
+            f'record_exc(): request_ts: {request_ts}, requests_info={requests_info} '
+            f'input_tokens_info={input_tokens_info} output_tokens_info={output_tokens_info}'
+        )
+        self.record(
+            request_ts=request_ts,
+            requests=requests_info,
+            input_tokens=input_tokens_info,
+            output_tokens=output_tokens_info,
+        )
+        self.has_exc = True
+
+        retry_after_str = exc.response.headers.get('retry-after')
+        if retry_after_str is not None:
+            _logger.debug(f'retry-after: {retry_after_str}')
+
+    def get_retry_delay(self, exc: Exception, attempt: int) -> float | None:
+        import anthropic
+
+        # deal with timeouts separately, they don't come with headers
+        if isinstance(exc, anthropic.APITimeoutError):
+            return exponential_backoff(attempt)
+
+        if not isinstance(exc, anthropic.APIStatusError):
+            return None
+        _logger.debug(f'headers={exc.response.headers}')
+        should_retry_str = exc.response.headers.get('x-should-retry', '')
+        if should_retry_str.lower() != 'true':
+            return None
+        return super().get_retry_delay(exc, attempt)
 
 
 @pxt.udf
-def messages(
+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[list[dict]] = None,
-    tools: Optional[dict] = None,
-    top_k: Optional[int] = None,
-    top_p: Optional[float] = None,
+    max_tokens: int,
+    model_kwargs: dict[str, Any] | None = None,
+    tools: list[dict[str, Any]] | None = None,
+    tool_choice: dict[str, Any] | None = None,
+    _runtime_ctx: env.RuntimeCtx | None = None,
 ) -> dict:
     """
     Create a Message.
@@ -57,6 +169,10 @@ def messages(
     Equivalent to the Anthropic `messages` API endpoint.
     For additional details, see: <https://docs.anthropic.com/en/api/messages>
 
+    Request throttling:
+    Uses the rate limit-related headers returned by the API to throttle requests adaptively, based on available
+    request and token capacity. No configuration is necessary.
+
     __Requirements:__
 
     - `pip install anthropic`
@@ -64,44 +180,107 @@ 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['response'] = messages(msgs, model='claude-3-haiku-20240307')
+        ... tbl.add_computed_column(response=messages(msgs, model='claude-3-5-sonnet-20241022'))
     """
-    return _retry(_anthropic_client().messages.create)(
-        messages=messages,
-        model=model,
-        max_tokens=max_tokens,
-        metadata=_opt(metadata),
-        stop_sequences=_opt(stop_sequences),
-        system=_opt(system),
-        temperature=_opt(temperature),
-        tool_choice=_opt(tool_choice),
-        tools=_opt(tools),
-        top_k=_opt(top_k),
-        top_p=_opt(top_p),
-    ).dict()
+    if model_kwargs is None:
+        model_kwargs = {}
 
+    if tools is not None:
+        # Reformat `tools` into Anthropic format
+        model_kwargs['tools'] = [
+            {
+                'name': tool['name'],
+                'description': tool['description'],
+                'input_schema': {
+                    'type': 'object',
+                    'properties': tool['parameters']['properties'],
+                    'required': tool['required'],
+                },
+            }
+            for tool in tools
+        ]
 
-_T = TypeVar('_T')
+    if tool_choice is not None:
+        if tool_choice['auto']:
+            model_kwargs['tool_choice'] = {'type': 'auto'}
+        elif tool_choice['required']:
+            model_kwargs['tool_choice'] = {'type': 'any'}
+        else:
+            assert tool_choice['tool'] is not None
+            model_kwargs['tool_choice'] = {'type': 'tool', 'name': tool_choice['tool']}
+        if not tool_choice['parallel_tool_calls']:
+            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}'
+    rate_limits_info = env.Env.get().get_resource_pool_info(resource_pool_id, AnthropicRateLimitsInfo)
+    assert isinstance(rate_limits_info, env.RateLimitsInfo)
 
-def _opt(arg: _T) -> Union[_T, 'anthropic.NotGiven']:
-    import anthropic
-    return arg if arg is not None else anthropic.NOT_GIVEN
+    # TODO: timeouts should be set system-wide and be user-configurable
+    from anthropic.types import MessageParam
+
+    start_ts = datetime.datetime.now(tz=datetime.timezone.utc)
+
+    result = await _anthropic_client().messages.with_raw_response.create(
+        messages=cast(Iterable[MessageParam], messages), model=model, max_tokens=max_tokens, **model_kwargs
+    )
+
+    requests_info, input_tokens_info, output_tokens_info = _get_header_info(result.headers)
+    # retry_after_str = result.headers.get('retry-after')
+    # if retry_after_str is not None:
+    #     _logger.debug(f'retry-after: {retry_after_str}')
+    is_retry = _runtime_ctx is not None and _runtime_ctx.is_retry
+    rate_limits_info.record(
+        request_ts=start_ts,
+        requests=requests_info,
+        input_tokens=input_tokens_info,
+        output_tokens=output_tokens_info,
+        reset_exc=is_retry,
+    )
+
+    result_dict = json.loads(result.text)
+    return result_dict
+
+
+@messages.resource_pool
+def _(model: str) -> str:
+    return f'rate-limits:anthropic:{model}'
+
+
+def invoke_tools(tools: Tools, response: exprs.Expr) -> exprs.InlineDict:
+    """Converts an Anthropic response dict to Pixeltable tool invocation format and calls `tools._invoke()`."""
+    return tools._invoke(_anthropic_response_to_pxt_tool_calls(response))
+
+
+@pxt.udf
+def _anthropic_response_to_pxt_tool_calls(response: dict) -> dict | None:
+    anthropic_tool_calls = [r for r in response['content'] if r['type'] == 'tool_use']
+    if len(anthropic_tool_calls) == 0:
+        return None
+    pxt_tool_calls: dict[str, list[dict[str, Any]]] = {}
+    for tool_call in anthropic_tool_calls:
+        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['input']})
+    return pxt_tool_calls
 
 
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/audio.py

@@ -1,30 +1,167 @@
 """
-Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs) for `AudioType`.
+Pixeltable UDFs for `AudioType`.
+"""
 
-Example:
-```python
-import pixeltable as pxt
-import pixeltable.functions as pxtf
+from typing import Any
 
-t = pxt.get_table(...)
-t.select(pxtf.audio.get_metadata()).collect()
-```
-"""
+import av
+import numpy as np
 
 import pixeltable as pxt
+import pixeltable.utils.av as av_utils
 from pixeltable.utils.code import local_public_names
+from pixeltable.utils.local_store import TempStore
 
 
 @pxt.udf(is_method=True)
 def get_metadata(audio: pxt.Audio) -> dict:
     """
     Gets various metadata associated with an audio file and returns it as a dictionary.
+
+    Args:
+        audio: The audio to get metadata for.
+
+    Returns:
+        A `dict` such as the following:
+
+            ```json
+            {
+                'size': 2568827,
+                'streams': [
+                    {
+                        'type': 'audio',
+                        'frames': 0,
+                        'duration': 2646000,
+                        'metadata': {},
+                        'time_base': 2.2675736961451248e-05,
+                        'codec_context': {
+                            'name': 'flac',
+                            'profile': None,
+                            'channels': 1,
+                            'codec_tag': '\\x00\\x00\\x00\\x00',
+                        },
+                        'duration_seconds': 60.0,
+                    }
+                ],
+                'bit_rate': 342510,
+                'metadata': {'encoder': 'Lavf61.1.100'},
+                'bit_exact': False,
+            }
+            ```
+
+    Examples:
+        Extract metadata for files in the `audio_col` column of the table `tbl`:
+
+        >>> tbl.select(tbl.audio_col.get_metadata()).collect()
+    """
+    return av_utils.get_metadata(audio)
+
+
+@pxt.udf()
+def encode_audio(
+    audio_data: pxt.Array[pxt.Float], *, input_sample_rate: int, format: str, output_sample_rate: int | None = None
+) -> pxt.Audio:
+    """
+    Encodes an audio clip represented as an array into a specified audio format.
+
+    Parameters:
+        audio_data: An array of sampled amplitudes. The accepted array shapes are `(N,)` or `(1, N)` for mono audio
+            or `(2, N)` for stereo.
+        input_sample_rate: The sample rate of the input audio data.
+        format: The desired output audio format. The supported formats are 'wav', 'mp3', 'flac', and 'mp4'.
+        output_sample_rate: The desired sample rate for the output audio. Defaults to the input sample rate if
+            unspecified.
+
+    Examples:
+        Add a computed column with encoded FLAC audio files to a table with audio data (as arrays of floats) and sample
+        rates:
+
+        >>> t.add_computed_column(
+        ...     audio_file=encode_audio(
+        ...         t.audio_data, input_sample_rate=t.sample_rate, format='flac'
+        ...     )
+        ... )
+    """
+    if format not in av_utils.AUDIO_FORMATS:
+        raise pxt.Error(f'Only the following formats are supported: {av_utils.AUDIO_FORMATS.keys()}')
+    if output_sample_rate is None:
+        output_sample_rate = input_sample_rate
+
+    codec, ext = av_utils.AUDIO_FORMATS[format]
+    output_path = str(TempStore.create_path(extension=f'.{ext}'))
+
+    match audio_data.shape:
+        case (_,):
+            # Mono audio as 1D array, reshape for pyav
+            layout = 'mono'
+            audio_data_transformed = audio_data[None, :]
+        case (1, _):
+            # Mono audio as 2D array, simply reshape and transpose the input for pyav
+            layout = 'mono'
+            audio_data_transformed = audio_data.reshape(-1, 1).transpose()
+        case (2, _):
+            # Stereo audio. Input layout: [[L0, L1, L2, ...],[R0, R1, R2, ...]],
+            # pyav expects: [L0, R0, L1, R1, L2, R2, ...]
+            layout = 'stereo'
+            audio_data_transformed = np.empty(audio_data.shape[1] * 2, dtype=audio_data.dtype)
+            audio_data_transformed[0::2] = audio_data[0]
+            audio_data_transformed[1::2] = audio_data[1]
+            audio_data_transformed = audio_data_transformed.reshape(1, -1)
+        case _:
+            raise pxt.Error(
+                f'Supported input array shapes are (N,), (1, N) for mono and (2, N) for stereo, got {audio_data.shape}'
+            )
+
+    with av.open(output_path, mode='w') as output_container:
+        stream = output_container.add_stream(codec, rate=output_sample_rate)
+        assert isinstance(stream, av.AudioStream)
+
+        frame = av.AudioFrame.from_ndarray(audio_data_transformed, format='flt', layout=layout)
+        frame.sample_rate = input_sample_rate
+
+        for packet in stream.encode(frame):
+            output_container.mux(packet)
+        for packet in stream.encode():
+            output_container.mux(packet)
+
+        return output_path
+
+
+def audio_splitter(
+    audio: Any, chunk_duration_sec: float, *, overlap_sec: float = 0.0, min_chunk_duration_sec: float = 0.0
+) -> tuple[type[pxt.iterators.ComponentIterator], dict[str, Any]]:
+    """
+    Iterator over chunks of an audio file. The audio file is split into smaller chunks,
+    where the duration of each chunk is determined by chunk_duration_sec.
+    The iterator yields audio chunks as pxt.Audio, along with the start and end time of each chunk.
+    If the input contains no audio, no chunks are yielded.
+
+    Args:
+        chunk_duration_sec: Audio chunk duration in seconds
+        overlap_sec: Overlap between consecutive chunks in seconds
+        min_chunk_duration_sec: Drop the last chunk if it is smaller than min_chunk_duration_sec
+
+    Examples:
+        This example assumes an existing table `tbl` with a column `audio` of type `pxt.Audio`.
+
+        Create a view that splits all audio files into chunks of 30 seconds with 5 seconds overlap:
+
+        >>> pxt.create_view(
+        ...     'audio_chunks',
+        ...     tbl,
+        ...     iterator=audio_splitter(tbl.audio, chunk_duration_sec=30.0, overlap_sec=5.0)
+        ... )
     """
-    return pxt.functions.video._get_metadata(audio)
+    kwargs: dict[str, Any] = {}
+    if overlap_sec != 0.0:
+        kwargs['overlap_sec'] = overlap_sec
+    if min_chunk_duration_sec != 0.0:
+        kwargs['min_chunk_duration_sec'] = min_chunk_duration_sec
+    return pxt.iterators.AudioSplitter._create(audio=audio, chunk_duration_sec=chunk_duration_sec, **kwargs)
 
 
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/bedrock.py

@@ -0,0 +1,137 @@
+"""
+Pixeltable UDFs for AWS Bedrock AI models.
+
+Provides integration with AWS Bedrock for accessing various foundation models
+including Anthropic Claude, Amazon Titan, and other providers.
+"""
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+import pixeltable as pxt
+from pixeltable import env, exprs
+from pixeltable.func import Tools
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    from botocore.client import BaseClient
+
+_logger = logging.getLogger('pixeltable')
+
+
+@env.register_client('bedrock')
+def _() -> 'BaseClient':
+    import boto3
+
+    return boto3.client(service_name='bedrock-runtime')
+
+
+# boto3 typing is weird; type information is dynamically defined, so the best we can do for the static checker is `Any`
+def _bedrock_client() -> Any:
+    return env.Env.get().get_client('bedrock')
+
+
+@pxt.udf
+def converse(
+    messages: list[dict[str, Any]],
+    *,
+    model_id: str,
+    system: list[dict[str, Any]] | None = None,
+    inference_config: dict | None = None,
+    additional_model_request_fields: dict | None = None,
+    tool_config: list[dict] | None = None,
+) -> dict:
+    """
+    Generate a conversation response.
+
+    Equivalent to the AWS Bedrock `converse` API endpoint.
+    For additional details, see: <https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html>
+
+    __Requirements:__
+
+    - `pip install boto3`
+
+    Args:
+        messages: Input messages.
+        model_id: The model that will complete your prompt.
+        system: An optional system prompt.
+        inference_config: Base inference parameters to use.
+        additional_model_request_fields: Additional inference parameters to use.
+
+    For details on the optional parameters, see:
+    <https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html>
+
+    Returns:
+        A dictionary containing the response and other metadata.
+
+    Examples:
+        Add a computed column that applies the model `anthropic.claude-3-haiku-20240307-v1:0`
+        to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
+
+        >>> msgs = [{'role': 'user', 'content': [{'text': tbl.prompt}]}]
+        ... tbl.add_computed_column(response=messages(msgs, model_id='anthropic.claude-3-haiku-20240307-v1:0'))
+    """
+
+    kwargs: dict[str, Any] = {'messages': messages, 'modelId': model_id}
+
+    if system is not None:
+        kwargs['system'] = system
+    if inference_config is not None:
+        kwargs['inferenceConfig'] = inference_config
+    if additional_model_request_fields is not None:
+        kwargs['additionalModelRequestFields'] = additional_model_request_fields
+
+    if tool_config is not None:
+        tool_config_ = {
+            'tools': [
+                {
+                    'toolSpec': {
+                        'name': tool['name'],
+                        'description': tool['description'],
+                        'inputSchema': {
+                            'json': {
+                                'type': 'object',
+                                'properties': tool['parameters']['properties'],
+                                'required': tool['required'],
+                            }
+                        },
+                    }
+                }
+                for tool in tool_config
+            ]
+        }
+        kwargs['toolConfig'] = tool_config_
+
+    return _bedrock_client().converse(**kwargs)
+
+
+def invoke_tools(tools: Tools, response: exprs.Expr) -> exprs.InlineDict:
+    """Converts an Anthropic response dict to Pixeltable tool invocation format and calls `tools._invoke()`."""
+    return tools._invoke(_bedrock_response_to_pxt_tool_calls(response))
+
+
+@pxt.udf
+def _bedrock_response_to_pxt_tool_calls(response: dict) -> dict | None:
+    if response.get('stopReason') != 'tool_use':
+        return None
+
+    pxt_tool_calls: dict[str, list[dict[str, Any]]] = {}
+    for message in response['output']['message']['content']:
+        if 'toolUse' in message:
+            tool_call = message['toolUse']
+            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['input']})
+
+    if len(pxt_tool_calls) == 0:
+        return None
+
+    return pxt_tool_calls
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__