pixeltable 0.3.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

pixeltable/func/function.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
-import abc
 import importlib
 import inspect
+from abc import abstractmethod, ABC
 from copy import copy
 from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence, cast
 
@@ -12,7 +12,6 @@ from typing_extensions import Self
 import pixeltable as pxt
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
-
 from .globals import resolve_symbol
 from .signature import Signature
 
@@ -20,7 +19,7 @@ if TYPE_CHECKING:
     from .expr_template_function import ExprTemplate, ExprTemplateFunction
 
 
-class Function(abc.ABC):
+class Function(ABC):
     """Base class for Pixeltable's function interface.
 
     A function in Pixeltable is an object that has a signature and implements __call__().
@@ -44,6 +43,12 @@ class Function(abc.ABC):
     # parameter names as the original function. Each parameter is going to be of type sql.ColumnElement.
     _to_sql: Callable[..., Optional[sql.ColumnElement]]
 
+    # Returns the resource pool to use for calling this function with the given arguments.
+    # Overriden for specific Function instances via the resource_pool() decorator. The override must accept a subset
+    # of the parameters of the original function, with the same type.
+    _resource_pool: Callable[..., Optional[str]]
+
+
     def __init__(
         self,
         signatures: list[Signature],
@@ -60,9 +65,9 @@ class Function(abc.ABC):
         self.is_method = is_method
         self.is_property = is_property
         self._conditional_return_type = None
-        self._to_sql = self.__default_to_sql
-
         self.__resolved_fns = []
+        self._to_sql = self.__default_to_sql
+        self._resource_pool = self.__default_resource_pool
 
     @property
     def name(self) -> str:
@@ -92,6 +97,10 @@ class Function(abc.ABC):
         assert not self.is_polymorphic
         return len(self.signature.parameters)
 
+    @property
+    @abstractmethod
+    def is_async(self) -> bool: ...
+
     def _docstring(self) -> Optional[str]:
         return None
 
@@ -184,6 +193,26 @@ class Function(abc.ABC):
         """Override this to do custom validation of the arguments"""
         assert not self.is_polymorphic
 
+    def _get_callable_args(self, callable: Callable, kwargs: dict[str, Any]) -> dict[str, Any]:
+        """Return the kwargs to pass to callable, given kwargs passed to this function"""
+        bound_args = self.signature.py_signature.bind(**kwargs).arguments
+        # add defaults to bound_args, if not already present
+        bound_args.update({
+            name: param.default
+            for name, param in self.signature.parameters.items() if name not in bound_args and param.has_default()
+        })
+        result: dict[str, Any] = {}
+        sig = inspect.signature(callable)
+        for param in sig.parameters.values():
+            if param.name in bound_args:
+                result[param.name] = bound_args[param.name]
+        return result
+
+    def call_resource_pool(self, kwargs: dict[str, Any]) -> str:
+        """Return the resource pool to use for calling this function with the given arguments"""
+        kw_args = self._get_callable_args(self._resource_pool, kwargs)
+        return self._resource_pool(**kw_args)
+
     def call_return_type(self, args: Sequence[Any], kwargs: dict[str, Any]) -> ts.ColumnType:
         """Return the type of the value returned by calling this function with the given arguments"""
         assert not self.is_polymorphic
@@ -268,10 +297,13 @@ class Function(abc.ABC):
 
         return ExprTemplate(call, new_signature)
 
-    @abc.abstractmethod
     def exec(self, args: Sequence[Any], kwargs: dict[str, Any]) -> Any:
         """Execute the function with the given arguments and return the result."""
-        pass
+        raise NotImplementedError()
+
+    async def aexec(self, *args: Any, **kwargs: Any) -> Any:
+        """Execute the function with the given arguments and return the result."""
+        raise NotImplementedError()
 
     def to_sql(self, fn: Callable[..., Optional[sql.ColumnElement]]) -> Callable[..., Optional[sql.ColumnElement]]:
         """Instance decorator for specifying the SQL translation of this function"""
@@ -282,6 +314,15 @@ class Function(abc.ABC):
         """The default implementation of SQL translation, which provides no translation"""
         return None
 
+    def resource_pool(self, fn: Callable[..., str]) -> Callable[..., str]:
+        """Instance decorator for specifying the resource pool of this function"""
+        # TODO: check that fn's parameters are a subset of our parameters
+        self._resource_pool = fn
+        return fn
+
+    def __default_resource_pool(self) -> Optional[str]:
+        return None
+
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, self.__class__):
             return False

pixeltable/func/query_template_function.py

@@ -18,6 +18,10 @@ if TYPE_CHECKING:
 
 class QueryTemplateFunction(Function):
     """A parameterized query/DataFrame from which an executable DataFrame is created with a function call."""
+    template_df: Optional['DataFrame']
+    self_name: Optional[str]
+    conn: Optional[sql.engine.Connection]
+    defaults: dict[str, exprs.Literal]
 
     @classmethod
     def create(
@@ -48,10 +52,10 @@ class QueryTemplateFunction(Function):
         # if we're running as part of an ongoing update operation, we need to use the same connection, otherwise
         # we end up with a deadlock
         # TODO: figure out a more general way to make execution state available
-        self.conn: Optional[sql.engine.Connection] = None
+        self.conn = None
 
         # convert defaults to Literals
-        self.defaults: dict[str, exprs.Literal] = {}  # key: param name, value: default value converted to a Literal
+        self.defaults = {}  # key: param name, value: default value converted to a Literal
         param_types = self.template_df.parameters()
         for param in [p for p in sig.parameters.values() if p.has_default()]:
             assert param.name in param_types
@@ -65,14 +69,18 @@ class QueryTemplateFunction(Function):
     def set_conn(self, conn: Optional[sql.engine.Connection]) -> None:
         self.conn = conn
 
-    def exec(self, args: Sequence[Any], kwargs: dict[str, Any]) -> Any:
-        assert not self.is_polymorphic
+    @property
+    def is_async(self) -> bool:
+        return True
+
+    async def aexec(self, *args: Any, **kwargs: Any) -> Any:
+        #assert not self.is_polymorphic
         bound_args = self.signature.py_signature.bind(*args, **kwargs).arguments
         # apply defaults, otherwise we might have Parameters left over
         bound_args.update(
             {param_name: default for param_name, default in self.defaults.items() if param_name not in bound_args})
         bound_df = self.template_df.bind(bound_args)
-        result = bound_df._collect(self.conn)
+        result = await bound_df._acollect(self.conn)
         return list(result)
 
     @property
@@ -93,11 +101,12 @@ class QueryTemplateFunction(Function):
 
 
 @overload
-def query(self, py_fn: Callable) -> QueryTemplateFunction: ...
+def query(py_fn: Callable) -> QueryTemplateFunction: ...
 
 @overload
 def query(
-    self, *, param_types: Optional[list[ts.ColumnType]] = None
+    *,
+    param_types: Optional[list[ts.ColumnType]] = None
 ) -> Callable[[Callable], QueryTemplateFunction]: ...
 
 def query(*args: Any, **kwargs: Any) -> Any:

pixeltable/func/udf.py

@@ -15,7 +15,7 @@ from .signature import Signature
 
 # Decorator invoked without parentheses: @pxt.udf
 @overload
-def udf(decorated_fn: Callable) -> Function: ...
+def udf(decorated_fn: Callable) -> CallableFunction: ...
 
 
 # Decorator schema invoked with parentheses: @pxt.udf(**kwargs)
@@ -26,6 +26,7 @@ def udf(
     substitute_fn: Optional[Callable] = None,
     is_method: bool = False,
     is_property: bool = False,
+    resource_pool: Optional[str] = None,
     type_substitutions: Optional[Sequence[dict]] = None,
     _force_stored: bool = False
 ) -> Callable[[Callable], CallableFunction]: ...
@@ -53,6 +54,7 @@ def udf(*args, **kwargs):
         substitute_fn = kwargs.pop('substitute_fn', None)
         is_method = kwargs.pop('is_method', None)
         is_property = kwargs.pop('is_property', None)
+        resource_pool = kwargs.pop('resource_pool', None)
         type_substitutions = kwargs.pop('type_substitutions', None)
         force_stored = kwargs.pop('_force_stored', False)
         if len(kwargs) > 0:
@@ -67,6 +69,7 @@ def udf(*args, **kwargs):
                 substitute_fn=substitute_fn,
                 is_method=is_method,
                 is_property=is_property,
+                resource_pool=resource_pool,
                 type_substitutions=type_substitutions,
                 force_stored=force_stored
             )
@@ -82,6 +85,7 @@ def make_function(
     substitute_fn: Optional[Callable] = None,
     is_method: bool = False,
     is_property: bool = False,
+    resource_pool: Optional[str] = None,
     type_substitutions: Optional[Sequence[dict]] = None,
     function_name: Optional[str] = None,
     force_stored: bool = False
@@ -162,6 +166,8 @@ def make_function(
         is_method=is_method,
         is_property=is_property
     )
+    if resource_pool is not None:
+        result.resource_pool(lambda: resource_pool)
 
     # If this function is part of a module, register it
     if function_path is not None:

pixeltable/functions/__init__.py

@@ -2,7 +2,7 @@ from pixeltable.utils.code import local_public_names
 
 from . import (anthropic, audio, fireworks, gemini, huggingface, image, json, llama_cpp, math, mistralai, ollama,
                openai, string, timestamp, together, video, vision, whisper)
-from .globals import *
+from .globals import count, max, mean, min, sum
 
 __all__ = local_public_names(__name__, exclude=['globals']) + local_public_names(globals.__name__)
 

pixeltable/functions/anthropic.py

@@ -5,12 +5,14 @@ first `pip install anthropic` and configure your Anthropic credentials, as descr
 the [Working with Anthropic](https://pixeltable.readme.io/docs/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, Optional, TypeVar, Union, cast, Iterable
 
-import tenacity
+import httpx
 
 import pixeltable as pxt
-import pixeltable.exceptions as excs
 from pixeltable import env, exprs
 from pixeltable.func import Tools
 from pixeltable.utils.code import local_public_names
@@ -18,28 +20,54 @@ from pixeltable.utils.code import local_public_names
 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)
+class AnthropicRateLimitsInfo(env.RateLimitsInfo):
+
+    def __init__(self):
+        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 get_retry_delay(self, exc: Exception) -> Optional[float]:
+        import anthropic
+
+        # deal with timeouts separately, they don't come with headers
+        if isinstance(exc, anthropic.APITimeoutError):
+            return 1.0
+
+        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
+        retry_after_str = exc.response.headers.get('retry-after', '1')
+        return int(retry_after_str)
 
 
 @pxt.udf
-def messages(
+async def messages(
     messages: list[dict[str, str]],
     *,
     model: str,
@@ -79,6 +107,9 @@ def messages(
         >>> msgs = [{'role': 'user', 'content': tbl.prompt}]
         ... tbl['response'] = messages(msgs, model='claude-3-haiku-20240307')
     """
+
+    # it doesn't look like count_tokens() actually exists in the current version of the library
+
     if tools is not None:
         # Reformat `tools` into Anthropic format
         tools = [
@@ -106,19 +137,62 @@ def messages(
         if not tool_choice['parallel_tool_calls']:
             tool_choice_['disable_parallel_tool_use'] = True
 
-    return _retry(_anthropic_client().messages.create)(
-        messages=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(metadata),
+        metadata=_opt(cast(Any, metadata)),
         stop_sequences=_opt(stop_sequences),
         system=_opt(system),
-        temperature=_opt(temperature),
-        tool_choice=_opt(tool_choice_),
-        tools=_opt(tools),
+        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),
-    ).dict()
+        timeout=10,
+    )
+
+    requests_limit_str = result.headers.get('anthropic-ratelimit-requests-limit')
+    requests_limit = int(requests_limit_str) if requests_limit_str is not None else None
+    requests_remaining_str = result.headers.get('anthropic-ratelimit-requests-remaining')
+    requests_remaining = int(requests_remaining_str) if requests_remaining_str is not None else None
+    requests_reset_str = result.headers.get('anthropic-ratelimit-requests-reset')
+    requests_reset = datetime.datetime.fromisoformat(requests_reset_str.replace('Z', '+00:00'))
+    input_tokens_limit_str = result.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 = result.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 = result.headers.get('anthropic-ratelimit-input-tokens-reset')
+    input_tokens_reset = datetime.datetime.fromisoformat(input_tokens_reset_str.replace('Z', '+00:00'))
+    output_tokens_limit_str = result.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 = result.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 = result.headers.get('anthropic-ratelimit-output-tokens-reset')
+    output_tokens_reset = datetime.datetime.fromisoformat(output_tokens_reset_str.replace('Z', '+00:00'))
+    retry_after_str = result.headers.get('retry-after')
+    if retry_after_str is not None:
+        _logger.debug(f'retry-after: {retry_after_str}')
+
+    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)
+    rate_limits_info.record(
+        requests=(requests_limit, requests_remaining, requests_reset),
+        input_tokens=(input_tokens_limit, input_tokens_remaining, input_tokens_reset),
+        output_tokens=(output_tokens_limit, output_tokens_remaining, output_tokens_reset))
+
+    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:

pixeltable/functions/gemini.py

@@ -13,7 +13,7 @@ from pixeltable import env
 
 @env.register_client('gemini')
 def _(api_key: str) -> None:
-    import google.generativeai as genai  # type: ignore[import-untyped]
+    import google.generativeai as genai
     genai.configure(api_key=api_key)
 
 
@@ -36,8 +36,6 @@ def generate_content(
     response_schema: Optional[dict] = None,
     presence_penalty: Optional[float] = None,
     frequency_penalty: Optional[float] = None,
-    response_logprobs: Optional[bool] = None,
-    logprobs: Optional[int] = None,
 ) -> dict:
     """
     Generate content from the specified model. For additional details, see:
@@ -60,7 +58,7 @@ def generate_content(
         Add a computed column that applies the model `gemini-1.5-flash`
         to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
 
-        >>> tbl['response'] = generate_content(tbl.prompt, model_name='gemini-1.5-flash')
+        >>> tbl.add_computed_column(response=generate_content(tbl.prompt, model_name='gemini-1.5-flash'))
     """
     env.Env.get().require_package('google.generativeai')
     _ensure_loaded()
@@ -78,8 +76,6 @@ def generate_content(
         response_schema=response_schema,
         presence_penalty=presence_penalty,
         frequency_penalty=frequency_penalty,
-        response_logprobs=response_logprobs,
-        logprobs=logprobs,
     )
     response = model.generate_content(contents, generation_config=gc)
     return response.to_dict()