furyhawk/pydantic-ai-harness
1
0
Fork 0
mirror of https://github.com/pydantic/pydantic-ai-harness.git synced 2026-07-21 10:55:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add opt-in deep_search tool and harden ExaSearch per review

Browse Source 
deep_search runs Exa's multi-step deep search (search type='deep' with a
text output schema) and returns a synthesized, cited answer in one tool
call -- unlike the start/check polling pair Exa's MCP server exposes. It
is off by default because deep calls are markedly slower and pricier and
the model decides when to invoke tools, so the spend is opt-in.

Review hardening: get_page retries on a result with no text, non-positive
num_results/max_text_chars are rejected at construction, and num_results
is re-applied to oversized responses so a custom client cannot exceed the
bounded-output contract. Docs no longer promise output "fits" the model
context, and now position ExaSearch against Exa's hosted MCP server.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Bill Easton
2026-07-16 15:15:48 -05:00
co-authored by Claude Fable 5
parent 1b0d1c3edc
commit 02540a2409
7 changed files with 402 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1 -1
View File
@@ -151,7 +151,7 @@ We studied leading coding agents, agent frameworks, and Claw-style assistants to
| | **Shell** | Execute commands with allowlists, denylists, and timeouts | :white_check_mark: [Docs](pydantic_ai_harness/shell/) | [pydantic-ai-backend](https://github.com/vstorm-co/pydantic-ai-backend) (vstorm&#8209;co) |
| | **Repo context injection** | Auto-load CLAUDE.md/AGENTS.md and repo structure | :white_check_mark: [Docs](pydantic_ai_harness/context/) | [pydantic-deep](https://github.com/vstorm-co/pydantic-deepagents) (vstorm&#8209;co) |
| | **Docs lookup** | On-demand `read_pyai_docs` tool for Pydantic AI docs | :white_check_mark: [Docs](pydantic_ai_harness/docs/) | |
| | **Web research** | Web search returning full page text, plus single-page retrieval, backed by [Exa](https://exa.ai) | :white_check_mark: [Docs](pydantic_ai_harness/exa/) | |
| | **Web research** | Web search returning full page text, single-page retrieval, and opt-in deep search with cited answers, backed by [Exa](https://exa.ai) | :white_check_mark: [Docs](pydantic_ai_harness/exa/) | |
| | **Verification loop** | Run tests after edits, auto-fix failures | :construction: [PR&nbsp;#169](https://github.com/pydantic/pydantic-ai-harness/pull/169) | |
| **Editor integration** | **ACP** | Serve an agent to editors (Zed, etc.) over the [Agent Client Protocol](https://agentclientprotocol.com) -- streamed text, diff-rendered edits, tool approval | :white_check_mark: [Docs](pydantic_ai_harness/experimental/acp/) (experimental) | |
| **Prompt management** | **Managed prompt** | Back an agent's instructions with a [Logfire](https://pydantic.dev/logfire)-managed prompt, editable without shipping code | :white_check_mark: [Docs](pydantic_ai_harness/logfire/) | |
docs/exa-search.md
+63 -15
View File
@@ -1,14 +1,15 @@
---
title: Exa Search
description: Give a Pydantic AI agent web research tools backed by the Exa search API -- search with page text and full-page retrieval, with output capped to fit model context.
description: Give a Pydantic AI agent web research tools backed by the Exa search API -- search with page text, full-page retrieval, and opt-in deep search, with page text capped per result to bound tool output.
---
# Exa Search
`ExaSearch` gives an agent web research tools backed by the
[Exa](https://exa.ai) search API: search that returns page text alongside each
hit, and full-page retrieval for digging into a specific URL. Page text is
capped per result so tool output fits the model's context.
hit, full-page retrieval for digging into a specific URL, and opt-in deep
search that synthesizes a cited answer in one call. Page text is capped per
result to bound tool output.
[Source](https://github.com/pydantic/pydantic-ai-harness/tree/main/pydantic_ai_harness/exa/)
@@ -21,7 +22,7 @@ context, and prompting the agent to research methodically is boilerplate every
research agent reinvents.
`ExaSearch` bundles that plumbing into a single
[capability](/ai/core-concepts/capabilities/): two research tools, a per-result
[capability](/ai/core-concepts/capabilities/): the research tools, a per-result
page-text cap, and short research guidance in the system prompt.
## Usage
@@ -47,30 +48,50 @@ print(result.output)
## Tools
`ExaSearch` contributes two tools to the agent:
`ExaSearch` contributes these tools to the agent:
| Tool | Purpose |
|---|---|
| `web_search` | Search the web and return the top `num_results` pages, each with title, URL, and page text. |
| `get_page` | Retrieve the text contents of one specific URL. |
| `deep_search` | Run Exa's multi-step deep search and return a synthesized, cited answer. Opt-in via `include_deep_search=True`. |
Page text is capped at `max_text_chars` characters per result: the cap is sent
to Exa as the contents limit and re-enforced when tool output is formatted, so
text stays bounded even with a custom client. When text is cut, the **head** is
kept (a page's lead carries the substance) and a
`[... page text truncated at N characters]` marker is appended.
text stays bounded even with a custom client. The result count is bounded the
same way (`num_results` is requested from Exa and re-applied to the response).
When text is cut, the **head** is kept (a page's lead carries the substance)
and a `[... page text truncated at N characters]` marker is appended.
A URL that returns no content surfaces to the model as a
A URL or question that returns no content surfaces to the model as a
[`ModelRetry`](/ai/tools-toolsets/tools-advanced/#tool-retries) rather than a
hard error: the run continues and the model can correct the URL or pick another
page.
hard error: the run continues and the model can correct the URL or rephrase.
## Deep search
`deep_search` calls Exa search with `type='deep'` and a plain-text output
schema: Exa expands the question into multiple queries, searches, and returns
an answer grounded in citations -- all in **one tool call**, with the cited
sources listed under the answer. It is markedly slower and more expensive per
call than `web_search`, and the model decides when to invoke tools, so the
tool is off by default -- enable it explicitly:
```python
from pydantic_ai_harness.exa import ExaSearch
ExaSearch(include_deep_search=True)
```
When enabled, the capability's instructions tell the model to treat it as an
escalation from `web_search`, not a replacement.
## Instructions
`ExaSearch` contributes short research guidance to the system prompt: search
wide with `web_search` first, read the most promising pages in full with
`get_page` before drawing conclusions, prefer primary sources, and cite the
URLs relied on.
URLs relied on. With `include_deep_search=True`, the guidance also covers when
to escalate to `deep_search`.
## Configuration
@@ -80,9 +101,10 @@ Every field of `ExaSearch` with its default:
from pydantic_ai_harness.exa import ExaSearch
ExaSearch(
num_results=5, # results per web_search call
max_text_chars=10_000, # page-text cap per result, in characters
client=None, # ExaClient -- None builds exa_py.AsyncExa from EXA_API_KEY
num_results=5, # results per web_search call
max_text_chars=10_000, # page-text cap per result, in characters
include_deep_search=False, # also expose the deep_search tool
client=None, # ExaClient -- None builds exa_py.AsyncExa from EXA_API_KEY
)
```
@@ -105,6 +127,32 @@ ExaSearch(client=AsyncExa(api_key='...'))
The API may change between releases while the capability settles; breaking
changes ship deprecation warnings where practical.
## ExaSearch vs Exa's MCP server
Exa also ships an official hosted MCP server at `https://mcp.exa.ai/mcp`
([exa-labs/exa-mcp-server](https://github.com/exa-labs/exa-mcp-server)) with
their full tool catalog: `web_search_exa`, `get_code_context_exa`,
`crawling_exa`, `company_research_exa`, `linkedin_search_exa`, and a
`deep_researcher_start` / `deep_researcher_check` pair.
`ExaSearch` is the curated, typed path: bounded output, a retry-on-empty
contract, bundled research instructions, and a client seam that makes it
testable offline. The MCP server is how you get Exa's full catalog with zero
wrapper code, via Pydantic AI core's
[MCP capability](/ai/core-concepts/capabilities/). Where the MCP deep
researcher is a start/check tool pair the model must poll across calls,
`deep_search` returns the answer in a single call. The two compose in one
`capabilities` list, and the tool names don't collide (the MCP tools are
`*_exa`-suffixed):
```python
from pydantic_ai import Agent
from pydantic_ai.capabilities import MCP
from pydantic_ai_harness.exa import ExaSearch
agent = Agent('anthropic:claude-sonnet-4-6', capabilities=[ExaSearch(), MCP('https://mcp.exa.ai/mcp')])
```
## Further reading
- [Pydantic AI capabilities](/ai/core-concepts/capabilities/)
docs/index.md
+1 -1
View File
@@ -117,7 +117,7 @@ Each capability is a self-contained battery you drop into an agent's `capabiliti
| [Shell](shell.md) | Command execution in a subprocess rooted at a working directory, gated by allowlists, denylists, timeouts, and optional environment-variable stripping (including a preset for common LLM provider credentials). | -- |
| [Context](context.md) | Auto-loads repo context -- `CLAUDE.md`/`AGENTS.md` and repository structure -- so the agent starts a run already oriented in the project. | -- |
| [Pydantic AI Docs](pydantic-ai-docs.md) | An on-demand `read_pyai_docs` tool that pulls Pydantic AI documentation into the run when the agent needs it, instead of preloading it. | -- |
| [Exa Search](exa-search.md) | Web research backed by the [Exa](https://exa.ai) search API: `web_search` returns results with page text and `get_page` reads a specific URL, with page text capped per result so tool output fits model context. | `exa` |
| [Exa Search](exa-search.md) | Web research backed by the [Exa](https://exa.ai) search API: `web_search` returns results with page text, `get_page` reads a specific URL, and opt-in `deep_search` synthesizes a cited answer in one call. Page text is capped per result to bound tool output. | `exa` |
| [Compaction](compaction.md) | Keeps a run within token limits: sliding-window trimming, LLM-powered summarization of older messages, and warnings before the context or iteration ceiling is hit. | -- |
| [Overflowing Tool Output](overflowing-tool-output.md) | Reduces an oversized tool return when it is produced -- truncate, spill to a queryable file, or summarize -- so a large payload does not persist in history and get re-sent every request. | -- |
| [Step Persistence](step-persistence.md) | Saves and restores full conversation state; snapshot, resume (`continue_run`), and fork (`fork_run`) a run. | -- |
pydantic_ai_harness/exa/README.md
+61 -12
View File
@@ -1,8 +1,9 @@
# Exa Search
Give an agent web research tools backed by the [Exa](https://exa.ai) search
API: search that returns page text alongside each hit, and full-page retrieval
for digging into a specific URL.
API: search that returns page text alongside each hit, full-page retrieval for
digging into a specific URL, and opt-in deep search that synthesizes a cited
answer in one call.
[Source](https://github.com/pydantic/pydantic-ai-harness/tree/main/pydantic_ai_harness/exa/)
@@ -26,7 +27,7 @@ research agent reinvents.
## The solution
`ExaSearch` bundles both tools with output capping and short research guidance
`ExaSearch` bundles the tools with output capping and short research guidance
in the system prompt:
```python
@@ -45,22 +46,44 @@ print(result.output)
|---|---|
| `web_search` | Search the web and return the top `num_results` pages, each with title, URL, and page text. |
| `get_page` | Retrieve the text contents of one specific URL. |
| `deep_search` | Run Exa's multi-step deep search and return a synthesized, cited answer. Opt-in via `include_deep_search=True`. |
Page text is capped at `max_text_chars` characters per result: the cap is sent
to Exa as the contents limit and re-enforced when tool output is formatted, so
text stays bounded even with a custom client. When text is cut, the **head** is
kept (a page's lead carries the substance) and a
`[... page text truncated at N characters]` marker is appended.
text stays bounded even with a custom client. The result count is bounded the
same way (`num_results` is requested from Exa and re-applied to the response).
When text is cut, the **head** is kept (a page's lead carries the substance)
and a `[... page text truncated at N characters]` marker is appended.
A URL that returns no content surfaces to the model as a `ModelRetry` (the
model can correct the URL or pick another page) rather than aborting the run.
A URL or question that returns no content surfaces to the model as a
`ModelRetry` (the model can correct the URL or rephrase) rather than aborting
the run.
## Deep search
`deep_search` calls Exa search with `type='deep'` and a plain-text output
schema: Exa expands the question into multiple queries, searches, and returns
an answer grounded in citations -- all in **one tool call**, with the cited
sources listed under the answer. It is markedly slower and more expensive per
call than `web_search`, and the model decides when to invoke tools, so the
tool is off by default -- enable it explicitly:
```python
from pydantic_ai_harness.exa import ExaSearch
ExaSearch(include_deep_search=True)
```
When enabled, the capability's instructions tell the model to treat it as an
escalation from `web_search`, not a replacement.
## Instructions
`ExaSearch` contributes short research guidance to the system prompt: search
wide with `web_search` first, read the most promising pages in full with
`get_page` before drawing conclusions, prefer primary sources, and cite the
URLs relied on.
URLs relied on. With `include_deep_search=True`, the guidance also covers when
to escalate to `deep_search`.
## Configuration
@@ -70,9 +93,10 @@ Every field of `ExaSearch` with its default:
from pydantic_ai_harness.exa import ExaSearch
ExaSearch(
num_results=5, # results per web_search call
max_text_chars=10_000, # page-text cap per result, in characters
client=None, # ExaClient -- None builds exa_py.AsyncExa from EXA_API_KEY
num_results=5, # results per web_search call
max_text_chars=10_000, # page-text cap per result, in characters
include_deep_search=False, # also expose the deep_search tool
client=None, # ExaClient -- None builds exa_py.AsyncExa from EXA_API_KEY
)
```
@@ -95,6 +119,31 @@ ExaSearch(client=AsyncExa(api_key='...'))
The API may change between releases while the capability settles; breaking
changes ship deprecation warnings where practical.
## ExaSearch vs Exa's MCP server
Exa also ships an official hosted MCP server at `https://mcp.exa.ai/mcp`
([exa-labs/exa-mcp-server](https://github.com/exa-labs/exa-mcp-server)) with
their full tool catalog: `web_search_exa`, `get_code_context_exa`,
`crawling_exa`, `company_research_exa`, `linkedin_search_exa`, and a
`deep_researcher_start` / `deep_researcher_check` pair.
`ExaSearch` is the curated, typed path: bounded output, a retry-on-empty
contract, bundled research instructions, and a client seam that makes it
testable offline. The MCP server is how you get Exa's full catalog with zero
wrapper code, via Pydantic AI core's MCP capability. Where the MCP deep
researcher is a start/check tool pair the model must poll across calls,
`deep_search` returns the answer in a single call. The two compose in one
`capabilities` list, and the tool names don't collide (the MCP tools are
`*_exa`-suffixed):
```python
from pydantic_ai import Agent
from pydantic_ai.capabilities import MCP
from pydantic_ai_harness.exa import ExaSearch
agent = Agent('anthropic:claude-sonnet-4-6', capabilities=[ExaSearch(), MCP('https://mcp.exa.ai/mcp')])
```
## Further reading
- [Pydantic AI capabilities](https://ai.pydantic.dev/capabilities/)
pydantic_ai_harness/exa/_capability.py
+25 -3
View File
@@ -16,6 +16,11 @@ _INSTRUCTIONS = (
'the pages you relied on in your answer.'
)
_DEEP_INSTRUCTIONS = _INSTRUCTIONS + (
' For questions that need synthesis across many sources, escalate to `deep_search`: it is '
'slower and costs more per call, but returns a cited answer in one step.'
)
@dataclass
class ExaSearch(AbstractCapability[AgentDepsT]):
@@ -23,6 +28,9 @@ class ExaSearch(AbstractCapability[AgentDepsT]):
Adds two tools: `web_search`, which returns search results together with
page text, and `get_page`, which retrieves the text of a specific URL.
Set `include_deep_search=True` to also expose `deep_search`, which runs
Exa's multi-step deep search and returns a synthesized, cited answer in
one tool call.
```python
from pydantic_ai import Agent
@@ -45,6 +53,15 @@ class ExaSearch(AbstractCapability[AgentDepsT]):
formatted, so page text stays bounded even with a custom `client`.
"""
include_deep_search: bool = False
"""Also expose the `deep_search` tool. Off by default.
Deep search (Exa search `type='deep'`) runs a multi-step agentic search
and synthesizes a cited answer in one call. It is markedly slower and more
expensive per call than `web_search`, and the model decides when to invoke
tools, so the extra spend is opt-in rather than the default.
"""
client: ExaClient | None = None
"""Exa client to use; when `None`, an `exa_py.AsyncExa` is built from `EXA_API_KEY`.
@@ -53,13 +70,18 @@ class ExaSearch(AbstractCapability[AgentDepsT]):
"""
def get_instructions(self) -> str:
"""Static research guidance: search wide, read the promising pages in full, cite URLs."""
return _INSTRUCTIONS
"""Static research guidance: search wide, read the promising pages in full, cite URLs.
When `include_deep_search` is set, the guidance also covers when to
escalate to `deep_search`.
"""
return _DEEP_INSTRUCTIONS if self.include_deep_search else _INSTRUCTIONS
def get_toolset(self) -> ExaSearchToolset[AgentDepsT]:
"""Build the toolset providing the `web_search` and `get_page` tools."""
"""Build the toolset providing the `web_search`, `get_page`, and optional `deep_search` tools."""
return ExaSearchToolset[AgentDepsT](
client=self.client,
num_results=self.num_results,
max_text_chars=self.max_text_chars,
include_deep_search=self.include_deep_search,
)
pydantic_ai_harness/exa/_toolset.py
+77 -12
View File
@@ -1,8 +1,9 @@
"""Exa toolset -- gives agents web search and page retrieval backed by the Exa API."""
"""Exa toolset -- gives agents web search, page retrieval, and deep search backed by the Exa API."""
from __future__ import annotations
from typing import Protocol
import json
from typing import Literal, Protocol
from pydantic_ai.exceptions import ModelRetry, UserError
from pydantic_ai.tools import AgentDepsT
@@ -10,7 +11,14 @@ from pydantic_ai.toolsets import FunctionToolset
try:
from exa_py import AsyncExa
from exa_py.api import ContentsOptions, Result, SearchResponse, TextContentsOptions
from exa_py.api import (
ContentsOptions,
DeepOutputSchema,
Result,
SearchResponse,
SearchType,
TextContentsOptions,
)
except ImportError as _import_error: # pragma: no cover
raise ImportError(
'exa-py is required for ExaSearch. Install it with: pip install "pydantic-ai-harness[exa]"'
@@ -29,8 +37,16 @@ class ExaClient(Protocol):
`AsyncExa` instance satisfies the protocol as-is.
"""
async def search(self, query: str, *, contents: ContentsOptions, num_results: int) -> SearchResponse[Result]:
"""Search the web and return results with page contents."""
async def search(
self,
query: str,
*,
contents: ContentsOptions | Literal[False],
num_results: int | None = None,
type: SearchType | None = None,
output_schema: DeepOutputSchema | None = None,
) -> SearchResponse[Result]:
"""Search the web and return results, optionally with page contents or a synthesized output."""
... # pragma: no cover
async def get_contents(self, urls: str, *, text: TextContentsOptions) -> SearchResponse[Result]:
@@ -53,19 +69,32 @@ class ExaSearchToolset(FunctionToolset[AgentDepsT]):
"""Gives an agent web research tools backed by the Exa search API.
`web_search` surveys the web and returns results together with page text,
and `get_page` retrieves the text of one specific URL. Page text is capped
at `max_text_chars` characters per result: the cap is sent to Exa as the
contents limit and re-enforced when output is formatted, so text stays
bounded even with a custom `client`.
and `get_page` retrieves the text of one specific URL. With
`include_deep_search=True`, `deep_search` runs Exa's multi-step deep search
and returns a synthesized, cited answer in one call.
Page text is capped at `max_text_chars` characters per result: the cap is
sent to Exa as the contents limit and re-enforced when output is formatted,
so text stays bounded even with a custom `client`. The result count is
bounded the same way: `num_results` is requested from Exa and re-applied to
the response.
"""
def __init__(self, *, client: ExaClient | None, num_results: int, max_text_chars: int) -> None:
def __init__(
self, *, client: ExaClient | None, num_results: int, max_text_chars: int, include_deep_search: bool
) -> None:
if num_results < 1:
raise ValueError(f'num_results must be at least 1, got {num_results}')
if max_text_chars < 1:
raise ValueError(f'max_text_chars must be at least 1, got {max_text_chars}')
super().__init__()
self._client = client if client is not None else _default_client()
self._num_results = num_results
self._max_text_chars = max_text_chars
self.add_function(self.web_search, name='web_search')
self.add_function(self.get_page, name='get_page')
if include_deep_search:
self.add_function(self.deep_search, name='deep_search')
async def web_search(self, query: str) -> str:
"""Search the web and return matching pages, each with its text content.
@@ -84,7 +113,8 @@ class ExaSearchToolset(FunctionToolset[AgentDepsT]):
)
if not response.results:
return f'No results found for {query!r}.'
sections = [_format_result(result, self._max_text_chars) for result in response.results]
results = response.results[: self._num_results]
sections = [_format_result(result, self._max_text_chars) for result in results]
plural = 's' if len(sections) != 1 else ''
joined = '\n\n---\n\n'.join(sections)
return f'Found {len(sections)} result{plural} for {query!r}:\n\n{joined}'
@@ -99,10 +129,45 @@ class ExaSearchToolset(FunctionToolset[AgentDepsT]):
The page's title, URL, and text content.
"""
response = await self._client.get_contents(url, text={'max_characters': self._max_text_chars})
if not response.results:
if not response.results or not response.results[0].text:
raise ModelRetry(f'No content could be retrieved for {url!r}. Check the URL or try another page.')
return _format_result(response.results[0], self._max_text_chars)
async def deep_search(self, question: str) -> str:
"""Run Exa's multi-step deep search and return a synthesized answer with its sources.
Slower and more expensive per call than `web_search`; suited to
questions that need synthesis across many sources rather than a quick
survey.
Args:
question: The research question to answer.
Returns:
The synthesized answer, followed by the sources it drew on.
"""
response = await self._client.search(
question,
contents=False,
type='deep',
output_schema={'type': 'text'},
)
output = response.output
if output is None or not output.content:
raise ModelRetry(
f'Deep search returned no answer for {question!r}. Rephrase the question, or use web_search.'
)
content = output.content
answer = content if isinstance(content, str) else json.dumps(content)
sources = {citation.url: citation.title for row in output.grounding for citation in row.citations}
if not sources:
sources = {result.url: result.title or '' for result in response.results}
lines = [_truncate(answer, self._max_text_chars)]
if sources:
lines.extend(['', 'Sources:'])
lines.extend(f'- {title or "(untitled)"}: {url}' for url, title in sources.items())
return '\n'.join(lines)
def _format_result(result: Result, max_text_chars: int) -> str:
"""Render one result as labelled metadata lines followed by the page text."""
tests/exa/test_exa.py
+174 -12
View File
@@ -2,10 +2,22 @@
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import Literal
import pytest
from exa_py.api import ContentsOptions, Result, SearchResponse, TextContentsOptions
from exa_py.api import (
ContentsOptions,
DeepOutputSchema,
DeepSearchOutput,
DeepSearchOutputGrounding,
DeepSearchOutputGroundingCitation,
Result,
SearchResponse,
SearchType,
TextContentsOptions,
)
from pydantic_ai import Agent
from pydantic_ai.exceptions import ModelRetry, UserError
from pydantic_ai.messages import ModelRequest, ModelResponse, ToolCallPart, ToolReturnPart
@@ -37,28 +49,82 @@ def _response(*results: Result) -> SearchResponse[Result]:
return SearchResponse(results=list(results), resolved_search_type=None, auto_date=None)
def _deep_response(
content: str | dict[str, object],
*,
citations: Sequence[tuple[str, str]] = (),
results: Sequence[Result] = (),
) -> SearchResponse[Result]:
"""A real `exa_py` deep-search response with a synthesized output."""
grounding = (
[
DeepSearchOutputGrounding(
field='answer',
citations=[DeepSearchOutputGroundingCitation(url=url, title=title) for url, title in citations],
confidence='high',
)
]
if citations
else []
)
return SearchResponse(
results=list(results),
resolved_search_type=None,
auto_date=None,
output=DeepSearchOutput(content=content, grounding=grounding),
)
@dataclass
class _FakeExaClient:
"""In-memory `ExaClient` double: canned responses, recorded call arguments."""
search_response: SearchResponse[Result] = field(default_factory=_response)
contents_response: SearchResponse[Result] = field(default_factory=_response)
search_calls: list[tuple[str, ContentsOptions, int]] = field(default_factory=list[tuple[str, ContentsOptions, int]])
deep_response: SearchResponse[Result] = field(default_factory=_response)
search_calls: list[dict[str, object]] = field(default_factory=list[dict[str, object]])
get_contents_calls: list[tuple[str, TextContentsOptions]] = field(
default_factory=list[tuple[str, TextContentsOptions]]
)
async def search(self, query: str, *, contents: ContentsOptions, num_results: int) -> SearchResponse[Result]:
self.search_calls.append((query, contents, num_results))
return self.search_response
async def search(
self,
query: str,
*,
contents: ContentsOptions | Literal[False],
num_results: int | None = None,
type: SearchType | None = None,
output_schema: DeepOutputSchema | None = None,
) -> SearchResponse[Result]:
self.search_calls.append(
{
'query': query,
'contents': contents,
'num_results': num_results,
'type': type,
'output_schema': output_schema,
}
)
return self.deep_response if type is not None else self.search_response
async def get_contents(self, urls: str, *, text: TextContentsOptions) -> SearchResponse[Result]:
self.get_contents_calls.append((urls, text))
return self.contents_response
def _toolset(client: _FakeExaClient, *, num_results: int = 5, max_text_chars: int = 10_000) -> ExaSearchToolset[None]:
return ExaSearch[None](num_results=num_results, max_text_chars=max_text_chars, client=client).get_toolset()
def _toolset(
client: _FakeExaClient,
*,
num_results: int = 5,
max_text_chars: int = 10_000,
include_deep_search: bool = False,
) -> ExaSearchToolset[None]:
return ExaSearch[None](
num_results=num_results,
max_text_chars=max_text_chars,
include_deep_search=include_deep_search,
client=client,
).get_toolset()
class TestWebSearch:
@@ -86,13 +152,34 @@ class TestWebSearch:
async def test_passes_num_results_and_contents_cap_to_client(self) -> None:
client = _FakeExaClient(search_response=_response(_result()))
await _toolset(client, num_results=3, max_text_chars=42).web_search('q')
assert client.search_calls == [('q', {'text': {'max_characters': 42}}, 3)]
assert client.search_calls == [
{
'query': 'q',
'contents': {'text': {'max_characters': 42}},
'num_results': 3,
'type': None,
'output_schema': None,
}
]
async def test_no_results(self) -> None:
client = _FakeExaClient()
output = await _toolset(client).web_search('nothing to see')
assert output == "No results found for 'nothing to see'."
async def test_num_results_enforced_on_oversized_response(self) -> None:
client = _FakeExaClient(
search_response=_response(
_result('https://a.dev', title='A', text=None),
_result('https://b.dev', title='B', text=None),
_result('https://c.dev', title='C', text=None),
)
)
output = await _toolset(client, num_results=2).web_search('q')
assert output == (
"Found 2 results for 'q':\n\nTitle: A\nURL: https://a.dev\n\n---\n\nTitle: B\nURL: https://b.dev"
)
async def test_text_at_cap_is_not_truncated(self) -> None:
client = _FakeExaClient(search_response=_response(_result(title='T', text='x' * 10)))
output = await _toolset(client, max_text_chars=10).web_search('q')
@@ -114,11 +201,69 @@ class TestGetPage:
assert client.get_contents_calls == [('https://a.dev', {'max_characters': 10_000})]
assert output == 'Title: A\nURL: https://a.dev\n\nalpha text'
async def test_no_content_raises_model_retry(self) -> None:
async def test_no_results_raises_model_retry(self) -> None:
client = _FakeExaClient()
with pytest.raises(ModelRetry, match='No content could be retrieved'):
await _toolset(client).get_page('https://gone.dev')
async def test_result_without_text_raises_model_retry(self) -> None:
client = _FakeExaClient(contents_response=_response(_result('https://a.dev', title='A', text=None)))
with pytest.raises(ModelRetry, match='No content could be retrieved'):
await _toolset(client).get_page('https://a.dev')
class TestDeepSearch:
async def test_returns_answer_with_deduplicated_sources(self) -> None:
client = _FakeExaClient(
deep_response=_deep_response(
'Deep answer.',
citations=[('https://a.dev', 'A'), ('https://b.dev', ''), ('https://a.dev', 'A')],
)
)
output = await _toolset(client, include_deep_search=True).deep_search('why is the sky blue?')
assert client.search_calls == [
{
'query': 'why is the sky blue?',
'contents': False,
'num_results': None,
'type': 'deep',
'output_schema': {'type': 'text'},
}
]
assert output == 'Deep answer.\n\nSources:\n- A: https://a.dev\n- (untitled): https://b.dev'
async def test_structured_content_is_rendered_as_json(self) -> None:
client = _FakeExaClient(deep_response=_deep_response({'finding': 'x'}, citations=[('https://a.dev', 'A')]))
output = await _toolset(client, include_deep_search=True).deep_search('q')
assert output == '{"finding": "x"}\n\nSources:\n- A: https://a.dev'
async def test_without_grounding_falls_back_to_result_sources(self) -> None:
client = _FakeExaClient(
deep_response=_deep_response('Answer.', results=[_result('https://a.dev', title='A', text=None)])
)
output = await _toolset(client, include_deep_search=True).deep_search('q')
assert output == 'Answer.\n\nSources:\n- A: https://a.dev'
async def test_without_any_sources_returns_answer_only(self) -> None:
client = _FakeExaClient(deep_response=_deep_response('Answer.'))
output = await _toolset(client, include_deep_search=True).deep_search('q')
assert output == 'Answer.'
async def test_missing_output_raises_model_retry(self) -> None:
client = _FakeExaClient(deep_response=_response())
with pytest.raises(ModelRetry, match='Deep search returned no answer'):
await _toolset(client, include_deep_search=True).deep_search('q')
async def test_empty_content_raises_model_retry(self) -> None:
client = _FakeExaClient(deep_response=_deep_response(''))
with pytest.raises(ModelRetry, match='Deep search returned no answer'):
await _toolset(client, include_deep_search=True).deep_search('q')
def test_tool_absent_by_default_and_present_when_enabled(self) -> None:
client = _FakeExaClient()
assert list(_toolset(client).tools) == ['web_search', 'get_page']
assert list(_toolset(client, include_deep_search=True).tools) == ['web_search', 'get_page', 'deep_search']
class TestExaSearch:
def test_default_client_requires_api_key(self, monkeypatch: pytest.MonkeyPatch) -> None:
@@ -131,18 +276,33 @@ class TestExaSearch:
toolset = ExaSearch[None]().get_toolset()
assert isinstance(toolset, ExaSearchToolset)
def test_non_positive_num_results_rejected(self) -> None:
with pytest.raises(ValueError, match='num_results must be at least 1, got 0'):
ExaSearch[None](num_results=0, client=_FakeExaClient()).get_toolset()
def test_non_positive_max_text_chars_rejected(self) -> None:
with pytest.raises(ValueError, match='max_text_chars must be at least 1, got -1'):
ExaSearch[None](max_text_chars=-1, client=_FakeExaClient()).get_toolset()
def test_instructions_reference_the_tools(self) -> None:
instructions = ExaSearch[None]().get_instructions()
assert 'web_search' in instructions
assert 'get_page' in instructions
assert 'cite' in instructions
assert 'deep_search' not in instructions
async def test_agent_run_uses_both_tools_and_instructions(self) -> None:
def test_instructions_cover_deep_search_when_enabled(self) -> None:
instructions = ExaSearch[None](include_deep_search=True).get_instructions()
assert instructions.startswith(ExaSearch[None]().get_instructions())
assert 'escalate to `deep_search`' in instructions
async def test_agent_run_uses_all_tools_and_instructions(self) -> None:
client = _FakeExaClient(
search_response=_response(_result('https://a.dev', title='A', text='alpha')),
contents_response=_response(_result('https://b.dev', title='B', text='beta')),
deep_response=_deep_response('Deep answer.', citations=[('https://c.dev', 'C')]),
)
agent = Agent(TestModel(), capabilities=[ExaSearch(client=client)])
agent = Agent(TestModel(), capabilities=[ExaSearch(include_deep_search=True, client=client)])
result = await agent.run('Research something.')
@@ -151,6 +311,7 @@ class TestExaSearch:
assert isinstance(first, ModelRequest)
assert first.instructions is not None
assert 'web_search' in first.instructions
assert 'deep_search' in first.instructions
calls = {
part.tool_name: part.args_as_dict()
@@ -166,7 +327,8 @@ class TestExaSearch:
for part in message.parts
if isinstance(part, ToolReturnPart)
}
assert set(returns) == {'web_search', 'get_page'}
assert set(returns) == {'web_search', 'get_page', 'deep_search'}
query = calls['web_search']['query']
assert returns['web_search'] == f'Found 1 result for {query!r}:\n\nTitle: A\nURL: https://a.dev\n\nalpha'
assert returns['get_page'] == 'Title: B\nURL: https://b.dev\n\nbeta'
assert returns['deep_search'] == 'Deep answer.\n\nSources:\n- C: https://c.dev'