Skip to content

Commit 1db487d

Browse files
authored
Explicitly set framework_wrappers for LangChain tools (#1820)
* Avoids a warning when using the `code_generation_tool` ## By Submitting this PR I confirm: - I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/NeMo-Agent-Toolkit/blob/develop/docs/source/resources/contributing/index.md). - We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license. - Any contribution which contains commits that are not Signed-Off will not be accepted. - When the PR is ready for review, new or existing tests cover these changes. - When the PR is ready for review, the documentation is up to date with these changes. ## Summary by CodeRabbit ## Release Notes * **Chores** * Updated framework registration for code generation, Tavily internet search, and Wikipedia search tools for improved LangChain compatibility. Authors: - David Gardner (https://github.com/dagardner-nv) Approvers: - https://github.com/mnajafian-nv URL: #1820
1 parent 7629460 commit 1db487d

3 files changed

Lines changed: 5 additions & 3 deletions

File tree

packages/nvidia_nat_langchain/src/nat/plugins/langchain/tools/code_generation_tool.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ class CodeGenerationTool(FunctionBaseConfig, name="code_generation"):
3636
"this tool!")
3737

3838

39-
@register_function(config_type=CodeGenerationTool)
39+
@register_function(config_type=CodeGenerationTool, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
4040
async def code_generation_tool(config: CodeGenerationTool, builder: Builder):
4141
from langchain_core.prompts.chat import ChatPromptTemplate
4242

packages/nvidia_nat_langchain/src/nat/plugins/langchain/tools/tavily_internet_search.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
from pydantic import Field
1818

1919
from nat.builder.builder import Builder
20+
from nat.builder.framework_enum import LLMFrameworkEnum
2021
from nat.builder.function_info import FunctionInfo
2122
from nat.cli.register_workflow import register_function
2223
from nat.data_models.common import SerializableSecretStr
@@ -39,7 +40,7 @@ class TavilyInternetSearchToolConfig(FunctionBaseConfig, name="tavily_internet_s
3940
description="Depth for relevance vs latency tradeoff - 'basic', 'advanced', 'fast', or 'ultra-fast'")
4041

4142

42-
@register_function(config_type=TavilyInternetSearchToolConfig)
43+
@register_function(config_type=TavilyInternetSearchToolConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
4344
async def tavily_internet_search(tool_config: TavilyInternetSearchToolConfig, builder: Builder):
4445
import os
4546

packages/nvidia_nat_langchain/src/nat/plugins/langchain/tools/wikipedia_search.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
# limitations under the License.
1515

1616
from nat.builder.builder import Builder
17+
from nat.builder.framework_enum import LLMFrameworkEnum
1718
from nat.builder.function_info import FunctionInfo
1819
from nat.cli.register_workflow import register_function
1920
from nat.data_models.function import FunctionBaseConfig
@@ -28,7 +29,7 @@ class WikiSearchToolConfig(FunctionBaseConfig, name="wiki_search"):
2829

2930

3031
# Wiki search
31-
@register_function(config_type=WikiSearchToolConfig)
32+
@register_function(config_type=WikiSearchToolConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
3233
async def wiki_search(tool_config: WikiSearchToolConfig, builder: Builder):
3334
from langchain_community.document_loaders import WikipediaLoader
3435

0 commit comments

Comments
 (0)