pixeltable 0.2.30__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
 
@@ -119,6 +128,7 @@ class Function(abc.ABC):
                 for idx in range(len(self.signatures)):
                     resolution = cast(Self, copy(self))
                     resolution.signatures = [self.signatures[idx]]
+                    resolution.__resolved_fns = [resolution]  # Resolves to itself
                     resolution._update_as_overload_resolution(idx)
                     self.__resolved_fns.append(resolution)
 
@@ -183,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
@@ -198,13 +228,12 @@ class Function(abc.ABC):
 
     def conditional_return_type(self, fn: Callable[..., ts.ColumnType]) -> Callable[..., ts.ColumnType]:
         """Instance decorator for specifying a conditional return type for this function"""
-        if self.is_polymorphic:
-            raise excs.Error('`conditional_return_type` is not supported for functions with multiple signatures')
         # verify that call_return_type only has parameters that are also present in the signature
-        sig = inspect.signature(fn)
-        for param in sig.parameters.values():
-            if param.name not in self.signature.parameters:
-                raise ValueError(f'`conditional_return_type` has parameter `{param.name}` that is not in the signature')
+        fn_sig = inspect.signature(fn)
+        for param in fn_sig.parameters.values():
+            for self_sig in self.signatures:
+                if param.name not in self_sig.parameters:
+                    raise ValueError(f'`conditional_return_type` has parameter `{param.name}` that is not in a signature')
         self._conditional_return_type = fn
         return fn
 
@@ -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

@@ -1,23 +1,31 @@
 from __future__ import annotations
 
 import inspect
-from typing import Any, Callable, Optional, Sequence
+from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence, overload
 
 import sqlalchemy as sql
 
-import pixeltable as pxt
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
 from pixeltable import exprs
 
 from .function import Function
 from .signature import Signature
 
+if TYPE_CHECKING:
+    from pixeltable import DataFrame
+
 
 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(
-        cls, template_callable: Callable, param_types: Optional[list[pxt.ColumnType]], path: str, name: str
+        cls, template_callable: Callable, param_types: Optional[list[ts.ColumnType]], path: str, name: str
     ) -> QueryTemplateFunction:
         # we need to construct a template df and a signature
         py_sig = inspect.signature(template_callable)
@@ -29,11 +37,11 @@ class QueryTemplateFunction(Function):
         from pixeltable import DataFrame
         assert isinstance(template_df, DataFrame)
         # we take params and return json
-        sig = Signature(return_type=pxt.JsonType(), parameters=params)
+        sig = Signature(return_type=ts.JsonType(), parameters=params)
         return QueryTemplateFunction(template_df, sig, path=path, name=name)
 
     def __init__(
-            self, template_df: Optional['pxt.DataFrame'], sig: Signature, path: Optional[str] = None,
+            self, template_df: Optional['DataFrame'], sig: Signature, path: Optional[str] = None,
             name: Optional[str] = None,
     ):
         assert sig is not None
@@ -44,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
@@ -61,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
@@ -86,3 +98,37 @@ class QueryTemplateFunction(Function):
     def _from_dict(cls, d: dict) -> Function:
         from pixeltable.dataframe import DataFrame
         return cls(DataFrame.from_dict(d['df']), Signature.from_dict(d['signature']), name=d['name'])
+
+
+@overload
+def query(py_fn: Callable) -> QueryTemplateFunction: ...
+
+@overload
+def query(
+    *,
+    param_types: Optional[list[ts.ColumnType]] = None
+) -> Callable[[Callable], QueryTemplateFunction]: ...
+
+def query(*args: Any, **kwargs: Any) -> Any:
+    def make_query_template(
+        py_fn: Callable, param_types: Optional[list[ts.ColumnType]]
+    ) -> QueryTemplateFunction:
+        if py_fn.__module__ != '__main__' and py_fn.__name__.isidentifier():
+            # this is a named function in a module
+            function_path = f'{py_fn.__module__}.{py_fn.__qualname__}'
+        else:
+            function_path = None
+        query_name = py_fn.__name__
+        query_fn = QueryTemplateFunction.create(
+            py_fn, param_types=param_types, path=function_path, name=query_name)
+        return query_fn
+
+        # TODO: verify that the inferred return type matches that of the template
+        # TODO: verify that the signature doesn't contain batched parameters
+
+    if len(args) == 1:
+        assert len(kwargs) == 0 and callable(args[0])
+        return make_query_template(args[0], None)
+    else:
+        assert len(args) == 0 and len(kwargs) == 1 and 'param_types' in kwargs
+        return lambda py_fn: make_query_template(py_fn, kwargs['param_types'])

pixeltable/func/tools.py

@@ -1,10 +1,9 @@
-from dataclasses import dataclass
-import dataclasses
-import json
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, TypeVar, Union
 
 import pydantic
 
+import pixeltable.exceptions as excs
+
 from .function import Function
 from .signature import Parameter
 from .udf import udf
@@ -75,6 +74,13 @@ class Tool(pydantic.BaseModel):
         assert False
 
 
+class ToolChoice(pydantic.BaseModel):
+    auto: bool
+    required: bool
+    tool: Optional[str]
+    parallel_tool_calls: bool
+
+
 class Tools(pydantic.BaseModel):
     tools: list[Tool]
 
@@ -92,25 +98,56 @@ class Tools(pydantic.BaseModel):
             for tool in self.tools
         })
 
+    def choice(
+        self,
+        auto: bool = False,
+        required: bool = False,
+        tool: Union[str, Function, None] = None,
+        parallel_tool_calls: bool = True,
+    ) -> ToolChoice:
+        if sum([auto, required, tool is not None]) != 1:
+            raise excs.Error('Exactly one of `auto`, `required`, or `tool` must be specified.')
+        tool_name: Optional[str] = None
+        if tool is not None:
+            try:
+                tool_obj = next(
+                    t for t in self.tools
+                    if (isinstance(tool, Function) and t.fn == tool)
+                    or (isinstance(tool, str) and (t.name or t.fn.name) == tool)
+                )
+                tool_name = tool_obj.name or tool_obj.fn.name
+            except StopIteration:
+                raise excs.Error(f'That tool is not in the specified list of tools: {tool}')
+        return ToolChoice(auto=auto, required=required, tool=tool_name, parallel_tool_calls=parallel_tool_calls)
+
 
 @udf
-def _extract_str_tool_arg(tool_calls: dict, func_name: str, param_name: str) -> Optional[str]:
-    return str(_extract_arg(tool_calls, func_name, param_name))
+def _extract_str_tool_arg(tool_calls: dict[str, Any], func_name: str, param_name: str) -> Optional[str]:
+    return _extract_arg(str, tool_calls, func_name, param_name)
+
 
 @udf
-def _extract_int_tool_arg(tool_calls: dict, func_name: str, param_name: str) -> Optional[int]:
-    return int(_extract_arg(tool_calls, func_name, param_name))
+def _extract_int_tool_arg(tool_calls: dict[str, Any], func_name: str, param_name: str) -> Optional[int]:
+    return _extract_arg(int, tool_calls, func_name, param_name)
+
 
 @udf
-def _extract_float_tool_arg(tool_calls: dict, func_name: str, param_name: str) -> Optional[float]:
-    return float(_extract_arg(tool_calls, func_name, param_name))
+def _extract_float_tool_arg(tool_calls: dict[str, Any], func_name: str, param_name: str) -> Optional[float]:
+    return _extract_arg(float, tool_calls, func_name, param_name)
+
 
 @udf
-def _extract_bool_tool_arg(tool_calls: dict, func_name: str, param_name: str) -> Optional[bool]:
-    return bool(_extract_arg(tool_calls, func_name, param_name))
+def _extract_bool_tool_arg(tool_calls: dict[str, Any], func_name: str, param_name: str) -> Optional[bool]:
+    return _extract_arg(bool, tool_calls, func_name, param_name)
+
+
+T = TypeVar('T')
+
 
-def _extract_arg(tool_calls: dict, func_name: str, param_name: str) -> Any:
+def _extract_arg(eval_fn: Callable[[Any], T], tool_calls: dict[str, Any], func_name: str, param_name: str) -> Optional[T]:
     if func_name in tool_calls:
         arguments = tool_calls[func_name]['args']
-        return arguments.get(param_name)
+        if param_name in arguments:
+            return eval_fn(arguments[param_name])
+        return None
     return None

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,9 +5,12 @@ 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
 from pixeltable import env, exprs
@@ -17,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,
@@ -47,7 +76,7 @@ def messages(
     stop_sequences: Optional[list[str]] = None,
     system: Optional[str] = None,
     temperature: Optional[float] = None,
-    tool_choice: Optional[list[dict]] = None,
+    tool_choice: Optional[dict] = None,
     tools: Optional[list[dict]] = None,
     top_k: Optional[int] = None,
     top_p: Optional[float] = None,
@@ -78,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 = [
@@ -93,19 +125,74 @@ def messages(
             for tool in tools
         ]
 
-    return _retry(_anthropic_client().messages.create)(
-        messages=messages,
+    tool_choice_: Optional[dict] = None
+    if tool_choice is not None:
+        if tool_choice['auto']:
+            tool_choice_ = {'type': 'auto'}
+        elif tool_choice['required']:
+            tool_choice_ = {'type': 'any'}
+        else:
+            assert tool_choice['tool'] is not None
+            tool_choice_ = {'type': 'tool', 'name': tool_choice['tool']}
+        if not tool_choice['parallel_tool_calls']:
+            tool_choice_['disable_parallel_tool_use'] = True
+
+    # 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()