Skip to content

Commit d9eb559

Browse files
dagardner-nvdagardner-nv
authored
Include plugins in the staged API tree for documentation builds (#909)
* Work around readthedocs/sphinx-autoapi#298 by copying the contents of plugin packages into the staged API tree. Also refactor this logic into a method. * Rename the `span_to_dfw_record` module to `span_to_dfw` this fixes an issue where the `span_to_dfw_record` contained a function of the same name, which was also imported into the package, making the rest of the module difficult to import. * Rename dataflyweel's `Request` class to `ESRequest` as this was conflicting with the `nat.data_models.api_server.Request` class * Fix Sphinx docstring parsing errors such as `**kwargs` in a parameter list is interpreted as a bold open, without a closing `**`. * Add missing `__init__.py` files * Since dataflywheel was not included in v1.2, I'm claiming this to be a non-breaking change. ## 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.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 - Refactor - Streamlined docs build with dynamic API tree generation, improved path handling, and path logging. - Renamed Elasticsearch request model to ESRequest; updated related components accordingly. - OTLP span redaction exporter now accepts an otlp_kwargs dict instead of variadic kwargs. - Documentation - Standardized docstrings across plugins: kwargs naming, “Example::” formatting, and spacing; no behavioral changes. - Tests - Updated tests to new trace conversion module paths; coverage preserved. - Chores - Added licensing boilerplate to a utilities package initializer. Authors: - David Gardner (https://github.com/dagardner-nv) Approvers: - Will Killian (https://github.com/willkill07) URL: #909
1 parent 4ee0dfd commit d9eb559

File tree

20 files changed

+124
-73
lines changed

20 files changed

+124
-73
lines changed

‎docs/source/conf.py‎

Lines changed: 44 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -27,32 +27,60 @@
2727
# add these directories to sys.path here. If the directory is relative to the
2828
# documentation root, use os.path.abspath to make it absolute, like shown here.
2929

30+
import glob
3031
import os
3132
import shutil
3233
import subprocess
3334
import typing
35+
from pathlib import Path
3436

3537
if typing.TYPE_CHECKING:
3638
from autoapi._objects import PythonObject
3739

38-
CUR_DIR = os.path.dirname(os.path.abspath(__file__))
39-
DOC_DIR = os.path.dirname(CUR_DIR)
40-
ROOT_DIR = os.path.dirname(os.path.dirname(CUR_DIR))
41-
NAT_DIR = os.path.join(ROOT_DIR, "src", "nat")
4240

43-
# Work-around for https://github.com/readthedocs/sphinx-autoapi/issues/298
44-
# AutoAPI support for implicit namespaces is broken, so we need to manually
45-
# construct an nat package with an __init__.py file
46-
BUILD_DIR = os.path.join(DOC_DIR, "build")
47-
API_TREE = os.path.join(BUILD_DIR, "_api_tree")
41+
def _build_api_tree() -> Path:
42+
# Work-around for https://github.com/readthedocs/sphinx-autoapi/issues/298
43+
# AutoAPI support for implicit namespaces is broken, so we need to manually
4844

49-
if os.path.exists(API_TREE):
50-
shutil.rmtree(API_TREE)
45+
cur_dir = Path(os.path.abspath(__file__)).parent
46+
docs_dir = cur_dir.parent
47+
root_dir = docs_dir.parent
48+
nat_dir = root_dir / "src" / "nat"
49+
plugins_dir = root_dir / "packages"
5150

52-
os.makedirs(API_TREE)
53-
shutil.copytree(NAT_DIR, os.path.join(API_TREE, "nat"))
54-
with open(os.path.join(API_TREE, "nat", "__init__.py"), "w") as f:
55-
f.write("")
51+
build_dir = docs_dir / "build"
52+
api_tree = build_dir / "_api_tree"
53+
dest_dir = api_tree / "nat"
54+
55+
if api_tree.exists():
56+
shutil.rmtree(api_tree.absolute())
57+
58+
os.makedirs(api_tree.absolute())
59+
shutil.copytree(nat_dir, dest_dir)
60+
dest_plugins_dir = dest_dir / "plugins"
61+
62+
for sub_dir in (dest_dir, dest_plugins_dir):
63+
with open(sub_dir / "__init__.py", "w", encoding="utf-8") as f:
64+
f.write("")
65+
66+
plugin_dirs = [Path(p) for p in glob.glob(f'{plugins_dir}/nvidia_nat_*')]
67+
for plugin_dir in plugin_dirs:
68+
src_dir = plugin_dir / 'src/nat/plugins'
69+
if src_dir.exists():
70+
for plugin_subdir in src_dir.iterdir():
71+
if plugin_subdir.is_dir():
72+
dest_subdir = dest_plugins_dir / plugin_subdir.name
73+
shutil.copytree(plugin_subdir, dest_subdir)
74+
package_file = dest_subdir / "__init__.py"
75+
if not package_file.exists():
76+
with open(package_file, "w", encoding="utf-8") as f:
77+
f.write("")
78+
79+
return api_tree
80+
81+
82+
API_TREE = _build_api_tree()
83+
print(f"API tree built at {API_TREE}")
5684

5785
# -- Project information -----------------------------------------------------
5886

@@ -87,7 +115,7 @@
87115
"sphinxmermaid"
88116
]
89117

90-
autoapi_dirs = [API_TREE]
118+
autoapi_dirs = [str(API_TREE.absolute())]
91119

92120
autoapi_root = "api"
93121
autoapi_python_class_content = "both"

‎packages/nvidia_nat_adk/src/nat/plugins/adk/adk_callback_handler.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ class ADKProfilerHandler(BaseProfilerCallback):
3838
A callback manager/handler for Google ADK that intercepts calls to:
3939
- Tools
4040
- LLMs
41+
4142
to collect usage statistics (tokens, inputs, outputs, time intervals, etc.)
4243
and store them in NeMo Agent Toolkit's usage_stats queue for subsequent analysis.
4344
"""

‎packages/nvidia_nat_agno/src/nat/plugins/agno/tool_wrapper.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -148,7 +148,7 @@ def execute_agno_tool(name: str,
148148
List of required fields for validation
149149
loop : asyncio.AbstractEventLoop
150150
The event loop to use for async execution
151-
**kwargs : Any
151+
kwargs : Any
152152
The arguments to pass to the function
153153
154154
Returns

‎packages/nvidia_nat_crewai/src/nat/plugins/crewai/crewai_callback_handler.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ class CrewAIProfilerHandler(BaseProfilerCallback):
4141
A callback manager/handler for CrewAI that intercepts calls to:
4242
- ToolUsage._use
4343
- LLM Calls
44+
4445
to collect usage statistics (tokens, inputs, outputs, time intervals, etc.)
4546
and store them in NAT's usage_stats queue for subsequent analysis.
4647
"""

‎packages/nvidia_nat_data_flywheel/src/nat/plugins/data_flywheel/observability/exporter/dfw_elasticsearch_exporter.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -43,11 +43,11 @@ def __init__(self,
4343
max_queue_size: The maximum queue size for exporting spans.
4444
drop_on_overflow: Whether to drop spans on overflow.
4545
shutdown_timeout: The shutdown timeout in seconds.
46-
**elasticsearch_kwargs: Additional arguments for ElasticsearchMixin:
47-
- endpoint: The elasticsearch endpoint.
48-
- index: The elasticsearch index name.
49-
- elasticsearch_auth: The elasticsearch authentication credentials.
50-
- headers: The elasticsearch headers.
46+
elasticsearch_kwargs: Additional arguments for ElasticsearchMixin:
47+
- endpoint: The elasticsearch endpoint.
48+
- index: The elasticsearch index name.
49+
- elasticsearch_auth: The elasticsearch authentication credentials.
50+
- headers: The elasticsearch headers.
5151
"""
5252
# Initialize both mixins - ElasticsearchMixin expects elasticsearch_kwargs,
5353
# DFWExporter expects the standard exporter parameters

‎packages/nvidia_nat_data_flywheel/src/nat/plugins/data_flywheel/observability/processor/trace_conversion/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
from .span_extractor import extract_timestamp
1717
from .span_extractor import extract_token_usage
1818
from .span_extractor import extract_usage_info
19-
from .span_to_dfw_record import span_to_dfw_record
19+
from .span_to_dfw import span_to_dfw_record
2020
from .trace_adapter_registry import TraceAdapterRegistry
2121
from .trace_adapter_registry import register_adapter
2222

‎packages/nvidia_nat_data_flywheel/src/nat/plugins/data_flywheel/observability/processor/trace_conversion/adapter/elasticsearch/openai_converter.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -24,12 +24,12 @@
2424
from nat.plugins.data_flywheel.observability.schema.provider.openai_trace_source import OpenAITraceSource
2525
from nat.plugins.data_flywheel.observability.schema.sink.elasticsearch.dfw_es_record import AssistantMessage
2626
from nat.plugins.data_flywheel.observability.schema.sink.elasticsearch.dfw_es_record import DFWESRecord
27+
from nat.plugins.data_flywheel.observability.schema.sink.elasticsearch.dfw_es_record import ESRequest
2728
from nat.plugins.data_flywheel.observability.schema.sink.elasticsearch.dfw_es_record import FinishReason
2829
from nat.plugins.data_flywheel.observability.schema.sink.elasticsearch.dfw_es_record import Function
2930
from nat.plugins.data_flywheel.observability.schema.sink.elasticsearch.dfw_es_record import FunctionDetails
3031
from nat.plugins.data_flywheel.observability.schema.sink.elasticsearch.dfw_es_record import FunctionMessage
3132
from nat.plugins.data_flywheel.observability.schema.sink.elasticsearch.dfw_es_record import Message
32-
from nat.plugins.data_flywheel.observability.schema.sink.elasticsearch.dfw_es_record import Request
3333
from nat.plugins.data_flywheel.observability.schema.sink.elasticsearch.dfw_es_record import RequestTool
3434
from nat.plugins.data_flywheel.observability.schema.sink.elasticsearch.dfw_es_record import Response
3535
from nat.plugins.data_flywheel.observability.schema.sink.elasticsearch.dfw_es_record import ResponseChoice
@@ -77,7 +77,7 @@ def create_message_by_role(role: str, content: str | None, **kwargs) -> Message:
7777
Args:
7878
role (str): The message role
7979
content (str): The message content
80-
**kwargs: Additional role-specific parameters
80+
kwargs: Additional role-specific parameters
8181
8282
Returns:
8383
Message: The appropriate message type for the role
@@ -314,11 +314,11 @@ def convert_langchain_openai(trace_source: TraceContainer) -> DFWESRecord:
314314
temperature = None
315315
max_tokens = None
316316

317-
request = Request(messages=messages,
318-
model=model_name,
319-
tools=request_tools if request_tools else None,
320-
temperature=temperature,
321-
max_tokens=max_tokens)
317+
request = ESRequest(messages=messages,
318+
model=model_name,
319+
tools=request_tools if request_tools else None,
320+
temperature=temperature,
321+
max_tokens=max_tokens)
322322

323323
# Transform chat responses
324324
response_choices = []

packages/nvidia_nat_data_flywheel/src/nat/plugins/data_flywheel/observability/processor/trace_conversion/span_to_dfw_record.py renamed to packages/nvidia_nat_data_flywheel/src/nat/plugins/data_flywheel/observability/processor/trace_conversion/span_to_dfw.py

File renamed without changes.

‎packages/nvidia_nat_data_flywheel/src/nat/plugins/data_flywheel/observability/schema/sink/elasticsearch/dfw_es_record.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,8 @@
2727
from pydantic import model_validator
2828

2929
from nat.plugins.data_flywheel.observability.schema.schema_registry import register_schema
30-
from nat.plugins.data_flywheel.observability.schema.sink.elasticsearch.contract_version import ContractVersion
30+
31+
from .contract_version import ContractVersion
3132

3233
logger = logging.getLogger(__name__)
3334

@@ -135,7 +136,7 @@ class RequestTool(BaseModel):
135136
function: FunctionDetails = Field(..., description="The function details.")
136137

137138

138-
class Request(BaseModel):
139+
class ESRequest(BaseModel):
139140
"""Request structure used in requests."""
140141

141142
model_config = ConfigDict(extra="allow") # Allow extra fields
@@ -199,7 +200,7 @@ class DFWESRecord(BaseModel):
199200
description="Contract version for compatibility tracking")
200201

201202
# Core fields (backward compatible)
202-
request: Request = Field(..., description="The OpenAI ChatCompletion request.")
203+
request: ESRequest = Field(..., description="The OpenAI ChatCompletion request.")
203204
response: Response = Field(..., description="The OpenAI ChatCompletion response.")
204205
client_id: str = Field(..., description="Identifier of the application or deployment that generated traffic.")
205206
workload_id: str = Field(..., description="Stable identifier for the logical task / route / agent node.")

‎packages/nvidia_nat_data_flywheel/src/nat/plugins/data_flywheel/observability/utils/__init__.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.

0 commit comments

Comments
 (0)