Skip to content

feat(knowledge): Tavily connector + web-research wrapper with caching#1234

Open
TravisHaa wants to merge 4 commits into
amd:mainfrom
TravisHaa:feature/tavily-connector-1141
Open

feat(knowledge): Tavily connector + web-research wrapper with caching#1234
TravisHaa wants to merge 4 commits into
amd:mainfrom
TravisHaa:feature/tavily-connector-1141

Conversation

@TravisHaa
Copy link
Copy Markdown

Why this matters

GAIA agents could only reach the web through keyless DuckDuckGo HTML scraping — no
Tavily, no result caching, and no way to track or cap API spend. This PR is the Phase 1
foundation for the Knowledge Agent: a _TAVILY connector that exposes Tavily
search/extract as MCP tools to all agents from one keyring-stored API key, plus a
gaia.web.tavily wrapper giving cached, credit-budgeted search/extract/crawl with an
automatic DuckDuckGo fallback when Tavily isn't configured. Agents get higher-quality web
research without re-paying for repeat queries or silently blowing past a credit budget.

Test plan

  • python -m pytest tests/unit/test_tavily_wrapper.py — 16 tests (mocked SDK): cache hit/TTL, credit ledger, budget warn/block, DuckDuckGo fallback
  • python -m pytest tests/unit/connectors/test_catalog_docs_url.py — the new connector's docs_url resolves
  • python util/lint.py --black --isort — clean
  • With a key: gaia connectors configure mcp-tavily --set TAVILY_API_KEY=tvly-… then gaia knowledge search "…" returns Tavily results; without it, the same command falls back to DuckDuckGo
  • gaia knowledge usage prints the credit ledger

Open questions for reviewers

  • Catalog scope: _TAVILY has no built-in agent consumer yet (the Knowledge Agent lands in a later phase), so there's no REQUIRED_CONNECTORS wiring. Acceptable for a foundation PR, or should the entry wait? (@kovtcharov-amd)
  • Async crawl: sync TavilyClient has crawl; AsyncTavilyClient doesn't yet — deferred unless you want parity now.

Phase 1 of #1141.

@TravisHaa TravisHaa requested a review from kovtcharov-amd as a code owner May 29, 2026 01:20
@github-actions github-actions Bot added documentation Documentation changes dependencies Dependency updates cli CLI changes tests Test changes labels May 29, 2026
@TravisHaa
Copy link
Copy Markdown
Author

@claude review this PR

@github-actions
Copy link
Copy Markdown
Contributor

The foundation is solid: the layered design (_TavilyBase → sync/async subclasses → CLI → connector catalog) is clean and the 16-test suite covers the critical paths well. One CLAUDE.md-required artifact is missing before this can merge.


Issues

🟡 Important — gaia knowledge not documented in docs/reference/cli.mdx (src/gaia/cli.py)

CLAUDE.md requires every new CLI subcommand to have an entry in both src/gaia/cli.py and docs/reference/cli.mdx. The knowledge subcommand and its three sub-actions (search, extract, usage) are wired in the parser but absent from the reference doc. The existing gaia cache block (around line 1645 in cli.mdx) is the closest precedent — follow that shape:

### `gaia knowledge`

Web research via Tavily, with SQLite caching, a credit budget, and a keyless DuckDuckGo fallback when Tavily isn't configured.

```bash
gaia knowledge search "AMD ROCm latest release" [--max-results N] [--depth basic|advanced] [--budget N] [--no-block]
gaia knowledge extract <url> [<url>…] [--depth basic|advanced] [--budget N] [--no-block]
gaia knowledge usage

search degrades to DuckDuckGo automatically when the mcp-tavily connector isn't configured. extract requires Tavily — see Tavily connector.


---

### 🟢 Minor — `AsyncTavilyClient` has no `aclose()` (`src/gaia/web/tavily.py:514`)

`close()` is synchronous and calls `close_db()` / `self._web_client.close()`, but the real `AsyncTavilyClient` from `tavily-python` wraps an `httpx.AsyncClient` that requires `await client.aclose()`. Callers holding an `AsyncTavilyClient` in a long-lived async context may leak connections.

The simplest fix is an `aclose()` coroutine that the caller can await, plus an `__aenter__`/`__aexit__` pair:

```suggestion
    async def aclose(self) -> None:
        """Close the async SDK client and the cache DB."""
        if self._sdk is not None and hasattr(self._sdk, "aclose"):
            await self._sdk.aclose()
        self.close()

    async def __aenter__(self) -> "AsyncTavilyClient":
        return self

    async def __aexit__(self, *_: object) -> None:
        await self.aclose()

The sync TavilyClient could get __enter__/__exit__ parity for consistency.


🟢 Minor — Sync SQLite I/O inside async methods (src/gaia/web/tavily.py:820–832)

AsyncTavilyClient.search() and extract() call _cache_get(), _check_budget(), and _record() directly — all of which do synchronous SQLite I/O. For a local .gaia/tavily_cache.db the latency is negligible in practice, but it blocks the event loop during those calls. Wrapping in asyncio.to_thread (same pattern already used for DuckDuckGo at line 812) would keep the async path fully non-blocking. Low urgency, but worth tracking before this is used at scale.


Strengths

  • Cache + ledger design is correct. The SHA-256 key over (op, normalized-payload, sorted-params) is collision-safe and deterministic. Stale rows are lazily replaced by _record() — no orphan accumulation.
  • Budget gate is fail-loud by default. TavilyBudgetExceeded propagates to the CLI which exits 1; block=False is an opt-in downgrade, not the default. Matches CLAUDE.md's no-silent-fallback rule.
  • Test doubles are injected, not patched. sdk_client= and web_client= constructor injection keeps tests deterministic without touching the module graph. The Clock + monkeypatch.setattr(tav, "time", clock) TTL test is a good pattern.
  • _ensure_schema uses IF NOT EXISTS. Safe to call on every init; no migration footgun.
  • Connector catalog entry follows the _GITHUB / _GIT pattern exactly. ConfigField(secret=True) ensures the key is keyring-backed, not written to disk.

Verdict

Request changes — the docs/reference/cli.mdx entry is a one-paragraph addition explicitly required by CLAUDE.md for every new CLI subcommand. The async lifecycle issue is worth fixing in the same pass given the async client is a public API surface. Everything else is ship-worthy.

@github-actions
Copy link
Copy Markdown
Contributor

Nice foundation — the _TavilyBase → sync/async split, the SHA-256 cache key over normalized payload + params, and the fail-loud budget gate are all clean, and the injected test doubles keep the suite deterministic. One blocker I want to flag that the earlier pass didn't catch, plus a confirm on the docs gap.

🟡 AsyncTavilyClient() crashes at construction when the connector is configured

_TavilyBase.__init__ resolves the key synchronously via _load_api_key()get_credential_sync() (src/gaia/web/tavily.py:495). But get_credential_sync raises RuntimeError when called inside a running event loop (src/gaia/connectors/handler.py:191-195). Since the whole point of AsyncTavilyClient is to be constructed inside an async context, the normal real-world path —

async def main():
    client = AsyncTavilyClient()   # mcp-tavily configured
    await client.search("…")

— hits is_mcp_server_configured() == Trueget_credential_sync()RuntimeError: ... called from inside a running event loop.

The tests miss it because every async test either injects sdk_client= (test_async_search_caches) or monkeypatches _load_api_key to return None (test_async_unconfigured_falls_back_to_ddg) — so no test exercises a keyring load inside a running loop. Worth a regression test that constructs the async client with a configured-connector stub.

Options: resolve the key lazily on first await via asyncio.to_thread(_load_api_key), expose an async def create(...) factory, or document that the async client must be constructed off-loop / with an explicit api_key=. The lazy-load path also pairs naturally with the aclose() fix from the prior review.

🟡 gaia knowledge still undocumented in docs/reference/cli.mdx

Confirming the earlier finding stands — CLAUDE.md requires every new subcommand to land in both src/gaia/cli.py and docs/reference/cli.mdx, and knowledge (search/extract/usage) is parser-only right now. One paragraph following the gaia cache block closes it.

Confirmations (not issues)

  • _ensure_schema's multi-statement execute() is fineDatabaseMixin.execute uses executescript() (src/gaia/database/mixin.py:235-269).
  • cred["env"]["TAVILY_API_KEY"] matches the mcp_server handler's return shape (src/gaia/connectors/mcp_server.py:182). ✅
  • docs_url="…/docs/connectors/tavily" matches the _GITHUB precedent exactly — consistent.
  • async_mode = "auto" is set in pyproject.toml, so the unmarked async def test_… cases will run.

Everything else looks ship-worthy once the async constructor path is sorted. Nice work, @TravisHaa.

@github-actions
Copy link
Copy Markdown
Contributor

Solid foundation for the Knowledge Agent phase — the four-layer design (cache → budget gate → Tavily SDK → DDG fallback) is clean and the test suite is thorough. One blocking usability bug in the async client that needs addressing before merge; everything else is minor polish.


Issues

🟡 Important — AsyncTavilyClient() raises RuntimeError when constructed from async code with a configured connector (src/gaia/web/tavily.py:555)

_TavilyBase.__init__ calls _load_api_key(), which calls get_credential_sync(). That function explicitly raises when called from a running event loop:

raise RuntimeError(
    "get_credential_sync() called from inside a running event loop. "
    "Use 'await get_credential(...)' instead."
)

So this pattern, which is exactly how library users would be expected to use the async client, breaks when the connector is configured:

async def my_research():
    async with AsyncTavilyClient() as client:   # RuntimeError if mcp-tavily is configured
        await client.search("query")

The class docstring says it's "for concurrent multi-query research" — which is inherently async usage. The failure only occurs when the connector IS configured (unconfigured → is_mcp_server_configured returns False → early return before get_credential_sync), so tests always pass because they inject sdk_client directly.

The minimal fix is a docstring note and a dedicated async factory. Alternatively, move key resolution into __aenter__:

    async def __aenter__(self) -> "AsyncTavilyClient":
        if not self._configured and self._sdk is None:
            # Defer key resolution to aenter so async callers don't hit
            # get_credential_sync() from inside a running event loop.
            from gaia.web.tavily import _load_api_key  # noqa: PLC0415
            key = _load_api_key()
            if key is not None and self._SDK_CLASS is not None:
                self._sdk = self._SDK_CLASS(api_key=key)
                self._configured = True
        return self

Or, as the lowest-cost fix: add a note to the class docstring that api_key must be supplied when constructing from async code:

class AsyncTavilyClient(_TavilyBase):
    """Asynchronous Tavily wrapper for concurrent multi-query research.

    Mirrors :class:`TavilyClient` but awaits the SDK; the synchronous
    DuckDuckGo fallback is offloaded with ``asyncio.to_thread`` so it doesn't
    block the event loop.

    When constructing from inside a running event loop with the mcp-tavily
    connector configured, supply ``api_key`` explicitly — the default key
    lookup uses :func:`get_credential_sync` which raises inside an event loop::

        async with AsyncTavilyClient(api_key=my_key) as client:
            results = await client.search("query")
    """

🟢 Minor — "optional dependency" comment contradicts install_requires (src/gaia/web/tavily.py:434-444)

The module comment says tavily-python is "an optional dependency", but it's in install_requires in setup.py:194. web/client.py uses the same guard pattern for beautifulsoup4 (also in install_requires) — so the pattern is fine, but the word "optional" is misleading. "Conditionally guarded" is more accurate.

# ``tavily-python`` is bundled with the standard install but the import is
# guarded so that a version-mismatched or stripped install still degrades
# to the DuckDuckGo fallback rather than crashing at import time.
# Mirrors how web/client.py handles beautifulsoup4.

🟢 Minor — Stale entries accumulate silently until the same query is re-fetched (src/gaia/web/tavily.py:615-625)

_cache_get returns None on stale entries without deleting them — they're only removed in _record when the same key is written again. For long-lived deployments with many unique queries, the cache table will grow unboundedly. The DB file on disk stays at ~/.gaia/tavily_cache.db and has no purge path.

Not a correctness bug, but worth a TODO comment or a periodic vacuum call. A gaia knowledge cache clear subcommand or a PRAGMA auto_vacuum setting would address this for the GA feature.


🟢 Nit — _check_budget warn-only log message (src/gaia/web/tavily.py:695)

log.warning("%s Proceeding (budget is in warn-only mode).", msg)

msg already ends with a period, producing "...cap is 5. Proceeding (budget is in warn-only mode)." — the double-sentence join reads awkwardly. The test catches "warn-only" so it passes, but clarity improves with:

            log.warning(
                "Tavily budget exceeded (%d/%d + ~%d = %d > cap %d), "
                "but proceeding because budget is in warn-only mode.",
                used, cap, est, projected, cap,
            )

Strengths

  • Clean abstraction hierarchy. _TavilyBase → TavilyClient / AsyncTavilyClient keeps the cache/ledger/budget logic in one place with no drift risk between the two clients.
  • Test depth. All meaningful paths are exercised with injected fakes — cache TTL, budget warn/block, DDG fallback, both context managers, credit from SDK response vs. estimate. test_budget_warns_near_threshold directly seeding the ledger via client.insert() is a nice use of the mixin API.
  • Atomic cache + ledger. Wrapping _record in self.transaction() is exactly right — a crash between the INSERT steps would otherwise leave the ledger inconsistent with the cache.
  • Documentation is complete. The Tavily connector page, CLI reference, and docs.json entry are all updated and consistent with the implementation.

Verdict

Request changes — the 🟡 async-init issue is a real usability bug that will surface as soon as the Knowledge Agent (the stated next phase) tries to use AsyncTavilyClient() from async code. A docstring-only fix is acceptable if a proper async factory is out of scope for this PR, but the constraint must be documented before merge.

@kovtcharov-amd
Copy link
Copy Markdown
Collaborator

@itomek-amd, please review this PR when you get a chance.

itomek
itomek previously approved these changes May 29, 2026
Copy link
Copy Markdown
Collaborator

@itomek itomek left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @TravisHaa — this is a genuinely strong first contribution. The four-layer design (cache -> budget gate -> Tavily SDK -> DuckDuckGo fallback) is clean, the SHA-256 cache key over normalized-query + sorted-params is correct, the budget gate is fail-loud by default (matching the no-silent-fallbacks rule), and injecting sdk_client=/web_client= in the tests instead of monkeypatching the module graph is exactly the pattern we like. Docs and the dep declaration are in place, and the description follows the house style well (why-first + test plan + "Phase 1 of #1141" + explicit open questions).

One real bug to land or document, noted inline: AsyncTavilyClient() raises RuntimeError inside a running event loop when the connector is configured, because key resolution happens synchronously in init via get_credential_sync(). The tests don't catch it because they all inject a client or hit the unconfigured early-return. Fix is to defer key resolution into aenter (or document that async callers pass api_key=), plus a regression test. There's also a low-urgency unbounded-stale-cache note inline.

On your catalog-scope open question: a foundation connector entry with no agent consumer yet is fine for a Phase 1 PR — just keep it called out as you did. Approving; happy to re-review once the async-init path is sorted.


Generated by Claude Code

Comment thread src/gaia/web/tavily.py
Comment thread src/gaia/web/tavily.py
kovtcharov-amd
kovtcharov-amd previously approved these changes May 29, 2026
@kovtcharov-amd
Copy link
Copy Markdown
Collaborator

@TravisHaa, thanks for the contribution. Approved on condition the following failing test is fixed:

FAILED tests/unit/connectors/test_catalog_ledger.py::test_no_unexpected_mcp_ids - AssertionError: Unexpected MCP catalog ids beyond the kept set: ['mcp-tavily']. Add the id to KEPT_IDS (with rationale) or remove it from the catalog.

@TravisHaa TravisHaa dismissed stale reviews from kovtcharov-amd and itomek via b860d5c May 29, 2026 19:02
@TravisHaa
Copy link
Copy Markdown
Author

@itomek can I get a approval for this so we can run CI/CD again? Ready for another pass when you have a moment :)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cli CLI changes dependencies Dependency updates documentation Documentation changes tests Test changes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants