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

inspect_ai/tool/_tool.py

@@ -224,13 +224,15 @@ def tool(
             tool_parallel = parallel
             tool_viewer = viewer
             tool_model_input = model_input
+            tool_options: dict[str, object] | None = None
             if is_registry_object(tool):
-                _, _, reg_parallel, reg_viewer, reg_model_input = tool_registry_info(
-                    tool
+                _, _, reg_parallel, reg_viewer, reg_model_input, options = (
+                    tool_registry_info(tool)
                 )
                 tool_parallel = parallel and reg_parallel
                 tool_viewer = viewer or reg_viewer
                 tool_model_input = model_input or reg_model_input
+                tool_options = options
 
             # tag the object
             registry_tag(
@@ -247,6 +249,7 @@ def tool(
                             tool_model_input
                             or getattr(tool, TOOL_INIT_MODEL_INPUT, None)
                         ),
+                        TOOL_OPTIONS: tool_options,
                     },
                 ),
                 *args,
@@ -267,6 +270,7 @@ TOOL_PROMPT = "prompt"
 TOOL_PARALLEL = "parallel"
 TOOL_VIEWER = "viewer"
 TOOL_MODEL_INPUT = "model_input"
+TOOL_OPTIONS = "options"
 
 
 TOOL_INIT_MODEL_INPUT = "__TOOL_INIT_MODEL_INPUT__"

inspect_ai/tool/_tool_def.py

@@ -16,6 +16,7 @@ from inspect_ai._util.registry import (
 
 from ._tool import (
     TOOL_MODEL_INPUT,
+    TOOL_OPTIONS,
     TOOL_PARALLEL,
     TOOL_PROMPT,
     TOOL_VIEWER,
@@ -44,6 +45,7 @@ class ToolDef:
         parallel: bool | None = None,
         viewer: ToolCallViewer | None = None,
         model_input: ToolCallModelInput | None = None,
+        options: dict[str, object] | None = None,
     ) -> None:
         """Create a tool definition.
 
@@ -59,6 +61,8 @@ class ToolDef:
           viewer: Optional tool call viewer implementation.
           model_input: Optional function that determines how
               tool call results are played back as model input.
+          options: Optional property bag that can be used by the model provider
+              to customize the implementation of the tool
 
         Returns:
           Tool definition.
@@ -82,6 +86,7 @@ class ToolDef:
             self.parallel = parallel if parallel is not None else tdef.parallel
             self.viewer = viewer or tdef.viewer
             self.model_input = model_input or tdef.model_input
+            self.options = options or tdef.options
 
         # if its not a tool then extract tool_info if all fields have not
         # been provided explicitly
@@ -112,6 +117,7 @@ class ToolDef:
             self.parallel = parallel is not False
             self.viewer = viewer
             self.model_input = model_input
+            self.options = options
 
     tool: Callable[..., Any]
     """Callable to execute tool."""
@@ -134,13 +140,20 @@ class ToolDef:
     model_input: ToolCallModelInput | None
     """Custom model input presenter for tool calls."""
 
+    options: dict[str, object] | None = None
+    """Optional property bag that can be used by the model provider to customize the implementation of the tool"""
+
     def as_tool(self) -> Tool:
         """Convert a ToolDef to a Tool."""
         tool = self.tool
         info = RegistryInfo(
             type="tool",
             name=self.name,
-            metadata={TOOL_PARALLEL: self.parallel, TOOL_VIEWER: self.viewer},
+            metadata={
+                TOOL_PARALLEL: self.parallel,
+                TOOL_VIEWER: self.viewer,
+                TOOL_OPTIONS: self.options,
+            },
         )
         set_registry_info(tool, info)
         set_registry_params(tool, {})
@@ -189,11 +202,12 @@ class ToolDefFields(NamedTuple):
     parallel: bool
     viewer: ToolCallViewer | None
     model_input: ToolCallModelInput | None
+    options: dict[str, object] | None
 
 
 def tool_def_fields(tool: Tool) -> ToolDefFields:
     # get tool_info
-    name, prompt, parallel, viewer, model_input = tool_registry_info(tool)
+    name, prompt, parallel, viewer, model_input, options = tool_registry_info(tool)
     tool_info = parse_tool_info(tool)
 
     # if there is a description then append any prompt to the
@@ -234,19 +248,28 @@ def tool_def_fields(tool: Tool) -> ToolDefFields:
         parallel=parallel,
         viewer=viewer,
         model_input=model_input,
+        options=options,
     )
 
 
 def tool_registry_info(
     tool: Tool,
-) -> tuple[str, str | None, bool, ToolCallViewer | None, ToolCallModelInput | None]:
+) -> tuple[
+    str,
+    str | None,
+    bool,
+    ToolCallViewer | None,
+    ToolCallModelInput | None,
+    dict[str, object] | None,
+]:
     info = registry_info(tool)
     name = info.name.split("/")[-1]
     prompt = info.metadata.get(TOOL_PROMPT, None)
     parallel = info.metadata.get(TOOL_PARALLEL, True)
     viewer = info.metadata.get(TOOL_VIEWER, None)
     model_input = info.metadata.get(TOOL_MODEL_INPUT, None)
-    return name, prompt, parallel, viewer, model_input
+    options = info.metadata.get(TOOL_OPTIONS, None)
+    return name, prompt, parallel, viewer, model_input, options
 
 
 def validate_tool_parameters(tool_name: str, parameters: dict[str, ToolParam]) -> None:

inspect_ai/tool/_tool_info.py

@@ -49,6 +49,8 @@ class ToolInfo(BaseModel):
     """Short description of tool."""
     parameters: ToolParams = Field(default_factory=ToolParams)
     """JSON Schema of tool parameters object."""
+    options: dict[str, object] | None = Field(default=None)
+    """Optional property bag that can be used by the model provider to customize the implementation of the tool"""
 
 
 def parse_tool_info(func: Callable[..., Any]) -> ToolInfo:

inspect_ai/tool/_tools/_web_search/_google.py

@@ -4,6 +4,7 @@ from typing import Awaitable, Callable
 import anyio
 import httpx
 from bs4 import BeautifulSoup, NavigableString
+from pydantic import BaseModel
 from tenacity import (
     retry,
     retry_if_exception,
@@ -23,10 +24,18 @@ Page Content: {text}
 """
 
 
+class GoogleOptions(BaseModel):
+    num_results: int | None = None
+    max_provider_calls: int | None = None
+    max_connections: int | None = None
+    model: str | None = None
+
+
 class SearchLink:
-    def __init__(self, url: str, snippet: str) -> None:
+    def __init__(self, url: str, snippet: str, title: str) -> None:
         self.url = url
         self.snippet = snippet
+        self.title = title
 
 
 def maybe_get_google_api_keys() -> tuple[str, str] | None:
@@ -42,11 +51,14 @@ def maybe_get_google_api_keys() -> tuple[str, str] | None:
 
 
 def google_search_provider(
-    num_results: int,
-    max_provider_calls: int,
-    max_connections: int,
-    model: str | None,
+    in_options: dict[str, object] | None = None,
 ) -> Callable[[str], Awaitable[str | None]]:
+    options = GoogleOptions.model_validate(in_options) if in_options else None
+    num_results = (options.num_results if options else None) or 3
+    max_provider_calls = (options.max_provider_calls if options else None) or 3
+    max_connections = (options.max_connections if options else None) or 10
+    model = options.model if options else None
+
     keys = maybe_get_google_api_keys()
     if not keys:
         raise PrerequisiteError(
@@ -60,8 +72,7 @@ def google_search_provider(
     async def search(query: str) -> str | None:
         # limit number of concurrent searches
         page_contents: list[str] = []
-        urls: list[str] = []
-        snippets: list[str] = []
+        processed_links: list[SearchLink] = []
         search_calls = 0
 
         # Paginate through search results until we have successfully extracted num_results pages or we have reached max_provider_calls
@@ -76,8 +87,7 @@ def google_search_provider(
                         page = await page_if_relevant(link.url, query, model, client)
                         if page:
                             page_contents.append(page)
-                            urls.append(link.url)
-                            snippets.append(link.snippet)
+                            processed_links.append(link)
                     # exceptions fetching pages are very common!
                     except Exception:
                         pass
@@ -87,8 +97,18 @@ def google_search_provider(
 
             search_calls += 1
 
-        all_page_contents = "\n\n".join(page_contents)
-        return None if all_page_contents == "" else all_page_contents
+        return (
+            "\n\n".join(
+                "[{title}]({url}):\n{page_content}".format(
+                    title=link.title, url=link.url, page_content=page_content
+                )
+                for link, page_content in zip(
+                    processed_links, page_contents, strict=True
+                )
+            )
+            if processed_links
+            else None
+        )
 
     async def _search(query: str, start_idx: int) -> list[SearchLink]:
         # List of allowed parameters can be found https://developers.google.com/custom-search/v1/reference/rest/v1/cse/list
@@ -110,13 +130,21 @@ def google_search_provider(
             before_sleep=log_httpx_retry_attempt(search_url),
         )
         async def execute_search() -> httpx.Response:
+            # See https://developers.google.com/custom-search/v1/reference/rest/v1/Search
             return await client.get(search_url)
 
         result = await execute_search()
         data = result.json()
 
         if "items" in data:
-            return [SearchLink(item["link"], item["snippet"]) for item in data["items"]]
+            return [
+                SearchLink(
+                    url=item["link"],
+                    snippet=item.get("snippet", ""),  # sometimes not present
+                    title=item["title"],
+                )
+                for item in data["items"]
+            ]
         else:
             return []
 
@@ -124,13 +152,13 @@ def google_search_provider(
 
 
 async def page_if_relevant(
-    link: str, query: str, relevance_model: str | None, client: httpx.AsyncClient
+    url: str, query: str, relevance_model: str | None, client: httpx.AsyncClient
 ) -> str | None:
     """
     Use parser model to determine if a web page contents is relevant to a query.
 
     Args:
-        link (str): Web page link.
+        url (str): Web page url.
         query (str): Search query.
         relevance_model (Model): Model used to parse web pages for relevance.
         client: (httpx.Client): HTTP client to use to fetch the page
@@ -145,7 +173,7 @@ async def page_if_relevant(
 
     # retrieve document
     try:
-        response = await client.get(link)
+        response = await client.get(url)
         response.raise_for_status()
     except httpx.HTTPError as exc:
         raise Exception(f"HTTP error occurred: {exc}")

inspect_ai/tool/_tools/_web_search/_tavily.py

@@ -1,5 +1,5 @@
 import os
-from typing import Awaitable, Callable
+from typing import Awaitable, Callable, Literal
 
 import httpx
 from pydantic import BaseModel, Field
@@ -16,6 +16,25 @@ from inspect_ai._util.httpx import httpx_should_retry, log_httpx_retry_attempt
 from inspect_ai.util._concurrency import concurrency
 
 
+class TavilyOptions(BaseModel):
+    topic: Literal["general", "news"] | None = None
+    search_depth: Literal["basic", "advanced"] | None = None
+    chunks_per_source: Literal[1, 2, 3] | None = None
+    max_results: int | None = None
+    time_range: Literal["day", "week", "month", "year", "d", "w", "m", "y"] | None = (
+        None
+    )
+    days: int | None = None
+    include_answer: bool | Literal["basic", "advanced"] | None = None
+    include_raw_content: bool | None = None
+    include_images: bool | None = None
+    include_image_descriptions: bool | None = None
+    include_domains: list[str] | None = None
+    exclude_domains: list[str] | None = None
+    # max_connections is not a Tavily API option, but an inspect option
+    max_connections: int | None = None
+
+
 class TavilySearchResult(BaseModel):
     title: str
     url: str
@@ -32,32 +51,37 @@ class TavilySearchResponse(BaseModel):
 
 
 def tavily_search_provider(
-    num_results: int, max_connections: int
+    in_options: dict[str, object] | None = None,
 ) -> Callable[[str], Awaitable[str | None]]:
+    options = TavilyOptions.model_validate(in_options) if in_options else None
+    # Separate max_connections (which is an inspect thing) from the rest of the
+    # options which will be passed in the request body
+    max_connections = (options.max_connections if options else None) or 10
+    api_options = (
+        options.model_dump(exclude={"max_connections"}, exclude_none=True)
+        if options
+        else {}
+    )
+    if not api_options.get("include_answer", False):
+        api_options["include_answer"] = True
+
     tavily_api_key = os.environ.get("TAVILY_API_KEY", None)
     if not tavily_api_key:
         raise PrerequisiteError(
             "TAVILY_API_KEY not set in the environment. Please ensure ths variable is defined to use Tavily with the web_search tool.\n\nLearn more about the Tavily web search provider at https://inspect.aisi.org.uk/tools.html#tavily-provider"
         )
-    if num_results > 20:
-        raise PrerequisiteError(
-            "The Tavily search provider is limited to 20 results per query."
-        )
 
     # Create the client within the provider
     client = httpx.AsyncClient(timeout=30)
 
     async def search(query: str) -> str | None:
+        # See https://docs.tavily.com/documentation/api-reference/endpoint/search
         search_url = "https://api.tavily.com/search"
         headers = {
             "Authorization": f"Bearer {tavily_api_key}",
         }
-        body = {
-            "query": query,
-            "max_results": 10,  # num_results,
-            # "search_depth": "advanced",
-            "include_answer": "advanced",
-        }
+
+        body = {"query": query, **api_options}
 
         # retry up to 5 times over a period of up to 1 minute
         @retry(
@@ -72,6 +96,15 @@ def tavily_search_provider(
             return response
 
         async with concurrency("tavily_web_search", max_connections):
-            return TavilySearchResponse.model_validate((await _search()).json()).answer
+            tavily_search_response = TavilySearchResponse.model_validate(
+                (await _search()).json()
+            )
+            results_str = "\n\n".join(
+                [
+                    f"[{result.title}]({result.url}):\n{result.content}"
+                    for result in tavily_search_response.results
+                ]
+            )
+            return f"Answer: {tavily_search_response.answer}\n\n{results_str}"
 
     return search

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.")