feat(tavily): Add Tavily tools for framework-neutral usage#1901
feat(tavily): Add Tavily tools for framework-neutral usage#1901lakshyaag-tavily wants to merge 7 commits intoNVIDIA:developfrom
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughAdds a new framework-neutral Tavily subpackage and example: implements AsyncTavilyClient construction, SSE research-stream parsing, a Changes
Sequence Diagram(s)sequenceDiagram
participant Agent as Agent / Workflow
participant FuncGroup as Tavily FunctionGroup
participant Client as AsyncTavilyClient
participant SDK as Tavily SDK (HTTP / SSE)
participant Parser as SSE Aggregator
Agent->>FuncGroup: invoke tool (e.g., tavily__research)
FuncGroup->>Client: build/resolve client, call method (stream=True)
Client->>SDK: open /research SSE stream
SDK-->>Parser: SSE byte chunks
Parser-->>FuncGroup: aggregated result (content, sources, model, trace)
FuncGroup-->>Agent: structured tool result
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
b31a254 to
62485c0
Compare
There was a problem hiding this comment.
Actionable comments posted: 8
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@examples/tavily_weather/README.md`:
- Line 8: Update the README sentence that currently claims the "full Tavily tool
group (`search`, `extract`, `crawl`, `map`)" to match the shipped example config
by removing the excluded tools (`crawl` and `map`) so it lists only the actual
enabled tools (e.g., `search`, `extract`), and replace any occurrence of the
phrase "NAT functions" with "NeMo Agent Toolkit functions"; target the line
containing the phrase "full Tavily tool group (`search`, `extract`, `crawl`,
`map`)" and any instances of "NAT functions" in the README to make these exact
wording changes.
In `@packages/nvidia_nat_tavily/src/nat/meta/pypi.md`:
- Line 21: Replace the phrase "NAT functions" in the pypi.md prose with the
spelled-out form "NeMo Agent Toolkit functions" (or use a backticked package
identifier if you intended to refer to a package), e.g., change the sentence
that currently reads "It provides framework-neutral NAT functions..." to "It
provides framework-neutral NeMo Agent Toolkit functions..." to comply with the
documentation guideline.
In `@packages/nvidia_nat_tavily/src/nat/plugins/tavily/_client.py`:
- Around line 25-28: The code currently treats whitespace-only API keys as
valid; update the client initialization in _client.py to strip and validate both
the config value (api_key via get_secret_value) and the environment fallback
(os.environ.get("TAVILY_API_KEY")) before using them: call .strip() on the
returned strings, treat empty or None after stripping as invalid, and only
assign to resolved if the stripped value is non-empty; ensure references to
api_key, get_secret_value, and the resolved variable are updated accordingly so
a whitespace-only key is rejected.
In `@packages/nvidia_nat_tavily/src/nat/plugins/tavily/parse_streaming.py`:
- Around line 67-107: The loop over chunks can leave a final SSE block in buffer
when the stream ends without a trailing "\n\n", so after the async for finishes,
if buffer is non-empty call _parse_sse_block(buffer) and run the same handling
logic used inside the loop: inspect event_type (handle "done" by returning
_finalize_research), skip if payload is None, raise on payload.object ==
"error", set model if None, extract delta and append content to content_parts or
assign structured_content, update sources when "sources" present and
"tool_calls" not present, and append tool_calls to trace when include_trace is
set; then proceed to the existing final return calling _finalize_research with
content_parts, structured_content, sources, model, trace and include_trace.
- Around line 1-5: This file is missing the required SPDX Apache-2.0 header; add
the repository's standard SPDX license header as the very first lines of
packages/nvidia_nat_tavily/src/nat/plugins/tavily/parse_streaming.py (above all
imports) so the file begins with the Apache-2.0 SPDX block before imports like
json, logging, and symbols such as logger and any functions in this module.
In `@packages/nvidia_nat_tavily/src/nat/plugins/tavily/tools.py`:
- Line 175: The public async generator tavily_tools currently lacks a return
type hint; update its signature to declare the return type as
AsyncIterator[FunctionGroup] (i.e., async def tavily_tools(config:
TavilyToolsGroupConfig, _builder: Builder) -> AsyncIterator[FunctionGroup]:),
and add the necessary import for AsyncIterator (from typing import
AsyncIterator) and ensure FunctionGroup is imported from its defining module so
the annotation resolves; keep the existing parameter types
(TavilyToolsGroupConfig, Builder) unchanged.
- Around line 58-60: The function _annotation_accepts_list currently only treats
typing.Union as a union origin, so PEP 604 unions (X | Y) are not recognized;
update the union check to accept both typing.Union and types.UnionType (or
otherwise detect PEP 604 unions) so the branch `if origin is Union:` also fires
for `types.UnionType`; specifically, import types and change the condition that
inspects `origin` (used alongside `typing.get_args(annotation)`) to accept
either `typing.Union` or `types.UnionType` so annotations like `list[str] | str`
are detected as list-compatible and fallback JSON list coercion runs.
In `@packages/nvidia_nat_tavily/tests/test_tools.py`:
- Line 69: Remove the unnecessary `@pytest.mark.asyncio` decorators from the async
test definitions in this test module: search for and delete every occurrence of
the "@pytest.mark.asyncio" decorator used above async test functions (they are
redundant because the async test runner auto-detects async tests); ensure the
async test function definitions remain unchanged (keep their async def names)
and run the test suite to verify no behavior changed.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 72494b49-1a4a-4452-b1a1-c4bc6a7e21e2
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (13)
examples/tavily_weather/README.mdexamples/tavily_weather/pyproject.tomlexamples/tavily_weather/src/nat_tavily_weather/__init__.pyexamples/tavily_weather/src/nat_tavily_weather/configs/config.ymlpackages/nvidia_nat_tavily/pyproject.tomlpackages/nvidia_nat_tavily/src/nat/meta/pypi.mdpackages/nvidia_nat_tavily/src/nat/plugins/tavily/__init__.pypackages/nvidia_nat_tavily/src/nat/plugins/tavily/_client.pypackages/nvidia_nat_tavily/src/nat/plugins/tavily/parse_streaming.pypackages/nvidia_nat_tavily/src/nat/plugins/tavily/register.pypackages/nvidia_nat_tavily/src/nat/plugins/tavily/tools.pypackages/nvidia_nat_tavily/tests/test_tools.pypyproject.toml
b83b2e9 to
d953f6f
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@packages/nvidia_nat_tavily/src/nat/plugins/tavily/tools.py`:
- Around line 164-167: The research_timeout_seconds Field currently permits
negative numbers; update its Pydantic validation to only allow positive floats
or None by adding a constraint (e.g., use Field(default=900.0, gt=0,
description=...) or replace the annotation with confloat(gt=0) | None) so
research_timeout_seconds rejects non-positive values while still allowing None;
adjust the Field declaration named research_timeout_seconds in the same module
accordingly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9ee035ed-d47d-4b39-bea1-36d7fa7b21cf
📒 Files selected for processing (6)
examples/tavily_weather/README.mdpackages/nvidia_nat_tavily/src/nat/meta/pypi.mdpackages/nvidia_nat_tavily/src/nat/plugins/tavily/_client.pypackages/nvidia_nat_tavily/src/nat/plugins/tavily/parse_streaming.pypackages/nvidia_nat_tavily/src/nat/plugins/tavily/tools.pypackages/nvidia_nat_tavily/tests/test_tools.py
✅ Files skipped from review due to trivial changes (2)
- packages/nvidia_nat_tavily/src/nat/meta/pypi.md
- examples/tavily_weather/README.md
🚧 Files skipped from review as they are similar to previous changes (1)
- packages/nvidia_nat_tavily/tests/test_tools.py
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
packages/nvidia_nat_tavily/src/nat/plugins/tavily/tools.py (1)
164-167:⚠️ Potential issue | 🟡 MinorReject non-positive
research_timeout_secondsin the model.
0and negative values are accepted today and only fail later when the HTTP client consumes them. Add a field constraint so invalid config is rejected at parse time.🛡️ Proposed fix
research_timeout_seconds: float | None = Field( default=900.0, + gt=0, description="HTTP timeout for the research SSE stream. None disables the client-side cap.", )As per coding guidelines: "Validate and sanitise all user input, especially in web or CLI interfaces."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_tavily/src/nat/plugins/tavily/tools.py` around lines 164 - 167, The model currently accepts zero or negative values for the research_timeout_seconds Field; reject non-positive values at parse time by adding a constraint (e.g. set gt=0 on the Field) or a Pydantic validator for the research_timeout_seconds attribute so only None or numbers > 0 are allowed; update the Field declaration for research_timeout_seconds (or add a `@validator` on the containing model) and include a clear error message if validation fails.
🧹 Nitpick comments (1)
packages/nvidia_nat_tavily/src/nat/plugins/tavily/parse_streaming.py (1)
54-55: Align thetraceannotation with the actual payload shape.
tool_callsis appended as one list per SSE event, sotraceis effectivelylist[list[dict]], notlist[dict]. Tightening the annotation here will make the accumulator contract match what the code returns.♻️ Proposed fix
def _finalize_research( *, content_parts: list[str], structured_content: dict | None, sources: list[dict], model: str | None, - trace: list[dict], + trace: list[list[dict]], include_trace: bool, ) -> dict: @@ - trace: list[dict] = [] + trace: list[list[dict]] = []Also applies to: 79-80, 109-110
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_tavily/src/nat/plugins/tavily/parse_streaming.py` around lines 54 - 55, The type annotation for the parameter named "trace" is incorrect: update its type from list[dict] to list[list[dict]] wherever it's declared (the function signatures in this module that accept a trace parameter) so the accumulator contract matches that tool_calls are appended as one list per SSE event; adjust all occurrences (including the signatures where include_trace is present) to use list[list[dict]].
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@packages/nvidia_nat_tavily/tests/test_tools.py`:
- Around line 1-3: Replace the short two-line SPDX pair at the top of the test
module with the repository's full standard Apache-2.0 SPDX header template: add
the complete multi-line header block (including copyright notice line,
SPDX-License-Identifier, and any required boilerplate lines) at the very start
of packages/nvidia_nat_tavily/tests/test_tools.py so the file begins with the
canonical Apache-2.0 header template instead of just the two SPDX lines.
---
Duplicate comments:
In `@packages/nvidia_nat_tavily/src/nat/plugins/tavily/tools.py`:
- Around line 164-167: The model currently accepts zero or negative values for
the research_timeout_seconds Field; reject non-positive values at parse time by
adding a constraint (e.g. set gt=0 on the Field) or a Pydantic validator for the
research_timeout_seconds attribute so only None or numbers > 0 are allowed;
update the Field declaration for research_timeout_seconds (or add a `@validator`
on the containing model) and include a clear error message if validation fails.
---
Nitpick comments:
In `@packages/nvidia_nat_tavily/src/nat/plugins/tavily/parse_streaming.py`:
- Around line 54-55: The type annotation for the parameter named "trace" is
incorrect: update its type from list[dict] to list[list[dict]] wherever it's
declared (the function signatures in this module that accept a trace parameter)
so the accumulator contract matches that tool_calls are appended as one list per
SSE event; adjust all occurrences (including the signatures where include_trace
is present) to use list[list[dict]].
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 38f523cf-e1d0-430d-b7bb-30eb2f3e274e
📒 Files selected for processing (6)
examples/tavily_weather/README.mdpackages/nvidia_nat_tavily/src/nat/meta/pypi.mdpackages/nvidia_nat_tavily/src/nat/plugins/tavily/_client.pypackages/nvidia_nat_tavily/src/nat/plugins/tavily/parse_streaming.pypackages/nvidia_nat_tavily/src/nat/plugins/tavily/tools.pypackages/nvidia_nat_tavily/tests/test_tools.py
✅ Files skipped from review due to trivial changes (2)
- packages/nvidia_nat_tavily/src/nat/meta/pypi.md
- examples/tavily_weather/README.md
🚧 Files skipped from review as they are similar to previous changes (1)
- packages/nvidia_nat_tavily/src/nat/plugins/tavily/_client.py
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
packages/nvidia_nat_tavily/src/nat/plugins/tavily/tools.py (1)
1-24: Add a module docstring for this public plugin module.This file exposes a public integration module, but it currently starts straight with imports after the SPDX header. Please add a short Google-style module docstring describing that it builds the framework-neutral
tavilyfunction group and reflects Tavily SDK signatures into NAT tool schemas.As per coding guidelines, "Provide Google-style docstrings for every public module, class, function and CLI command."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_tavily/src/nat/plugins/tavily/tools.py` around lines 1 - 24, Add a Google-style module docstring at the top of this public plugin module (immediately after the SPDX header and before imports) that briefly describes the module's purpose: it builds the framework-neutral tavily function group and reflects Tavily SDK signatures into NAT tool schemas for use by the NAT integration; mention that this module exposes the public tavily API surface and any important usage notes or exported symbols (e.g., the tavily function group) so documentation tools and linters recognize the module as documented.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@packages/nvidia_nat_tavily/src/nat/plugins/tavily/tools.py`:
- Around line 181-208: The wrapper coroutines _search, _extract, _crawl, _map,
and _research are registered via group.add_function but have empty __doc__
strings, so the LLM-facing tool descriptions are missing; update each wrapper's
docstring (assign or declare a docstring on
_search/_extract/_crawl/_map/_research) to the matching entry from _DESCRIPTIONS
(and ensure the text describes the LLM-facing parameters/result format rather
than internal wrapper args) before calling group.add_function so the registered
tool descriptions are populated.
---
Nitpick comments:
In `@packages/nvidia_nat_tavily/src/nat/plugins/tavily/tools.py`:
- Around line 1-24: Add a Google-style module docstring at the top of this
public plugin module (immediately after the SPDX header and before imports) that
briefly describes the module's purpose: it builds the framework-neutral tavily
function group and reflects Tavily SDK signatures into NAT tool schemas for use
by the NAT integration; mention that this module exposes the public tavily API
surface and any important usage notes or exported symbols (e.g., the tavily
function group) so documentation tools and linters recognize the module as
documented.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2cbbd327-caef-42e4-a2db-f2369fbc2b77
📒 Files selected for processing (1)
packages/nvidia_nat_tavily/src/nat/plugins/tavily/tools.py
4c7a10e to
708e6c2
Compare
This commit introduces the `nvidia-nat-tavily` package, which integrates Tavily with the NeMo Agent Toolkit. It includes a new example workflow for weather queries using Tavily and LangChain, demonstrating the functionality of the new package. Key changes: - Added `nvidia-nat-tavily` package with necessary configurations and dependencies. - Created example workflow `nat_tavily_weather` showcasing Tavily's web search capabilities. - Updated `pyproject.toml` and `uv.lock` to include the new package. - Added tests for the Tavily search functionality to ensure proper integration. Signed-off-by: Lakshya Agarwal <lakshya.agarwal@tavily.com>
…ction group This commit updates the Tavily integration by consolidating the search, extract, crawl, and map functionalities into a single function group. Signed-off-by: Lakshya Agarwal <lakshya.agarwal@tavily.com>
Signed-off-by: Lakshya Agarwal <lakshya.agarwal@tavily.com>
Signed-off-by: Lakshya Agarwal <lakshya.agarwal@tavily.com>
Signed-off-by: Lakshya Agarwal <lakshya.agarwal@tavily.com>
Signed-off-by: Lakshya Agarwal <lakshya.agarwal@tavily.com>
Signed-off-by: Lakshya Agarwal <lakshya.agarwal@tavily.com>
708e6c2 to
e9d6ae1
Compare
Description
Adds a new
nvidia-nat-tavilypackage that exposes the Tavily Python SDK as framework-neutral NAT functions inside a single function group.tavilyfunction group (FunctionGroupBaseConfig, nametavily) — five tools, one sharedAsyncTavilyClient, one API-key resolution path:tavily__searchtavily__extracttavily__crawltavily__maptavily__research(consumes the/researchSSE stream and returns an aggregated{content, sources, model}dict;stream=Trueis locked internally)Example workflow —
examples/tavily_weather/adds a config that consumes the group via a singletool_names: [tavily]referenceWiring —
pyproject.tomladds thetavilyextra, thenvidia-nat-tavilyworkspace source, andnat_tavily_weatherexample.The existing
tavily_internet_searchinnvidia_nat_langchainis left untouched for backwards compatibility.Closes #1900.
By Submitting this PR I confirm:
Summary by CodeRabbit
New Features
Documentation
Tests
Chores