inspect-ai 0.3.99__py3-none-any.whl → 0.3.100__py3-none-any.whl

inspect_ai/tool/_tools/_web_search/_web_search.py

@@ -1,68 +1,123 @@
-from typing import Literal
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Literal,
+    TypeAlias,
+    TypedDict,
+    get_args,
+)
+
+from typing_extensions import Unpack
 
 from inspect_ai._util.deprecation import deprecation_warning
+from inspect_ai.tool._tool_def import ToolDef
 
 from ..._tool import Tool, ToolResult, tool
-from ._google import google_search_provider, maybe_get_google_api_keys
-from ._tavily import tavily_search_provider
+from ._google import GoogleOptions, google_search_provider
+from ._tavily import TavilyOptions, tavily_search_provider
+
+Provider: TypeAlias = Literal["openai", "tavily", "google"]  # , "gemini", "anthropic"
+valid_providers = set(get_args(Provider))
+
+
+# It would have been nice if the values below were TypedDicts. The problem is
+# that if the caller creates a literal dict variable (rather than passing the
+# dict inline), the type checker will erase the type of the literal to something
+# that doesn't conform the the required TypedDict when passed. This is lame, but
+# we'll do runtime validation instead.
+#
+# If the caller uses this dict form and uses a value of `None`, it means that
+# they want to use that provider and to use the default options.
+class Providers(TypedDict, total=False):
+    google: dict[str, Any] | None
+    tavily: dict[str, Any] | None
+    openai: dict[str, Any] | None
+
+
+class WebSearchDeprecatedArgs(TypedDict, total=False):
+    provider: Literal["tavily", "google"] | None
+    num_results: int | None
+    max_provider_calls: int | None
+    max_connections: int | None
+    model: str | None
 
 
 @tool
 def web_search(
-    provider: Literal["tavily", "google"] | None = None,
-    num_results: int = 3,
-    max_provider_calls: int = 3,
-    max_connections: int = 10,
-    model: str | None = None,
+    providers: Provider | Providers | list[Provider | Providers] | None = None,
+    **deprecated: Unpack[WebSearchDeprecatedArgs],
 ) -> Tool:
     """Web search tool.
 
-    A tool that can be registered for use by models to search the web. Use
-    the `use_tools()` solver to make the tool available (e.g.
-    `use_tools(web_search(provider="tavily"))`))
+    Web searches are executed using a provider. Providers are split
+    into two categories:
 
-    A web search is conducted using the specified provider.
-    - When using Tavily, all logic for relevance and summarization is handled by
-    the Tavily API.
-    - When using Google, the results are parsed for relevance using the specified
-    model, and the top 'num_results' relevant pages are returned.
+    - Internal providers: "openai" - these use the model's built-in search
+      capability and do not require separate API keys. These work only for
+      their respective model provider (e.g. the "openai" search provider
+      works only for `openai/*` models).
+
+    - External providers: "tavily" and "google". These are external services
+      that work with any m odel and require separate accounts and API keys.
+
+    Internal providers will be prioritized if running on the corresponding model
+    (e.g., "openai" provider will be used when running on `openai` models). If an
+    internal provider is specified but the evaluation is run with a different
+    model, a fallback external provider must also be specified.
 
     See further documentation at <https://inspect.aisi.org.uk/tools-standard.html#sec-web-search>.
 
     Args:
-      provider: Search provider to use:
-        - "tavily": Uses Tavily's Research API.
-        - "google": Uses Google Custom Search.
-        Note: The `| None` type is only for backwards compatibility. Passing
-        `None` is deprecated.
-      num_results: The number of search result pages used to provide information
-        back to the model.
-      max_provider_calls: Maximum number of search calls to make to the search
-        provider.
-      max_connections: Maximum number of concurrent connections to API endpoint
-        of search provider.
-      model: Model used to parse web pages for relevance - used only by the
-        `google` provider.
+      providers: Configuration for the search providers to use. Currently supported
+        providers are "openai","tavily", and "google", The `providers` parameter
+        supports several formats based on either a `str` specifying a provider or
+        a `dict` whose keys are the provider names and whose values are the
+        provider-specific options. A single value or a list of these can be passed.
+        This arg is optional just for backwards compatibility. New code should
+        always provide this argument.
+
+        Single provider:
+        ```
+        web_search("tavily")
+        web_search({"tavily": {"max_results": 5}})  # Tavily-specific options
+        ```
+
+        Multiple providers:
+        ```
+        # "openai" used for OpenAI models, "tavily" as fallback
+        web_search(["openai", "tavily"])
+
+        # The None value means to use the provider with default options
+        web_search({"openai": None, "tavily": {"max_results": 5}}
+        ```
+
+        Mixed format:
+        ```
+        web_search(["openai", {"tavily": {"max_results": 5}}])
+        ```
+
+        When specified in the `dict` format, the `None` value for a provider means
+        to use the provider with default options.
+
+        Provider-specific options:
+        - openai: Supports OpenAI's web search parameters.
+          See https://platform.openai.com/docs/guides/tools-web-search?api-mode=responses
+
+        - tavily: Supports options like `max_results`, `search_depth`, etc.
+          See https://docs.tavily.com/documentation/api-reference/endpoint/search
+
+        - google: Supports options like `num_results`, `max_provider_calls`,
+          `max_connections`, and `model`
+
+      **deprecated: Deprecated arguments.
 
     Returns:
        A tool that can be registered for use by models to search the web.
     """
-    if provider is None:
-        if maybe_get_google_api_keys():
-            deprecation_warning(
-                "The `google` `web_search` provider was inferred based on the presence of environment variables. Please specify the provider explicitly to avoid this warning."
-            )
-            provider = "google"
-        else:
-            raise ValueError(
-                "Omitting `provider` is no longer supported. Please specify the `web_search` provider explicitly to avoid this error."
-            )
+    normalized_providers = _normalize_config(providers, **deprecated)
 
-    search_provider = (
-        google_search_provider(num_results, max_provider_calls, max_connections, model)
-        if provider == "google"
-        else tavily_search_provider(num_results, max_connections)
-    )
+    search_provider: Callable[[str], Awaitable[str | None]] | None = None
 
     async def execute(query: str) -> ToolResult:
         """
@@ -71,6 +126,9 @@ def web_search(
         Args:
             query (str): Search query.
         """
+        nonlocal search_provider
+        if not search_provider:
+            search_provider = _create_external_provider(normalized_providers)
         search_result = await search_provider(query)
 
         return (
@@ -82,4 +140,115 @@ def web_search(
             else ("I'm sorry, I couldn't find any relevant information on the web.")
         )
 
-    return execute
+    return ToolDef(
+        execute, name="web_search", options=dict(normalized_providers)
+    ).as_tool()
+
+
+def _normalize_config(
+    providers: Provider | Providers | list[Provider | Providers] | None,
+    **deprecated: Unpack[WebSearchDeprecatedArgs],
+) -> Providers:
+    """
+    Deal with breaking changes in the web_search parameter list.
+
+    This function adapts (hopefully) all of the old variants of how the tool
+    factory may have been called converts to the new config format.
+    """
+    # Cases to handle:
+    # 1. Both deprecated_provider and providers are set
+    #     ValueError
+    # 2. Neither deprecated_provider nor providers is set
+    #     act as if they passed provider="google"
+    # 3. Only providers is set
+    #     if any of the other deprecated parameters is set, then ValueError
+    #     else Happy path
+    # 4. Only deprecated_provider is set
+    #     convert to new config format - including processing old other params
+
+    deprecated_provider = deprecated.get("provider", None)
+    # Case 1.
+    if deprecated_provider and providers:
+        raise ValueError("`provider` is deprecated. Please only specify `providers`.")
+
+    # Case 2.
+    if providers is None and deprecated_provider is None:
+        deprecated_provider = "google"
+
+    num_results = deprecated.get("num_results", None)
+    max_provider_calls = deprecated.get("max_provider_calls", None)
+    max_connections = deprecated.get("max_connections", None)
+    model = deprecated.get("model", None)
+
+    # Getting here means that we have either a providers or a deprecated_provider
+    if deprecated_provider:
+        return _get_config_via_back_compat(
+            deprecated_provider,
+            num_results=num_results,
+            max_provider_calls=max_provider_calls,
+            max_connections=max_connections,
+            model=model,
+        )
+
+    assert providers, "providers should not be None here"
+    normalized: Providers = {}
+    for entry in providers if isinstance(providers, list) else [providers]:
+        if isinstance(entry, str):
+            if entry not in valid_providers:
+                raise ValueError(f"Invalid provider: '{entry}'")
+            normalized[entry] = None  # type: ignore
+        else:
+            for key, value in entry.items():
+                if key not in valid_providers:
+                    raise ValueError(f"Invalid provider: '{key}'")
+                normalized[key] = value  # type: ignore
+    return normalized
+
+
+def _get_config_via_back_compat(
+    provider: Literal["tavily", "google"],
+    num_results: int | None,
+    max_provider_calls: int | None,
+    max_connections: int | None,
+    model: str | None,
+) -> Providers:
+    if (
+        num_results is None
+        and max_provider_calls is None
+        and max_connections is None
+        and model is None
+    ):
+        return {"google": None} if provider == "google" else {"tavily": None}
+
+    # If we get here, we have at least one old school parameter
+    deprecation_warning(
+        "The `num_results`, `max_provider_calls`, `max_connections`, and `model` parameters are deprecated. Please use the `config` parameter instead."
+    )
+
+    if provider == "google":
+        return {
+            "google": GoogleOptions(
+                num_results=num_results,
+                max_provider_calls=max_provider_calls,
+                max_connections=max_connections,
+                model=model,
+            ).model_dump(exclude_none=True)
+        }
+    else:
+        return {
+            "tavily": TavilyOptions(
+                max_results=num_results, max_connections=max_connections
+            ).model_dump(exclude_none=True)
+        }
+
+
+def _create_external_provider(
+    providers: Providers,
+) -> Callable[[str], Awaitable[str | None]]:
+    if "tavily" in providers:
+        return tavily_search_provider(providers.get("tavily", None))
+
+    if "google" in providers:
+        return google_search_provider(providers.get("google", None))
+
+    raise ValueError("No valid provider found.")

inspect_ai/util/__init__.py

@@ -6,7 +6,9 @@ from inspect_ai.util._limit import (
     LimitScope,
     apply_limits,
     message_limit,
+    time_limit,
     token_limit,
+    working_limit,
 )
 
 from ._collect import collect
@@ -81,6 +83,8 @@ __all__ = [
     "subtask",
     "throttle",
     "token_limit",
+    "time_limit",
+    "working_limit",
     "trace_action",
     "trace_message",
     "RegistryType",

inspect_ai/util/_json.py

@@ -3,6 +3,7 @@ import typing
 from copy import deepcopy
 from dataclasses import is_dataclass
 from datetime import date, datetime, time
+from enum import EnumMeta
 from typing import (
     Any,
     Dict,
@@ -101,6 +102,8 @@ def json_schema(t: Type[Any]) -> JSONSchema:
             or (isinstance(t, type) and issubclass(t, BaseModel))
         ):
             return cls_json_schema(t)
+        elif isinstance(t, EnumMeta):
+            return JSONSchema(enum=[item.value for item in t])
         elif t is type(None):
             return JSONSchema(type="null")
         else:

inspect_ai/util/_limit.py

@@ -7,6 +7,7 @@ from contextvars import ContextVar
 from types import TracebackType
 from typing import TYPE_CHECKING, Generic, Iterator, Literal, TypeVar
 
+import anyio
 from typing_extensions import Self
 
 from inspect_ai._util.logger import warn_once
@@ -33,22 +34,23 @@ class LimitExceededError(Exception):
        value: Value compared to.
        limit: Limit applied.
        message (str | None): Optional. Human readable message.
-       source (Limit | None): Optional. The `Limit` instance which was responsible for
-         raising this error.
+       source (Limit | None): Optional. The `Limit` instance which was responsible for raising this error.
     """
 
     def __init__(
         self,
         type: Literal["message", "time", "working", "token", "operator", "custom"],
         *,
-        value: int,
-        limit: int,
+        value: float,
+        limit: float,
         message: str | None = None,
         source: Limit | None = None,
     ) -> None:
         self.type = type
         self.value = value
+        self.value_str = self._format_float_or_int(value)
         self.limit = limit
+        self.limit_str = self._format_float_or_int(limit)
         self.message = f"Exceeded {type} limit: {limit:,}"
         self.source = source
         super().__init__(message)
@@ -60,6 +62,12 @@ class LimitExceededError(Exception):
         )
         return self
 
+    def _format_float_or_int(self, value: float | int) -> str:
+        if isinstance(value, int):
+            return f"{value:,}"
+        else:
+            return f"{value:,.2f}"
+
 
 class Limit(abc.ABC):
     """Base class for all limit context managers."""
@@ -80,6 +88,12 @@ class Limit(abc.ABC):
     ) -> None:
         pass
 
+    @property
+    @abc.abstractmethod
+    def usage(self) -> float:
+        """The current usage of the resource being limited."""
+        pass
+
     def _check_reuse(self) -> None:
         if self._entered:
             raise RuntimeError(
@@ -112,18 +126,20 @@ def apply_limits(
         False, all `LimitExceededError` exceptions will be allowed to propagate.
     """
     limit_scope = LimitScope()
-    with ExitStack() as stack:
-        for limit in limits:
-            stack.enter_context(limit)
-        try:
+    # Try scope is outside the `with ExitStack()` so that we can catch any errors raised
+    # when exiting it (which will be where time_limit() would raise LimitExceededError).
+    try:
+        with ExitStack() as stack:
+            for limit in limits:
+                stack.enter_context(limit)
             yield limit_scope
-        except LimitExceededError as e:
-            # If it was not one of the limits we applied.
-            if e.source is None or e.source not in limits:
-                raise
-            limit_scope.limit_error = e
-            if not catch_errors:
-                raise
+    except LimitExceededError as e:
+        # If it was not one of the limits we applied.
+        if e.source is None or e.source not in limits:
+            raise
+        limit_scope.limit_error = e
+        if not catch_errors:
+            raise
 
 
 class LimitScope:
@@ -140,8 +156,6 @@ def token_limit(limit: int | None) -> _TokenLimit:
     """Limits the total number of tokens which can be used.
 
     The counter starts when the context manager is opened and ends when it is closed.
-    The context manager can be opened multiple times, even in different execution
-    contexts.
 
     These limits can be stacked.
 
@@ -186,8 +200,7 @@ def message_limit(limit: int | None) -> _MessageLimit:
     """Limits the number of messages in a conversation.
 
     The total number of messages in the conversation are compared to the limit (not just
-    "new" messages). The context manager can be opened multiple times, even in different
-    execution contexts.
+    "new" messages).
 
     These limits can be stacked.
 
@@ -220,6 +233,62 @@ def check_message_limit(count: int, raise_for_equal: bool) -> None:
     node.check(count, raise_for_equal)
 
 
+def time_limit(limit: float | None) -> _TimeLimit:
+    """Limits the wall clock time which can elapse.
+
+    The timer starts when the context manager is opened and stops when it is closed.
+
+    These limits can be stacked.
+
+    When a limit is exceeded, the code block is cancelled and a `LimitExceededError` is
+    raised.
+
+    Uses anyio's cancellation scopes meaning that the operations within the context
+    manager block are cancelled if the limit is exceeded. The `LimitExceededError` is
+    therefore raised at the level that the `time_limit()` context manager was opened,
+    not at the level of the operation which caused the limit to be exceeded (e.g. a call
+    to `generate()`). Ensure you handle `LimitExceededError` at the level of opening the context manager.
+
+    Args:
+      limit: The maximum number of seconds that can pass while the context manager is
+        open. A value of None means unlimited time.
+    """
+    return _TimeLimit(limit)
+
+
+def working_limit(limit: float | None) -> _WorkingLimit:
+    """Limits the working time which can elapse.
+
+    Working time is the wall clock time minus any waiting time e.g. waiting before
+    retrying in response to rate limits or waiting on a semaphore.
+
+    The timer starts when the context manager is opened and stops when it is closed.
+
+    These limits can be stacked.
+
+    When a limit is exceeded, a `LimitExceededError` is raised.
+
+    Args:
+      limit: The maximum number of seconds of working that can pass while the context
+        manager is open. A value of None means unlimited time.
+    """
+    return _WorkingLimit(limit)
+
+
+def record_waiting_time(waiting_time: float) -> None:
+    node = working_limit_tree.get()
+    if node is None:
+        return
+    node.record_waiting_time(waiting_time)
+
+
+def check_working_limit() -> None:
+    node = working_limit_tree.get()
+    if node is None:
+        return
+    node.check()
+
+
 class _Tree(Generic[TNode]):
     """A tree data structure of limit nodes.
 
@@ -253,6 +322,7 @@ token_limit_tree: _Tree[_TokenLimit] = _Tree("token_limit_tree")
 # Store the message limit leaf node so that we know which limit to check in
 # check_message_limit().
 message_limit_tree: _Tree[_MessageLimit] = _Tree("message_limit_tree")
+working_limit_tree: _Tree[_WorkingLimit] = _Tree("working_limit_tree")
 
 
 class _Node:
@@ -296,6 +366,10 @@ class _TokenLimit(Limit, _Node):
     ) -> None:
         self._pop_and_check_identity(token_limit_tree)
 
+    @property
+    def usage(self) -> float:
+        return self._usage.total_tokens
+
     @property
     def limit(self) -> int | None:
         """Get the configured token limit value."""
@@ -312,7 +386,7 @@ class _TokenLimit(Limit, _Node):
         self._limit = value
 
     def record(self, usage: ModelUsage) -> None:
-        """Record model usage for this node and its parent nodes."""
+        """Record model usage for this node and its ancestor nodes."""
         if self.parent is not None:
             self.parent.record(usage)
         self._usage += usage
@@ -369,6 +443,13 @@ class _MessageLimit(Limit, _Node):
     ) -> None:
         self._pop_and_check_identity(message_limit_tree)
 
+    @property
+    def usage(self) -> float:
+        raise NotImplementedError(
+            "Retrieving the message count from a limit is not supported. Please query "
+            "the messages property on the task or agent state instead."
+        )
+
     @property
     def limit(self) -> int | None:
         """Get the configured message limit value."""
@@ -414,3 +495,132 @@ class _MessageLimit(Limit, _Node):
             raise ValueError(
                 f"Message limit value must be a non-negative integer or None: {value}"
             )
+
+
+class _TimeLimit(Limit):
+    def __init__(self, limit: float | None) -> None:
+        super().__init__()
+        _validate_time_limit("Time", limit)
+        self._limit = limit
+        self._start_time: float | None = None
+        self._end_time: float | None = None
+
+    def __enter__(self) -> Limit:
+        super()._check_reuse()
+        # Unlike the other limits, this one is not stored in a tree. Anyio handles all
+        # of the state.
+        self._cancel_scope = anyio.move_on_after(self._limit)
+        self._cancel_scope.__enter__()
+        self._start_time = anyio.current_time()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        from inspect_ai.log._transcript import SampleLimitEvent, transcript
+
+        self._cancel_scope.__exit__(exc_type, exc_val, exc_tb)
+        self._end_time = anyio.current_time()
+        if self._cancel_scope.cancel_called and self._limit is not None:
+            message = f"Time limit exceeded. limit: {self._limit} seconds"
+            assert self._start_time is not None
+            # Note we've measured the elapsed time independently of anyio's cancel scope
+            # so this is an approximation.
+            time_elapsed = self._end_time - self._start_time
+            transcript()._event(
+                SampleLimitEvent(type="time", message=message, limit=self._limit)
+            )
+            raise LimitExceededError(
+                "time",
+                value=time_elapsed,
+                limit=self._limit,
+                message=message,
+                source=self,
+            ) from exc_val
+
+    @property
+    def usage(self) -> float:
+        if self._start_time is None:
+            return 0.0
+        if self._end_time is None:
+            return anyio.current_time() - self._start_time
+        return self._end_time - self._start_time
+
+
+class _WorkingLimit(Limit, _Node):
+    def __init__(self, limit: float | None) -> None:
+        super().__init__()
+        _validate_time_limit("Working time", limit)
+        self._limit = limit
+        self.parent: _WorkingLimit | None = None
+        self._start_time: float | None = None
+        self._end_time: float | None = None
+
+    def __enter__(self) -> Limit:
+        super()._check_reuse()
+        self._start_time = anyio.current_time()
+        self._waiting_time = 0.0
+        working_limit_tree.push(self)
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        self._end_time = anyio.current_time()
+        self._pop_and_check_identity(working_limit_tree)
+
+    @property
+    def usage(self) -> float:
+        if self._start_time is None:
+            return 0.0
+        if self._end_time is None:
+            return anyio.current_time() - self._start_time - self._waiting_time
+        return self._end_time - self._start_time - self._waiting_time
+
+    def record_waiting_time(self, waiting_time: float) -> None:
+        """Record waiting time for this node and its ancestor nodes."""
+        if self.parent is not None:
+            self.parent.record_waiting_time(waiting_time)
+        self._waiting_time += waiting_time
+
+    def check(self) -> None:
+        """Check if this working time limit or any ancestor limits have been exceeded.
+
+        The checks occur from root to leaf. This is so that if multiple limits are
+        simultaneously exceeded, the outermost (closest to root) one raises the error,
+        preventing certain sub-agent architectures from ending up in an infinite loop.
+        """
+        if self.parent is not None:
+            self.parent.check()
+        self._check_self()
+
+    def _check_self(self) -> None:
+        from inspect_ai.log._transcript import SampleLimitEvent, transcript
+
+        if self._limit is None:
+            return
+        if self.usage > self._limit:
+            message = f"Working time limit exceeded. limit: {self._limit} seconds"
+            transcript()._event(
+                SampleLimitEvent(type="working", message=message, limit=self._limit)
+            )
+            raise LimitExceededError(
+                "working",
+                value=self.usage,
+                limit=self._limit,
+                message=message,
+                source=self,
+            )
+
+
+def _validate_time_limit(name: str, value: float | None) -> None:
+    if value is not None and value < 0:
+        raise ValueError(
+            f"{name} limit value must be a non-negative float or None: {value}"
+        )