nvidia-nat-weave 1.2.0rc5__py3-none-any.whl → 1.2.0rc7__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {aiq → nat}/meta/pypi.md +1 -1
- {aiq → nat}/plugins/weave/register.py +4 -4
- nat/plugins/weave/weave_exporter.py +233 -0
- {nvidia_nat_weave-1.2.0rc5.dist-info → nvidia_nat_weave-1.2.0rc7.dist-info}/METADATA +4 -4
- nvidia_nat_weave-1.2.0rc7.dist-info/RECORD +9 -0
- nvidia_nat_weave-1.2.0rc7.dist-info/entry_points.txt +2 -0
- nvidia_nat_weave-1.2.0rc7.dist-info/top_level.txt +1 -0
- aiq/plugins/weave/mixins/_init__.py +0 -14
- aiq/plugins/weave/mixins/weave_mixin.py +0 -283
- aiq/plugins/weave/weave_exporter.py +0 -33
- nvidia_nat_weave-1.2.0rc5.dist-info/RECORD +0 -11
- nvidia_nat_weave-1.2.0rc5.dist-info/entry_points.txt +0 -2
- nvidia_nat_weave-1.2.0rc5.dist-info/top_level.txt +0 -1
- {aiq → nat}/plugins/weave/__init__.py +0 -0
- {nvidia_nat_weave-1.2.0rc5.dist-info → nvidia_nat_weave-1.2.0rc7.dist-info}/WHEEL +0 -0
{aiq → nat}/meta/pypi.md
RENAMED
@@ -15,7 +15,7 @@ See the License for the specific language governing permissions and
|
|
15
15
|
limitations under the License.
|
16
16
|
-->
|
17
17
|
|
18
|
-

|
19
19
|
|
20
20
|
# NVIDIA NeMo Agent Toolkit Subpackage
|
21
21
|
This is a subpackage for Weights and Biases Weave integration for observability.
|
@@ -17,9 +17,9 @@ import logging
|
|
17
17
|
|
18
18
|
from pydantic import Field
|
19
19
|
|
20
|
-
from
|
21
|
-
from
|
22
|
-
from
|
20
|
+
from nat.builder.builder import Builder
|
21
|
+
from nat.cli.register_workflow import register_telemetry_exporter
|
22
|
+
from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig
|
23
23
|
|
24
24
|
logger = logging.getLogger(__name__)
|
25
25
|
|
@@ -43,7 +43,7 @@ class WeaveTelemetryExporter(TelemetryExporterBaseConfig, name="weave"):
|
|
43
43
|
async def weave_telemetry_exporter(config: WeaveTelemetryExporter, builder: Builder): # pylint: disable=unused-argument
|
44
44
|
import weave
|
45
45
|
|
46
|
-
from
|
46
|
+
from nat.plugins.weave.weave_exporter import WeaveExporter
|
47
47
|
|
48
48
|
weave_settings = {}
|
49
49
|
|
@@ -0,0 +1,233 @@
|
|
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.
|
15
|
+
|
16
|
+
import logging
|
17
|
+
from collections.abc import Generator
|
18
|
+
from contextlib import contextmanager
|
19
|
+
|
20
|
+
from weave.trace.context import weave_client_context
|
21
|
+
from weave.trace.context.call_context import get_current_call
|
22
|
+
from weave.trace.context.call_context import set_call_stack
|
23
|
+
from weave.trace.weave_client import Call
|
24
|
+
|
25
|
+
from nat.data_models.intermediate_step import IntermediateStep
|
26
|
+
from nat.data_models.span import Span
|
27
|
+
from nat.observability.exporter.base_exporter import IsolatedAttribute
|
28
|
+
from nat.observability.exporter.span_exporter import SpanExporter
|
29
|
+
from nat.utils.log_utils import LogFilter
|
30
|
+
from nat.utils.type_utils import override
|
31
|
+
|
32
|
+
logger = logging.getLogger(__name__)
|
33
|
+
|
34
|
+
# Use LogFilter to filter out specific message patterns
|
35
|
+
presidio_filter = LogFilter([
|
36
|
+
"nlp_engine not provided",
|
37
|
+
"Created NLP engine",
|
38
|
+
"registry not provided",
|
39
|
+
"Loaded recognizer",
|
40
|
+
"Recognizer not added to registry"
|
41
|
+
])
|
42
|
+
|
43
|
+
|
44
|
+
class WeaveExporter(SpanExporter[Span, Span]):
|
45
|
+
"""A Weave exporter that exports telemetry traces to Weights & Biases Weave using OpenTelemetry."""
|
46
|
+
|
47
|
+
_weave_calls: IsolatedAttribute[dict[str, Call]] = IsolatedAttribute(dict)
|
48
|
+
|
49
|
+
def __init__(self,
|
50
|
+
context_state=None,
|
51
|
+
entity: str | None = None,
|
52
|
+
project: str | None = None,
|
53
|
+
verbose: bool = False):
|
54
|
+
super().__init__(context_state=context_state)
|
55
|
+
self._entity = entity
|
56
|
+
self._project = project
|
57
|
+
self._gc = weave_client_context.require_weave_client()
|
58
|
+
|
59
|
+
# Optionally, set log filtering for presidio-analyzer to reduce verbosity
|
60
|
+
if not verbose:
|
61
|
+
presidio_logger = logging.getLogger('presidio-analyzer')
|
62
|
+
presidio_logger.addFilter(presidio_filter)
|
63
|
+
|
64
|
+
@override
|
65
|
+
async def export_processed(self, item: Span | list[Span]) -> None:
|
66
|
+
"""Dummy implementation of export_processed.
|
67
|
+
|
68
|
+
Args:
|
69
|
+
item (Span | list[Span]): The span or list of spans to export.
|
70
|
+
"""
|
71
|
+
pass
|
72
|
+
|
73
|
+
def _process_start_event(self, event: IntermediateStep):
|
74
|
+
"""Process the start event for a Weave call.
|
75
|
+
|
76
|
+
Args:
|
77
|
+
event (IntermediateStep): The intermediate step event.
|
78
|
+
"""
|
79
|
+
super()._process_start_event(event)
|
80
|
+
span = self._span_stack.get(event.UUID, None)
|
81
|
+
if span is None:
|
82
|
+
logger.warning("No span found for event %s", event.UUID)
|
83
|
+
return
|
84
|
+
self._create_weave_call(event, span)
|
85
|
+
|
86
|
+
def _process_end_event(self, event: IntermediateStep):
|
87
|
+
"""Process the end event for a Weave call.
|
88
|
+
|
89
|
+
Args:
|
90
|
+
event (IntermediateStep): The intermediate step event.
|
91
|
+
"""
|
92
|
+
super()._process_end_event(event)
|
93
|
+
self._finish_weave_call(event)
|
94
|
+
|
95
|
+
@contextmanager
|
96
|
+
def parent_call(self, trace_id: str, parent_call_id: str) -> Generator[None]:
|
97
|
+
"""Create a dummy Weave call for the parent span.
|
98
|
+
|
99
|
+
Args:
|
100
|
+
trace_id (str): The trace ID of the parent span.
|
101
|
+
parent_call_id (str): The ID of the parent call.
|
102
|
+
|
103
|
+
Yields:
|
104
|
+
None: The dummy Weave call.
|
105
|
+
"""
|
106
|
+
dummy_call = Call(trace_id=trace_id, id=parent_call_id, _op_name="", project_id="", parent_id=None, inputs={})
|
107
|
+
with set_call_stack([dummy_call]):
|
108
|
+
yield
|
109
|
+
|
110
|
+
def _create_weave_call(self, step: IntermediateStep, span: Span) -> Call:
|
111
|
+
"""
|
112
|
+
Create a Weave call directly from the span and step data,
|
113
|
+
connecting to existing framework traces if available.
|
114
|
+
|
115
|
+
Args:
|
116
|
+
step (IntermediateStep): The intermediate step event.
|
117
|
+
span (Span): The span associated with the intermediate step.
|
118
|
+
|
119
|
+
Returns:
|
120
|
+
Call: The Weave call created from the span and step data.
|
121
|
+
"""
|
122
|
+
# Check for existing Weave trace/call
|
123
|
+
existing_call = get_current_call()
|
124
|
+
|
125
|
+
# Extract parent call if applicable
|
126
|
+
parent_call = None
|
127
|
+
|
128
|
+
# If we have an existing Weave call from another framework (e.g., LangChain),
|
129
|
+
# use it as the parent
|
130
|
+
if existing_call is not None:
|
131
|
+
parent_call = existing_call
|
132
|
+
logger.debug("Found existing Weave call: %s from trace: %s", existing_call.id, existing_call.trace_id)
|
133
|
+
# Otherwise, check our internal stack for parent relationships
|
134
|
+
elif len(self._weave_calls) > 0 and len(self._span_stack) > 1:
|
135
|
+
# Get the parent span using stack position (one level up)
|
136
|
+
parent_span_id = self._span_stack[-2].context.span_id
|
137
|
+
# Find the corresponding weave call for this parent span
|
138
|
+
for call in self._weave_calls.values():
|
139
|
+
if getattr(call, "span_id", None) == parent_span_id:
|
140
|
+
parent_call = call
|
141
|
+
break
|
142
|
+
|
143
|
+
# Generate a meaningful operation name based on event type
|
144
|
+
event_type = step.payload.event_type.split(".")[-1]
|
145
|
+
if step.payload.name:
|
146
|
+
op_name = f"aiq.{event_type}.{step.payload.name}"
|
147
|
+
else:
|
148
|
+
op_name = f"aiq.{event_type}"
|
149
|
+
|
150
|
+
# Create input dictionary
|
151
|
+
inputs = {}
|
152
|
+
if step.payload.data and step.payload.data.input is not None:
|
153
|
+
try:
|
154
|
+
# Add the input to the Weave call
|
155
|
+
inputs["input"] = step.payload.data.input
|
156
|
+
except Exception:
|
157
|
+
# If serialization fails, use string representation
|
158
|
+
inputs["input"] = str(step.payload.data.input)
|
159
|
+
|
160
|
+
# Create the Weave call
|
161
|
+
call = self._gc.create_call(
|
162
|
+
op_name,
|
163
|
+
inputs=inputs,
|
164
|
+
parent=parent_call,
|
165
|
+
attributes=span.attributes,
|
166
|
+
display_name=op_name,
|
167
|
+
)
|
168
|
+
|
169
|
+
# Store the call with step UUID as key
|
170
|
+
self._weave_calls[step.UUID] = call
|
171
|
+
|
172
|
+
# Store span ID for parent reference
|
173
|
+
if span.context is not None:
|
174
|
+
setattr(call, "span_id", span.context.span_id)
|
175
|
+
else:
|
176
|
+
logger.warning("Span has no context, skipping span_id setting")
|
177
|
+
|
178
|
+
return call
|
179
|
+
|
180
|
+
def _finish_weave_call(self, step: IntermediateStep) -> None:
|
181
|
+
"""
|
182
|
+
Finish a previously created Weave call.
|
183
|
+
|
184
|
+
Args:
|
185
|
+
step (IntermediateStep): The intermediate step event.
|
186
|
+
"""
|
187
|
+
# Find the call for this step
|
188
|
+
call = self._weave_calls.pop(step.UUID, None)
|
189
|
+
|
190
|
+
if call is None:
|
191
|
+
logger.warning("No Weave call found for step %s", step.UUID)
|
192
|
+
return
|
193
|
+
|
194
|
+
# Create output dictionary
|
195
|
+
outputs = {}
|
196
|
+
if step.payload.data and step.payload.data.output is not None:
|
197
|
+
try:
|
198
|
+
# Add the output to the Weave call
|
199
|
+
outputs["output"] = step.payload.data.output
|
200
|
+
except Exception:
|
201
|
+
# If serialization fails, use string representation
|
202
|
+
outputs["output"] = str(step.payload.data.output)
|
203
|
+
|
204
|
+
# Add usage information if available
|
205
|
+
usage_info = step.payload.usage_info
|
206
|
+
if usage_info:
|
207
|
+
if usage_info.token_usage:
|
208
|
+
outputs["prompt_tokens"] = usage_info.token_usage.prompt_tokens
|
209
|
+
outputs["completion_tokens"] = usage_info.token_usage.completion_tokens
|
210
|
+
outputs["total_tokens"] = usage_info.token_usage.total_tokens
|
211
|
+
|
212
|
+
if usage_info.num_llm_calls:
|
213
|
+
outputs["num_llm_calls"] = usage_info.num_llm_calls
|
214
|
+
|
215
|
+
if usage_info.seconds_between_calls:
|
216
|
+
outputs["seconds_between_calls"] = usage_info.seconds_between_calls
|
217
|
+
|
218
|
+
# Finish the call with outputs
|
219
|
+
self._gc.finish_call(call, outputs)
|
220
|
+
|
221
|
+
async def _cleanup_weave_calls(self) -> None:
|
222
|
+
"""
|
223
|
+
Clean up any lingering Weave calls.
|
224
|
+
"""
|
225
|
+
if self._weave_calls:
|
226
|
+
for _, call in list(self._weave_calls.items()):
|
227
|
+
self._gc.finish_call(call, {"status": "incomplete"})
|
228
|
+
self._weave_calls.clear()
|
229
|
+
|
230
|
+
async def _cleanup(self) -> None:
|
231
|
+
"""Perform cleanup once the exporter is stopped."""
|
232
|
+
await self._cleanup_weave_calls()
|
233
|
+
await super()._cleanup()
|
@@ -1,12 +1,12 @@
|
|
1
1
|
Metadata-Version: 2.4
|
2
2
|
Name: nvidia-nat-weave
|
3
|
-
Version: 1.2.
|
4
|
-
Summary: Subpackage for Weave integration in
|
3
|
+
Version: 1.2.0rc7
|
4
|
+
Summary: Subpackage for Weave integration in NeMo Agent toolkit
|
5
5
|
Keywords: ai,observability,wandb,pii
|
6
6
|
Classifier: Programming Language :: Python
|
7
7
|
Requires-Python: <3.13,>=3.11
|
8
8
|
Description-Content-Type: text/markdown
|
9
|
-
Requires-Dist: nvidia-nat
|
9
|
+
Requires-Dist: nvidia-nat==v1.2.0-rc7
|
10
10
|
Requires-Dist: presidio-analyzer~=2.2
|
11
11
|
Requires-Dist: presidio-anonymizer~=2.2
|
12
12
|
Requires-Dist: weave~=0.51
|
@@ -28,7 +28,7 @@ See the License for the specific language governing permissions and
|
|
28
28
|
limitations under the License.
|
29
29
|
-->
|
30
30
|
|
31
|
-

|
32
32
|
|
33
33
|
# NVIDIA NeMo Agent Toolkit Subpackage
|
34
34
|
This is a subpackage for Weights and Biases Weave integration for observability.
|
@@ -0,0 +1,9 @@
|
|
1
|
+
nat/meta/pypi.md,sha256=FVQR5lfZjqZHm4VWMmQJYgZbwJJjrfQZgEqHscDuMR8,1121
|
2
|
+
nat/plugins/weave/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
3
|
+
nat/plugins/weave/register.py,sha256=VSecYad8SJDzlmfqc3y7CvsvtaCs9-yMaMZScJ59-Ck,3216
|
4
|
+
nat/plugins/weave/weave_exporter.py,sha256=NC_q3wc_KKbHTN7vgLEHmbP8PV8HQnwfapnZk-ZzkXI,8789
|
5
|
+
nvidia_nat_weave-1.2.0rc7.dist-info/METADATA,sha256=U8xzu0ZrDppzKhINiT4Og3041ztYxMb9y5tph_EvEkA,1542
|
6
|
+
nvidia_nat_weave-1.2.0rc7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
|
7
|
+
nvidia_nat_weave-1.2.0rc7.dist-info/entry_points.txt,sha256=xg4vW3wKsOLfHa-QR6JqWnq3DmdfI_z9IZfg4I9thTY,56
|
8
|
+
nvidia_nat_weave-1.2.0rc7.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
|
9
|
+
nvidia_nat_weave-1.2.0rc7.dist-info/RECORD,,
|
@@ -0,0 +1 @@
|
|
1
|
+
nat
|
@@ -1,14 +0,0 @@
|
|
1
|
-
# SPDX-FileCopyrightText: Copyright (c) 2024-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.
|
@@ -1,283 +0,0 @@
|
|
1
|
-
# SPDX-FileCopyrightText: Copyright (c) 2024-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.
|
15
|
-
|
16
|
-
import logging
|
17
|
-
from collections.abc import Generator
|
18
|
-
from contextlib import contextmanager
|
19
|
-
|
20
|
-
from weave.trace.context import weave_client_context
|
21
|
-
from weave.trace.context.call_context import get_current_call
|
22
|
-
from weave.trace.context.call_context import set_call_stack
|
23
|
-
from weave.trace.weave_client import Call
|
24
|
-
|
25
|
-
from aiq.data_models.span import Span
|
26
|
-
from aiq.data_models.span import SpanAttributes
|
27
|
-
from aiq.observability.exporter.base_exporter import IsolatedAttribute
|
28
|
-
from aiq.utils.log_utils import LogFilter
|
29
|
-
|
30
|
-
logger = logging.getLogger(__name__)
|
31
|
-
|
32
|
-
# Alternative: Use LogFilter to filter specific message patterns
|
33
|
-
presidio_filter = LogFilter([
|
34
|
-
"nlp_engine not provided",
|
35
|
-
"Created NLP engine",
|
36
|
-
"registry not provided",
|
37
|
-
"Loaded recognizer",
|
38
|
-
"Recognizer not added to registry"
|
39
|
-
])
|
40
|
-
|
41
|
-
|
42
|
-
class WeaveMixin:
|
43
|
-
"""Mixin for Weave exporters.
|
44
|
-
|
45
|
-
This mixin provides a default implementation of the export method for Weave exporters.
|
46
|
-
It uses the weave_client_context to create and finish Weave calls.
|
47
|
-
|
48
|
-
Args:
|
49
|
-
project (str): The project name to group the telemetry traces.
|
50
|
-
entity (str | None): The entity name to group the telemetry traces.
|
51
|
-
"""
|
52
|
-
|
53
|
-
_weave_calls: IsolatedAttribute[dict[int, Call]] = IsolatedAttribute(dict)
|
54
|
-
_in_flight_calls: IsolatedAttribute[set[int]] = IsolatedAttribute(set)
|
55
|
-
|
56
|
-
def __init__(self, *args, project: str, entity: str | None = None, verbose: bool = False, **kwargs):
|
57
|
-
"""Initialize the Weave exporter with the specified project and entity.
|
58
|
-
|
59
|
-
Args:
|
60
|
-
project (str): The project name to group the telemetry traces.
|
61
|
-
entity (str | None): The entity name to group the telemetry traces.
|
62
|
-
"""
|
63
|
-
self._gc = weave_client_context.require_weave_client()
|
64
|
-
self._project = project
|
65
|
-
self._entity = entity
|
66
|
-
|
67
|
-
# Optionally, set log filtering for presidio-analyzer to reduce verbosity
|
68
|
-
if not verbose:
|
69
|
-
presidio_logger = logging.getLogger('presidio-analyzer')
|
70
|
-
presidio_logger.addFilter(presidio_filter)
|
71
|
-
|
72
|
-
super().__init__(*args, **kwargs)
|
73
|
-
|
74
|
-
async def export_processed(self, item: Span | list[Span]) -> None:
|
75
|
-
"""Export a batch of spans.
|
76
|
-
|
77
|
-
Args:
|
78
|
-
item (Span | list[Span]): The span or list of spans to export.
|
79
|
-
"""
|
80
|
-
if not isinstance(item, list):
|
81
|
-
spans = [item]
|
82
|
-
else:
|
83
|
-
spans = item
|
84
|
-
|
85
|
-
for span in spans:
|
86
|
-
self._export_processed(span)
|
87
|
-
|
88
|
-
def _export_processed(self, span: Span) -> None:
|
89
|
-
"""Export a single span.
|
90
|
-
|
91
|
-
Args:
|
92
|
-
span (Span): The span to export.
|
93
|
-
"""
|
94
|
-
span_id = span.context.span_id if span.context else None # type: ignore
|
95
|
-
if span_id is None:
|
96
|
-
logger.warning("Span has no context or span_id, skipping export")
|
97
|
-
return
|
98
|
-
|
99
|
-
try:
|
100
|
-
call = self._create_weave_call(span)
|
101
|
-
self._finish_weave_call(call, span)
|
102
|
-
except Exception as e:
|
103
|
-
logger.error("Error exporting spans: %s", e, exc_info=True)
|
104
|
-
# Clean up in-flight tracking if call creation/finishing failed
|
105
|
-
self._in_flight_calls.discard(span_id)
|
106
|
-
|
107
|
-
@contextmanager
|
108
|
-
def parent_call(self, trace_id: str, parent_call_id: str) -> Generator[None]:
|
109
|
-
"""Create a dummy Weave call for the parent span.
|
110
|
-
|
111
|
-
Args:
|
112
|
-
trace_id (str): The trace ID of the parent span.
|
113
|
-
parent_call_id (str): The ID of the parent call.
|
114
|
-
|
115
|
-
Yields:
|
116
|
-
None: The dummy Weave call.
|
117
|
-
"""
|
118
|
-
dummy_call = Call(trace_id=trace_id, id=parent_call_id, _op_name="", project_id="", parent_id=None, inputs={})
|
119
|
-
with set_call_stack([dummy_call]):
|
120
|
-
yield
|
121
|
-
|
122
|
-
def _create_weave_call(self, span: Span) -> Call:
|
123
|
-
"""
|
124
|
-
Create a Weave call directly from the span and step data, connecting to existing framework traces if available.
|
125
|
-
|
126
|
-
Args:
|
127
|
-
span (Span): The span to create a Weave call for.
|
128
|
-
|
129
|
-
Returns:
|
130
|
-
Call: The Weave call.
|
131
|
-
"""
|
132
|
-
span_id = span.context.span_id if span.context else None # type: ignore
|
133
|
-
if span_id is None:
|
134
|
-
raise ValueError("Span has no context or span_id")
|
135
|
-
|
136
|
-
# Mark this call as in-flight to prevent premature cleanup
|
137
|
-
self._in_flight_calls.add(span_id)
|
138
|
-
|
139
|
-
# Check for existing Weave trace/call
|
140
|
-
existing_call = get_current_call()
|
141
|
-
|
142
|
-
# Extract parent call if applicable
|
143
|
-
parent_call = None
|
144
|
-
|
145
|
-
# If we have an existing Weave call from another framework (e.g., LangChain),
|
146
|
-
# use it as the parent (but only if it's actually a Call object)
|
147
|
-
if existing_call is not None and isinstance(existing_call, Call):
|
148
|
-
parent_call = existing_call
|
149
|
-
logger.debug("Found existing Weave call: %s from trace: %s", existing_call.id, existing_call.trace_id)
|
150
|
-
# Otherwise, check our internal stack for parent relationships
|
151
|
-
elif len(self._weave_calls) > 0:
|
152
|
-
# Get the parent span using stack position (one level up)
|
153
|
-
parent_span_id = (span.parent.context.span_id if span.parent and span.parent.context else None
|
154
|
-
) # type: ignore
|
155
|
-
if parent_span_id:
|
156
|
-
# Find the corresponding weave call for this parent span
|
157
|
-
for call in self._weave_calls.values():
|
158
|
-
if getattr(call, "span_id", None) == parent_span_id:
|
159
|
-
parent_call = call
|
160
|
-
break
|
161
|
-
|
162
|
-
# Generate a meaningful operation name based on event type
|
163
|
-
span_event_type = span.attributes.get(SpanAttributes.AIQ_EVENT_TYPE.value, "unknown")
|
164
|
-
event_type = span_event_type.split(".")[-1]
|
165
|
-
if span.name:
|
166
|
-
op_name = f"aiq.{event_type}.{span.name}"
|
167
|
-
else:
|
168
|
-
op_name = f"aiq.{event_type}"
|
169
|
-
|
170
|
-
# Create input dictionary
|
171
|
-
inputs = {}
|
172
|
-
input_value = span.attributes.get(SpanAttributes.INPUT_VALUE.value)
|
173
|
-
if input_value is not None:
|
174
|
-
try:
|
175
|
-
# Add the input to the Weave call
|
176
|
-
inputs["input"] = input_value
|
177
|
-
except Exception:
|
178
|
-
# If serialization fails, use string representation
|
179
|
-
inputs["input"] = str(input_value)
|
180
|
-
|
181
|
-
# Create the Weave call
|
182
|
-
call = self._gc.create_call(
|
183
|
-
op_name,
|
184
|
-
inputs=inputs,
|
185
|
-
parent=parent_call,
|
186
|
-
attributes=span.attributes,
|
187
|
-
display_name=op_name,
|
188
|
-
)
|
189
|
-
|
190
|
-
# Store the call with span span ID as key
|
191
|
-
self._weave_calls[span_id] = call
|
192
|
-
|
193
|
-
# Store span ID for parent reference
|
194
|
-
setattr(call, "span_id", span_id)
|
195
|
-
|
196
|
-
return call
|
197
|
-
|
198
|
-
def _finish_weave_call(self, call: Call, span: Span):
|
199
|
-
"""Finish a previously created Weave call.
|
200
|
-
|
201
|
-
Args:
|
202
|
-
call (Call): The Weave call to finish.
|
203
|
-
span (Span): The span to finish the call for.
|
204
|
-
"""
|
205
|
-
span_id = span.context.span_id if span.context else None # type: ignore
|
206
|
-
if span_id is None:
|
207
|
-
logger.warning("Span has no context or span_id")
|
208
|
-
return
|
209
|
-
|
210
|
-
if call is None:
|
211
|
-
logger.warning("No Weave call found for span %s", span_id)
|
212
|
-
# Still remove from in-flight tracking
|
213
|
-
self._in_flight_calls.discard(span_id)
|
214
|
-
return
|
215
|
-
|
216
|
-
# Check if this call was already finished by cleanup (race condition protection)
|
217
|
-
if span_id not in self._weave_calls:
|
218
|
-
logger.debug("Call for span %s was already finished (likely by cleanup)", span_id)
|
219
|
-
self._in_flight_calls.discard(span_id)
|
220
|
-
return
|
221
|
-
|
222
|
-
# Create output dictionary
|
223
|
-
outputs = {}
|
224
|
-
output = span.attributes.get(SpanAttributes.OUTPUT_VALUE.value)
|
225
|
-
if output is not None:
|
226
|
-
try:
|
227
|
-
# Add the output to the Weave call
|
228
|
-
outputs["output"] = output
|
229
|
-
except Exception:
|
230
|
-
# If serialization fails, use string representation
|
231
|
-
outputs["output"] = str(output)
|
232
|
-
|
233
|
-
# Add usage information
|
234
|
-
outputs["prompt_tokens"] = span.attributes.get(SpanAttributes.LLM_TOKEN_COUNT_PROMPT.value)
|
235
|
-
outputs["completion_tokens"] = span.attributes.get(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION.value)
|
236
|
-
outputs["total_tokens"] = span.attributes.get(SpanAttributes.LLM_TOKEN_COUNT_TOTAL.value)
|
237
|
-
outputs["num_llm_calls"] = span.attributes.get(SpanAttributes.AIQ_USAGE_NUM_LLM_CALLS.value)
|
238
|
-
outputs["seconds_between_calls"] = span.attributes.get(SpanAttributes.AIQ_USAGE_SECONDS_BETWEEN_CALLS.value)
|
239
|
-
|
240
|
-
try:
|
241
|
-
# Finish the call with outputs
|
242
|
-
self._gc.finish_call(call, outputs)
|
243
|
-
logger.debug("Successfully finished call for span %s", span_id)
|
244
|
-
except Exception as e:
|
245
|
-
logger.warning("Error finishing call for span %s: %s", span_id, e)
|
246
|
-
finally:
|
247
|
-
# Always clean up tracking regardless of finish success/failure
|
248
|
-
self._weave_calls.pop(span_id, None)
|
249
|
-
self._in_flight_calls.discard(span_id)
|
250
|
-
|
251
|
-
async def _cleanup_weave_calls(self) -> None:
|
252
|
-
"""Clean up any lingering unfinished Weave calls.
|
253
|
-
|
254
|
-
This method should only be called during exporter shutdown to handle
|
255
|
-
calls that weren't properly finished during normal operation.
|
256
|
-
|
257
|
-
CRITICAL: Only cleans up calls that are NOT currently in-flight to prevent
|
258
|
-
race conditions with background export tasks.
|
259
|
-
"""
|
260
|
-
if self._gc is not None and self._weave_calls:
|
261
|
-
# Only clean up calls that are not currently being processed
|
262
|
-
abandoned_calls = {}
|
263
|
-
for span_id, call in self._weave_calls.items():
|
264
|
-
if span_id not in self._in_flight_calls:
|
265
|
-
abandoned_calls[span_id] = call
|
266
|
-
|
267
|
-
if abandoned_calls:
|
268
|
-
logger.debug("Cleaning up %d truly abandoned Weave calls (out of %d total)",
|
269
|
-
len(abandoned_calls),
|
270
|
-
len(self._weave_calls))
|
271
|
-
|
272
|
-
for span_id, call in abandoned_calls.items():
|
273
|
-
try:
|
274
|
-
# Finish any remaining calls with incomplete status
|
275
|
-
self._gc.finish_call(call, {"status": "incomplete"})
|
276
|
-
logger.debug("Finished abandoned call for span %s", span_id)
|
277
|
-
except Exception as e:
|
278
|
-
logger.warning("Error finishing abandoned call for span %s: %s", span_id, e)
|
279
|
-
finally:
|
280
|
-
# Remove from tracking
|
281
|
-
self._weave_calls.pop(span_id, None)
|
282
|
-
else:
|
283
|
-
logger.debug("No abandoned calls to clean up (%d calls still in-flight)", len(self._in_flight_calls))
|
@@ -1,33 +0,0 @@
|
|
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.
|
15
|
-
|
16
|
-
import logging
|
17
|
-
|
18
|
-
from aiq.data_models.span import Span
|
19
|
-
from aiq.observability.exporter.span_exporter import SpanExporter
|
20
|
-
from aiq.plugins.weave.mixins.weave_mixin import WeaveMixin
|
21
|
-
|
22
|
-
logger = logging.getLogger(__name__)
|
23
|
-
|
24
|
-
|
25
|
-
class WeaveExporter(WeaveMixin, SpanExporter[Span, Span]): # pylint: disable=R0901
|
26
|
-
"""A Weave exporter that exports telemetry traces to Weights & Biases Weave using OpenTelemetry."""
|
27
|
-
|
28
|
-
def __init__(self, context_state=None, **weave_kwargs):
|
29
|
-
super().__init__(context_state=context_state, **weave_kwargs)
|
30
|
-
|
31
|
-
async def _cleanup(self) -> None:
|
32
|
-
await self._cleanup_weave_calls()
|
33
|
-
await super()._cleanup()
|
@@ -1,11 +0,0 @@
|
|
1
|
-
aiq/meta/pypi.md,sha256=cVwHYnvOVN3K-yyZzrMa2pMObusdyEE0DYkxC6R4Dys,1131
|
2
|
-
aiq/plugins/weave/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
3
|
-
aiq/plugins/weave/register.py,sha256=WYRkxPaptb9itFJVlWMzFA4hpfPu0iH4RkPsz4Sq0uE,3216
|
4
|
-
aiq/plugins/weave/weave_exporter.py,sha256=QT8FsQ5Ec8bRv8CVrwgmBWHtVucr3kWhTXPdX3hZcSg,1334
|
5
|
-
aiq/plugins/weave/mixins/_init__.py,sha256=Xs1JQ16L9btwreh4pdGKwskffAw1YFO48jKrU4ib_7c,685
|
6
|
-
aiq/plugins/weave/mixins/weave_mixin.py,sha256=DfwWF4ioLEcepbrGhHDr2B81pPaEeEQu1xzsT6XZQIU,11561
|
7
|
-
nvidia_nat_weave-1.2.0rc5.dist-info/METADATA,sha256=RX-O4NHSBzGqI8H3KF2F0CZzYm7ZzFIPfmv6uHWG_-w,1537
|
8
|
-
nvidia_nat_weave-1.2.0rc5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
|
9
|
-
nvidia_nat_weave-1.2.0rc5.dist-info/entry_points.txt,sha256=1VTpfzbrsKofAPZlmzUF8gERdyvMgK91dngr7NneYIM,56
|
10
|
-
nvidia_nat_weave-1.2.0rc5.dist-info/top_level.txt,sha256=fo7AzYcNhZ_tRWrhGumtxwnxMew4xrT1iwouDy_f0Kc,4
|
11
|
-
nvidia_nat_weave-1.2.0rc5.dist-info/RECORD,,
|
@@ -1 +0,0 @@
|
|
1
|
-
aiq
|
File without changes
|
File without changes
|