Feature/oracle agent spec integration#1566
Feature/oracle agent spec integration#1566afourniernv wants to merge 5 commits intoNVIDIA:developfrom
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the
✨ Finishing touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Implement nvidia-nat-agent-spec plugin that bridges Oracle Agent-Spec YAML configurations with NeMo Agent Toolkit, enabling Agent-Spec agents to run as NAT Functions with full access to NAT's evaluation, profiling, and observability. Implementation: - AgentSpecWrapperFunction: Wraps LangGraph CompiledStateGraph as NAT Function - AgentSpecWrapperConfig: Configuration model supporting tool_registry, components_registry, and auto checkpointer detection - Two-layer adapter pattern: Oracle's AgentSpecLoader (YAML→LangGraph) + NAT wrapper (LangGraph→NAT Function) - Supports multiple input formats (strings, lists, message objects) - Auto-detects ClientTool and creates MemorySaver checkpointer when needed Features: - Tool registry: Wire up custom functions/tools via config - Components registry: Override LLMs/components (e.g., inject NIM LLMs) - Multiple input formats: Flexible input handling via convert_to_messages - Full NAT integration: Works with eval, profiler, observability, middleware Test Coverage: - Unit tests for AgentSpecWrapperFunction methods (_convert_input, _ainvoke, _astream, convert_to_str) - Registration function tests (tool registry, checkpointer, components registry) - Integration tests (tool usage, NIM LLM integration) - Test fixtures: Minimal agent, NIM agent, weather agent with tools Dependencies: - Resolves langchain-core/langgraph version conflicts via override-dependencies - Adds agent-spec extra to pyproject.toml - Updates uv.lock with new dependencies Signed-off-by: afourniernv <afournier@nvidia.com>
- Add nvidia-nat-agent-spec package to workspace with proper dependency management - Configure plugin to use setuptools_dynamic_dependencies pattern (matching langchain plugin) - Update plugin to depend on nvidia-nat-core instead of root nvidia-nat package - Add conflict declarations in root pyproject.toml for agent-spec vs langchain/vanna/most extras (pyagentspec requires langchain-core<1.0.0 while nvidia-nat-langchain requires >=1.2.6) - Add comprehensive README documenting installation, usage, known limitations, and testing Signed-off-by: afourniernv <afournier@nvidia.com>
…rations Add support for providing Agent-Spec configurations inline as YAML or JSON strings in addition to file paths. This enables more flexible configuration options for users who want to embed Agent-Spec configs directly in workflow definitions. Changes: - Add spec_yaml and spec_json fields to AgentSpecWrapperConfig - Make spec_file optional (FilePath | None) - Add validator to ensure exactly one source (file/yaml/json) is provided - Update register() to handle all three formats with proper format detection - Support both .yaml and .json file extensions - Maintain backward compatibility with existing spec_file usage Tests: - Add 8 new test cases covering inline YAML/JSON scenarios - Test validator enforcement (exactly one source required) - Test format detection for JSON files - Test components_registry integration with inline formats - Test ClientTool auto-detection in inline YAML/JSON All 21 tests passing (13 existing + 8 new) Signed-off-by: afourniernv <afournier@nvidia.com>
…istory trimming - Add inline YAML/JSON support: Allow Agent-Spec configs via spec_yaml and spec_json fields - Make spec_file optional with validator ensuring exactly one source is provided - Support both YAML and JSON formats for inline content and file-based configs - Add tool_names integration: Enable NAT tool discovery via FunctionRef/FunctionGroupRef - Integrate with builder.get_tools() to automatically resolve NAT-registered tools - Merge discovered tools with manual tool_registry (tool_registry takes precedence) - Support both individual function references and function group references - Add max_history message trimming: Implement conversation history truncation - Add max_history field (default: 15) to AgentSpecWrapperConfig - Use langchain_core.messages.trim_messages to limit message count before execution - Preserve full input state while only trimming message history - Consolidate tests: Use pytest.mark.parametrize for YAML/JSON test pairs - Reduce code duplication by parameterizing format-based tests - Maintain same test coverage with cleaner, more maintainable code - Add end-to-end smoke tests: Verify full integration with real LLM calls - Test Agent-Spec YAML file loading and execution - Test inline YAML support end-to-end - Marked as integration tests that skip gracefully without NVIDIA_API_KEY - Update minimal_agent.yaml fixture: Use OpenAiCompatibleConfig for NIM compatibility Tests: All 29 unit tests + 2 end-to-end tests passing Signed-off-by: afourniernv <afournier@nvidia.com>
- Enhanced E2E test validation (graph type, message types, content checks) - Added explicit return type annotation (-> Self) to _validate_sources - Documented component reuse limitation in class docstring and field description Signed-off-by: afourniernv <afournier@nvidia.com>
afourniernv
force-pushed
the
feature/oracle-agent-spec-integration
branch
from
10672ac to
c4b8b71
Compare
willkill07
added
feature request
| description="Path to the Agent-Spec YAML/JSON configuration file. " | ||
| "Recommended for NAT YAML workflow configurations." | ||
| ) | ||
| spec_yaml: str | None = Field( |
There was a problem hiding this comment.
here, the workflow config object contains only a stringified version of an agentspec config. This means that when using the optimizer, the optimizer cannot navigate in this config to discover parameters that can be optimized nor the values/ranges they should be optimized on.
I see two ways to make this work.
Using a dynamic basemodel
This is what I'm currently experimenting with. I basically added 2 attributes to this workflow config, one "dynamic_optimizeable_parameters" that's empty initially, and a "optimizeable_parameters_dict" dict.
It looks like this in the config:
workflow:
_type: agent_spec_wrapper
spec_file: stage_2.yaml
# this is required for adding dynamically the optimizable parameters in the config object
optimizable_params:
- dynamic_optimizable_parameters
# dict for optimize-able parameters with values
optimizable_parameters_dict:
temperature_ref:
range:
low: 0.0
high: 1.0
step: 0.1
max_tokens_ref:
range:
low: 2048
high: 8096
step: 2048
Then, as a post-validation of the AgentSpecWrapperConfig object, I create a dynamic base-model:
@model_validator(mode="after")
def _build_prompts_opt(self) -> Self:
fields: dict[str, tuple[type, Any]] = {}
for var_name, var in self.optimizable_parameters_dict.items():
if var.values is not None:
fields[var_name] = (
dict,
OptimizableField(
default=var.values[0],
space=SearchSpace(values=var.values),
),
)
elif var.range is not None:
fields[var_name] = (
dict,
OptimizableField(
default=var.range.low,
space=SearchSpace(low=var.range.low, high=var.range.high, step=var.range.step),
),
)
else:
raise ValueError('Not supported')
DynamicPrompts = create_model("DynamicPrompts", **fields)
self.dynamic_optimizable = DynamicPrompts()
return self
This allows the optimizer to still see all the parameters inside the config, even if these attributes are not decorated with OptimizableField.
Cleaner approach: patch our config objects with the "optimizable" annotation
We also represent config objects as base models, so one way would be to extend the loading of the workflow config with our own loading logic, and patch our base models with the optimizable field annotator, for example on temperature like in https://github.com/oracle/agent-spec/blob/main/pyagentspec/src/pyagentspec/llms/llmgenerationconfig.py#L23.
This would allow a more seemless and native use of the optimizer with agentspec configs.
Do you see another potential solution, and what do you think about these?
| "ClientTool in Agent-Spec YAML. If not provided and ClientTool is detected, " | ||
| "MemorySaver will be used automatically.", | ||
| ) | ||
| max_history: int = Field( |
There was a problem hiding this comment.
we support this concept natively in agentspec, I don't think we need this here (we can remove the part about pruning the history before the invoke/stream calls too)
| logger = logging.getLogger(__name__) | ||
|
|
||
|
|
||
| class AgentSpecWrapperInput(BaseModel): |
There was a problem hiding this comment.
we would need to extend to be also a dict of inputs, since soem agentspec config takes inputs and not a list of messages.
| messages: list[BaseMessage] | ||
|
|
||
|
|
||
| class AgentSpecWrapperOutput(BaseModel): |
There was a problem hiding this comment.
similarly, we would need to support outputting a dict of outputs. Do you already support handling such input/output dictionnaries, and how does it plain with the evaluator (the mappings question -> WrapperInput and WrapperOutput -> output are defined below in AgentSpecWrapperFunction, is that what we always want?
| if config.tool_registry: | ||
| tool_registry.update(config.tool_registry) |
There was a problem hiding this comment.
We should raise an error on any intersection of tool names rather than blindly take precedence
| # Determine the format and read the Agent-Spec content | ||
| if config.spec_file: | ||
| # Read from file | ||
| spec_path = Path(config.spec_file) | ||
| if not spec_path.exists(): | ||
| raise ValueError(f"Agent-Spec file '{spec_path}' does not exist.") | ||
|
|
||
| with open(spec_path, "r", encoding="utf-8") as f: | ||
| spec_content = f.read() | ||
|
|
||
| # Determine format from file extension | ||
| ext = spec_path.suffix.lower() | ||
| spec_format = "json" if ext == ".json" else "yaml" | ||
| source_description = f"file '{spec_path}'" | ||
| elif config.spec_yaml: | ||
| # Use inline YAML | ||
| assert config.spec_yaml is not None # Type narrowing: validator ensures this is set | ||
| spec_content = config.spec_yaml | ||
| spec_format = "yaml" | ||
| source_description = "inline YAML" | ||
| else: | ||
| # Use inline JSON (config.spec_json) | ||
| assert config.spec_json is not None # Type narrowing: validator ensures this is set | ||
| spec_content = config.spec_json | ||
| spec_format = "json" | ||
| source_description = "inline JSON" |
There was a problem hiding this comment.
Hoist as a _load_spec function?
| "Failed to import pyagentspec.adapters.langgraph. " | ||
| "Install pyagentspec with langgraph extras:\n" | ||
| " uv pip install \"pyagentspec[langgraph,langgraph_mcp]>=26.1.0\"\n\n" | ||
| "Note: The root pyproject.toml includes override-dependencies to resolve " | ||
| "langchain-core version conflicts between pyagentspec and nvidia-nat-langchain." |
There was a problem hiding this comment.
This should be a message to install nvidia-nat[agent_spec]
| # version when adding a new package. If unsure, default to using `~=` instead of `==`. Does not apply to nvidia-nat packages. | ||
| # Keep sorted!!! | ||
| "nvidia-nat-core == {version}", | ||
| "pyagentspec[langgraph,langgraph_mcp]>=26.1.0", |
There was a problem hiding this comment.
Is unbounded upper-bound really safe? I'd rather see: ~=26.1 (>=26.1,<27.0)
| # These version ranges are incompatible and cannot coexist in the same workspace | ||
| # Users must choose one or the other, or install agent-spec in a separate environment |
There was a problem hiding this comment.
remove unnecessary wording
| # These version ranges are incompatible and cannot coexist in the same workspace | |
| # Users must choose one or the other, or install agent-spec in a separate environment |
| @pytest.mark.skipif( | ||
| not os.getenv("NVIDIA_API_KEY"), | ||
| reason="Requires NVIDIA_API_KEY environment variable" | ||
| ) |
There was a problem hiding this comment.
Prefer nvidia_api_key fixture as parameter
| @pytest.mark.skipif( | ||
| not os.getenv("NVIDIA_API_KEY"), | ||
| reason="Requires NVIDIA_API_KEY environment variable" | ||
| ) |
There was a problem hiding this comment.
Prefer nvidia_api_key fixture as parameter
| @pytest.mark.skipif( | ||
| not os.getenv("NVIDIA_API_KEY"), | ||
| reason="Requires NVIDIA_API_KEY environment variable" | ||
| ) |
There was a problem hiding this comment.
Prefer nvidia_api_key fixture as parameter
| component_type: OpenAiCompatibleConfig | ||
| name: "NIM LLM" | ||
| model_id: "meta/llama-3.1-8b-instruct" | ||
| # NIM endpoint - will use NVIDIA_API_KEY from environment |
There was a problem hiding this comment.
Is this true? because it seems like you have to set OPENAI_API_KEY
| name: "NIM LLM" | ||
| model_id: "meta/llama-3.1-8b-instruct" | ||
| # NIM endpoint - will use NVIDIA_API_KEY from environment | ||
| url: "https://integrate.api.nvidia.com/v1" |
There was a problem hiding this comment.
Just a note on how messed up their implementation is...
OpenAiCompatibleConfig has an optional api_key.
But ironically, they don't even map it!
3 participants
Overview
This PR represents a major refactoring and enhancement of the Agent-Spec integration work originally proposed in PR #1432 (YashaPushak:feat/agentspec). This branch addresses review feedback from @willkill07 and implements a more robust, production-ready Agent-Spec plugin that follows NeMo Agent Toolkit's standard plugin architecture patterns.
What This Branch Does
This branch implements a complete Oracle Agent-Spec integration plugin for NeMo Agent Toolkit, allowing users to run Agent-Spec YAML/JSON configurations as native NAT workflows with full observability, profiling, and evaluation capabilities.
Key Features
Agent-Spec Workflow Wrapper: Converts Agent-Spec YAML/JSON configurations to LangGraph
CompiledStateGraphcomponents usingpyagentspec's LangGraph adapter, then wraps them as NAT Functions following the same pattern asLanggraphWrapperFunction.Flexible Configuration Support:
spec_filefor referencing Agent-Spec YAML/JSON filesspec_yamlandspec_jsonfor programmatic/dynamic configurationsNAT Tool Integration:
tool_namesfield supports referencing NAT tools viaFunctionReforFunctionGroupRefbuilder.get_tools()Component Reuse Support:
components_registryfor manually mapping component IDs to NAT-managed components (e.g., LLMs)Conversation History Management:
max_historyfield for trimming message history before executionlangchain_core.messages.trim_messagesto limit context sizeAuto-detection Features:
ClientToolusage and configuresMemorySavercheckpointer when neededArchitecture
This implementation follows NeMo Agent Toolkit's standard plugin architecture:
@register_functiondecorator with proper framework wrappers (LLMFrameworkEnum.LANGCHAIN)AgentSpecWrapperFunctionfollows the same pattern asLanggraphWrapperFunction, handling input/output conversion and delegating execution to the underlying LangGraphnvidia-nat-agent-specpackage with proper dependencies (nvidia-nat-core,pyagentspec[langgraph,langgraph_mcp])AgentSpecWrapperConfig) and input/output (AgentSpecWrapperInput,AgentSpecWrapperOutput)Changes from Original PR #1432
This branch addresses several architectural concerns raised in the review:
-> Self) to_validate_sourcesvalidatorTesting
55 unit tests covering:
tool_names,tool_registry)max_history)2 end-to-end integration tests (smoke tests):
NVIDIA_API_KEYUsage Example
functions:
agent_spec_workflow:
_type: agent_spec_wrapper
spec_file: path/to/agent_spec.yaml
tool_names:
- FunctionRef("wiki_search")
- FunctionGroupRef("search_tools")
max_history: 20
description: "Agent-Spec workflow with NAT tool integration"
The
agent-specextra cannot be installed alongsidelangchain,most, orvannaextras due to incompatiblelangchain-core/langgraphversion requirements:pyagentspecrequires:langchain-core<1.0.0, langgraph<1.0.0nvidia-nat-langchainrequires:langchain-core>=1.2.6, langgraph>=1.0.5Users must choose one or the other, or use separate environments.
Lock File Updates
The
uv.lockfile has been updated to include dependencies for the newnvidia-nat-agent-specpackage. This is expected and necessary when adding a new package.