Skip to content

Commit 175de3b

Browse files
akash-vijay-kvakash-vijay-kv
committed
[NET-995] feat: Add utility for adding streaming output on root span (#303)
1 parent e055fa1 commit 175de3b

12 files changed

Lines changed: 403 additions & 27 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -285,4 +285,4 @@ Users can be now overwrite the input and ouput attributes of spans created by in
285285

286286
- Added utility to set input and output data for any active span in a trace
287287

288-
[0.1.87]: https://github.com/KeyValueSoftwareSystems/netra-sdk-py/tree/main
288+
[0.1.89]: https://github.com/KeyValueSoftwareSystems/netra-sdk-py/tree/main

‎netra/__init__.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -508,6 +508,28 @@ def set_root_output(cls, value: Any) -> None:
508508
"""
509509
SessionManager.set_root_output(value)
510510

511+
@classmethod
512+
def set_root_output_stream(cls, value: Any) -> Any:
513+
"""
514+
Wrap a stream so the accumulated output is set on the root span when iteration ends.
515+
516+
The returned object is a transparent proxy — iterate over it instead of the original::
517+
518+
stream = Netra.set_root_output_stream(stream)
519+
for chunk in stream:
520+
...
521+
522+
Supports both sync and async iterables. Returns *value* unchanged if no active trace
523+
context exists or if *value* is not iterable.
524+
525+
Args:
526+
value: The stream to wrap (Netra-instrumented or any generic iterable).
527+
528+
Returns:
529+
A wrapped stream proxy, or *value* unchanged if wrapping is not possible.
530+
"""
531+
return SessionManager.set_root_output_stream(value)
532+
511533
@classmethod
512534
def start_span(
513535
cls,

‎netra/instrumentation/agno/utils.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -941,7 +941,7 @@ def set_request_attributes(
941941
span.set_attribute("input", input_content)
942942

943943

944-
def set_response_attributes(span: Span, response: Any) -> None:
944+
def set_response_attributes(span: Span, response: Any) -> Optional[str]:
945945
"""Set response-side span attributes from an Agno response object.
946946
947947
Writes token usage, output content, response ID, and output type.
@@ -951,7 +951,7 @@ def set_response_attributes(span: Span, response: Any) -> None:
951951
response: The Agno response object (RunResponse, TeamRunOutput, etc.).
952952
"""
953953
if not span.is_recording():
954-
return
954+
return None
955955

956956
usage = extract_token_usage(response)
957957
if usage:
@@ -965,6 +965,8 @@ def set_response_attributes(span: Span, response: Any) -> None:
965965
if response_id:
966966
span.set_attribute(ATTR_RESPONSE_ID, response_id)
967967

968+
return output
969+
968970

969971
def sanitize_headers(raw_headers: List[Tuple[bytes, bytes]]) -> Dict[str, str]:
970972
"""Convert ASGI raw header pairs to a dict with sensitive values redacted.

‎netra/instrumentation/agno/wrappers.py‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,8 @@ def _set_common_span_attributes(span: Span, entity_type: str) -> None:
157157
class _BaseStreamWrapper:
158158
"""Shared base for all span streaming wrappers."""
159159

160+
_netra_stream_wrapper = True
161+
160162
def __init__(self, span: Span, response: Any, ctx_token: Any = None) -> None:
161163
"""Initialise the streaming wrapper.
162164
@@ -222,10 +224,14 @@ class _AgentStreamOutputMixin:
222224

223225
def _set_output_on_success(self) -> None:
224226
"""Write accumulated run content and token usage to the span before it closes."""
227+
self._netra_output = ""
225228
if self._last_response is not None:
226-
set_response_attributes(self._span, self._last_response)
229+
output = set_response_attributes(self._span, self._last_response)
230+
self._netra_output = output if output else ""
227231
if self._content_chunks:
228-
self._span.set_attribute("output", "".join(self._content_chunks))
232+
output = "".join(self._content_chunks)
233+
self._span.set_attribute("output", output)
234+
self._netra_output = output
229235

230236

231237
class _LlmStreamOutputMixin:
@@ -239,9 +245,11 @@ class _LlmStreamOutputMixin:
239245
def _set_output_on_success(self) -> None:
240246
"""Write accumulated LLM content, token usage, and timing metrics to the span."""
241247
output_str = None
248+
self._netra_output = ""
242249
if self._content_chunks:
243250
content = "".join(self._content_chunks)
244251
output_str = json.dumps([{"role": "assistant", "content": content}])
252+
self._netra_output = content
245253
elif self._tool_calls:
246254
try:
247255
tc_serialized = serialize_value(self._tool_calls, clean=True)
@@ -251,10 +259,12 @@ def _set_output_on_success(self) -> None:
251259
except (json.JSONDecodeError, ValueError):
252260
tc_data = tc_serialized
253261
output_str = json.dumps([{"role": "assistant", "tool_calls": tc_data}])
262+
self._netra_output = tc_serialized
254263
except Exception as e:
255264
logger.debug("netra.instrumentation.agno: failed to serialize tool_calls for LLM output: %s", e)
256265
elif self._last_response is not None:
257266
output_str = format_response_as_output(self._last_response)
267+
self._netra_output = output_str if output_str else ""
258268
if output_str:
259269
self._span.set_attribute("output", output_str)
260270
set_llm_completion_attributes(self._span, output_str)

‎netra/instrumentation/cerebras/wrappers.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@ def _detect_streaming(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> bool:
4040
class StreamingWrapper(ObjectProxy): # type: ignore[misc]
4141
"""Wrapper for streaming responses"""
4242

43+
_netra_stream_wrapper = True
44+
4345
def __init__(self, span: Span, response: Iterator[Any], request_kwargs: Dict[str, Any]) -> None:
4446
super().__init__(response)
4547
self._span = span
@@ -59,6 +61,15 @@ def _ensure_choice(self, index: int) -> None:
5961
else:
6062
self._complete_response["choices"].append({"text": ""})
6163

64+
def _extract_content_text(self) -> str:
65+
"""Extract the plain text content from the accumulated response."""
66+
parts = []
67+
for choice in self._complete_response.get("choices", []):
68+
msg = choice.get("message", {})
69+
if content := msg.get("content"):
70+
parts.append(content)
71+
return "".join(parts)
72+
6273
def __iter__(self) -> Iterator[Any]:
6374
return self
6475

@@ -129,13 +140,16 @@ def _finalize_span(self) -> None:
129140
"""Finalize span when streaming is complete"""
130141
record_span_timing(self._span, LLM_RESPONSE_DURATION)
131142
set_response_attributes(self._span, self._complete_response)
143+
self._netra_output = self._extract_content_text()
132144
self._span.set_status(Status(StatusCode.OK))
133145
self._span.end()
134146

135147

136148
class AsyncStreamingWrapper(ObjectProxy): # type: ignore[misc]
137149
"""Async wrapper for streaming responses"""
138150

151+
_netra_stream_wrapper = True
152+
139153
def __init__(self, span: Span, response: AsyncIterator[Any], request_kwargs: Dict[str, Any]) -> None:
140154
super().__init__(response)
141155
self._span = span
@@ -155,6 +169,15 @@ def _ensure_choice(self, index: int) -> None:
155169
else:
156170
self._complete_response["choices"].append({"text": ""})
157171

172+
def _extract_content_text(self) -> str:
173+
"""Extract the plain text content from the accumulated response."""
174+
parts = []
175+
for choice in self._complete_response.get("choices", []):
176+
msg = choice.get("message", {})
177+
if content := msg.get("content"):
178+
parts.append(content)
179+
return "".join(parts)
180+
158181
def __aiter__(self) -> AsyncIterator[Any]:
159182
return self
160183

@@ -227,6 +250,7 @@ def _finalize_span(self) -> None:
227250
"""Finalize span when streaming is complete"""
228251
record_span_timing(self._span, LLM_RESPONSE_DURATION)
229252
set_response_attributes(self._span, self._complete_response)
253+
self._netra_output = self._extract_content_text()
230254
self._span.set_status(Status(StatusCode.OK))
231255
self._span.end()
232256

‎netra/instrumentation/google_genai/wrappers.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,8 @@ async def wrapper(wrapped: Callable[..., Any], instance: Any, args: Tuple[Any, .
235235

236236

237237
class StreamingWrapper:
238+
_netra_stream_wrapper = True
239+
238240
def __init__(self, span: Span, response: Iterator[Any]) -> None:
239241
self._span = span
240242
self._buffer: dict[Any, Any] = {"chunk": None, "content": ""}
@@ -272,11 +274,14 @@ def _process_chunk(self, chunk: Any) -> None:
272274
def _finalize_span(self) -> None:
273275
record_span_timing(self._span, LLM_RESPONSE_DURATION)
274276
set_response_attributes(self._span, self._buffer)
277+
self._netra_output = self._buffer.get("content", "") if isinstance(self._buffer, dict) else ""
275278
self._span.set_status(Status(StatusCode.OK))
276279
self._span.end()
277280

278281

279282
class AsyncStreamingWrapper:
283+
_netra_stream_wrapper = True
284+
280285
def __init__(self, span: Span, response: AsyncIterator[Any]) -> None:
281286
self._span = span
282287
self._buffer: dict[Any, Any] = {"chunk": None, "content": ""}
@@ -313,5 +318,6 @@ def _process_chunk(self, chunk: Any) -> None:
313318
def _finalize_span(self) -> None:
314319
record_span_timing(self._span, LLM_RESPONSE_DURATION)
315320
set_response_attributes(self._span, self._buffer)
321+
self._netra_output = self._buffer.get("content", "") if isinstance(self._buffer, dict) else ""
316322
self._span.set_status(Status(StatusCode.OK))
317323
self._span.end()

‎netra/instrumentation/groq/wrappers.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@
2626
class StreamingWrapper(ObjectProxy): # type: ignore[misc]
2727
"""Wrapper for streaming responses (OpenAI-style)."""
2828

29+
_netra_stream_wrapper = True
30+
2931
def __init__(self, span: Span, response: Iterator[Any], request_kwargs: Dict[str, Any]) -> None:
3032
super().__init__(response)
3133
self._span = span
@@ -43,6 +45,15 @@ def _ensure_choice(self, index: int) -> None:
4345
else:
4446
self._complete_response["choices"].append({"text": ""})
4547

48+
def _extract_content_text(self) -> str:
49+
"""Extract the plain text content from the accumulated response."""
50+
parts = []
51+
for choice in self._complete_response.get("choices", []):
52+
msg = choice.get("message", {})
53+
if content := msg.get("content"):
54+
parts.append(content)
55+
return "".join(parts)
56+
4657
def __iter__(self) -> Iterator[Any]:
4758
return self
4859

@@ -98,13 +109,16 @@ def _process_chunk(self, chunk: Any) -> None:
98109
def _finalize_span(self) -> None:
99110
record_span_timing(self._span, LLM_RESPONSE_DURATION)
100111
set_response_attributes(self._span, self._complete_response)
112+
self._netra_output = self._extract_content_text()
101113
self._span.set_status(Status(StatusCode.OK))
102114
self._span.end()
103115

104116

105117
class AsyncStreamingWrapper(ObjectProxy): # type: ignore[misc]
106118
"""Async wrapper for streaming responses (OpenAI-style)."""
107119

120+
_netra_stream_wrapper = True
121+
108122
def __init__(self, span: Span, response: AsyncIterator[Any], request_kwargs: Dict[str, Any]) -> None:
109123
super().__init__(response)
110124
self._span = span
@@ -122,6 +136,15 @@ def _ensure_choice(self, index: int) -> None:
122136
else:
123137
self._complete_response["choices"].append({"text": ""})
124138

139+
def _extract_content_text(self) -> str:
140+
"""Extract the plain text content from the accumulated response."""
141+
parts = []
142+
for choice in self._complete_response.get("choices", []):
143+
msg = choice.get("message", {})
144+
if content := msg.get("content"):
145+
parts.append(content)
146+
return "".join(parts)
147+
125148
def __aiter__(self) -> AsyncIterator[Any]:
126149
return self
127150

@@ -177,6 +200,7 @@ def _process_chunk(self, chunk: Any) -> None:
177200
def _finalize_span(self) -> None:
178201
record_span_timing(self._span, LLM_RESPONSE_DURATION)
179202
set_response_attributes(self._span, self._complete_response)
203+
self._netra_output = self._extract_content_text()
180204
self._span.set_status(Status(StatusCode.OK))
181205
self._span.end()
182206

‎netra/instrumentation/litellm/wrappers.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -335,6 +335,8 @@ async def wrapper(
335335
class StreamingWrapper(ObjectProxy): # type: ignore[misc]
336336
"""Wrapper for streaming responses"""
337337

338+
_netra_stream_wrapper = True
339+
338340
def __init__(self, span: Span, response: Iterator[Any], request_kwargs: Dict[str, Any]) -> None:
339341
super().__init__(response)
340342
self._span = span
@@ -354,6 +356,15 @@ def _ensure_choice(self, index: int) -> None:
354356
else:
355357
self._complete_response["choices"].append({"text": ""})
356358

359+
def _extract_content_text(self) -> str:
360+
"""Extract the plain text content from the accumulated response."""
361+
parts = []
362+
for choice in self._complete_response.get("choices", []):
363+
msg = choice.get("message", {})
364+
if content := msg.get("content"):
365+
parts.append(content)
366+
return "".join(parts)
367+
357368
def __enter__(self) -> "StreamingWrapper":
358369
if hasattr(self.__wrapped__, "__enter__"):
359370
self.__wrapped__.__enter__()
@@ -444,13 +455,16 @@ def _finalize_span(self) -> None:
444455
"""Finalize span when streaming is complete"""
445456
record_span_timing(self._span, LLM_RESPONSE_DURATION)
446457
set_response_attributes(self._span, self._complete_response)
458+
self._netra_output = self._extract_content_text()
447459
self._span.set_status(Status(StatusCode.OK))
448460
self._span.end()
449461

450462

451463
class AsyncStreamingWrapper(ObjectProxy): # type: ignore[misc]
452464
"""Async wrapper for streaming responses"""
453465

466+
_netra_stream_wrapper = True
467+
454468
def __init__(self, span: Span, response: AsyncIterator[Any], request_kwargs: Dict[str, Any]) -> None:
455469
super().__init__(response)
456470
self._span = span
@@ -470,6 +484,15 @@ def _ensure_choice(self, index: int) -> None:
470484
else:
471485
self._complete_response["choices"].append({"text": ""})
472486

487+
def _extract_content_text(self) -> str:
488+
"""Extract the plain text content from the accumulated response."""
489+
parts = []
490+
for choice in self._complete_response.get("choices", []):
491+
msg = choice.get("message", {})
492+
if content := msg.get("content"):
493+
parts.append(content)
494+
return "".join(parts)
495+
473496
async def __aenter__(self) -> "AsyncStreamingWrapper":
474497
if hasattr(self.__wrapped__, "__aenter__"):
475498
await self.__wrapped__.__aenter__()
@@ -560,5 +583,6 @@ def _finalize_span(self) -> None:
560583
"""Finalize span when streaming is complete"""
561584
record_span_timing(self._span, LLM_RESPONSE_DURATION)
562585
set_response_attributes(self._span, self._complete_response)
586+
self._netra_output = self._extract_content_text()
563587
self._span.set_status(Status(StatusCode.OK))
564588
self._span.end()

0 commit comments

Comments
 (0)