netra-sdk 0.1.34__tar.gz → 0.1.36__tar.gz

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.

Potentially problematic release.


This version of netra-sdk might be problematic. Click here for more details.

Files changed (57) hide show
  1. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/PKG-INFO +1 -1
  2. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/__init__.py +13 -0
  3. netra_sdk-0.1.36/netra/decorators.py +441 -0
  4. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/session_manager.py +50 -7
  5. netra_sdk-0.1.36/netra/version.py +1 -0
  6. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/pyproject.toml +1 -1
  7. netra_sdk-0.1.34/netra/decorators.py +0 -218
  8. netra_sdk-0.1.34/netra/version.py +0 -1
  9. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/LICENCE +0 -0
  10. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/README.md +0 -0
  11. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/anonymizer/__init__.py +0 -0
  12. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/anonymizer/anonymizer.py +0 -0
  13. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/anonymizer/base.py +0 -0
  14. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/anonymizer/fp_anonymizer.py +0 -0
  15. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/config.py +0 -0
  16. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/exceptions/__init__.py +0 -0
  17. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/exceptions/injection.py +0 -0
  18. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/exceptions/pii.py +0 -0
  19. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/input_scanner.py +0 -0
  20. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/__init__.py +0 -0
  21. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/aiohttp/__init__.py +0 -0
  22. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/aiohttp/version.py +0 -0
  23. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/cohere/__init__.py +0 -0
  24. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/cohere/version.py +0 -0
  25. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/fastapi/__init__.py +0 -0
  26. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/fastapi/version.py +0 -0
  27. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/google_genai/__init__.py +0 -0
  28. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/google_genai/config.py +0 -0
  29. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/google_genai/utils.py +0 -0
  30. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/google_genai/version.py +0 -0
  31. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/httpx/__init__.py +0 -0
  32. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/httpx/version.py +0 -0
  33. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/instruments.py +0 -0
  34. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/litellm/__init__.py +0 -0
  35. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/litellm/version.py +0 -0
  36. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/litellm/wrappers.py +0 -0
  37. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/mistralai/__init__.py +0 -0
  38. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/mistralai/config.py +0 -0
  39. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/mistralai/utils.py +0 -0
  40. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/mistralai/version.py +0 -0
  41. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/openai/__init__.py +0 -0
  42. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/openai/version.py +0 -0
  43. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/openai/wrappers.py +0 -0
  44. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/pydantic_ai/__init__.py +0 -0
  45. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/pydantic_ai/utils.py +0 -0
  46. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/pydantic_ai/version.py +0 -0
  47. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/pydantic_ai/wrappers.py +0 -0
  48. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/weaviate/__init__.py +0 -0
  49. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/instrumentation/weaviate/version.py +0 -0
  50. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/pii.py +0 -0
  51. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/processors/__init__.py +0 -0
  52. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/processors/instrumentation_span_processor.py +0 -0
  53. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/processors/scrubbing_span_processor.py +0 -0
  54. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/processors/session_span_processor.py +0 -0
  55. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/scanner.py +0 -0
  56. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/span_wrapper.py +0 -0
  57. {netra_sdk-0.1.34 → netra_sdk-0.1.36}/netra/tracer.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: netra-sdk
3
- Version: 0.1.34
3
+ Version: 0.1.36
4
4
  Summary: A Python SDK for AI application observability that provides OpenTelemetry-based monitoring, tracing, and PII protection for LLM and vector database applications. Enables easy instrumentation, session tracking, and privacy-focused data collection for AI systems in production environments.
5
5
  License: Apache-2.0
6
6
  Keywords: netra,tracing,observability,sdk,ai,llm,vector,database
@@ -284,5 +284,18 @@ class Netra:
284
284
  if value:
285
285
  SessionManager.set_attribute_on_target_span(f"{Config.LIBRARY_NAME}.span.output", value, span_name)
286
286
 
287
+ @classmethod
288
+ def set_prompt(cls, value: Any, span_name: Optional[str] = None) -> None:
289
+ """
290
+ Set custom attribute `netra.span.prompt` on a target span.
291
+
292
+ Args:
293
+ value: Prompt payload to record (string or JSON-serializable object)
294
+ span_name: Optional. When provided, sets the attribute on the span registered
295
+ with this name. Otherwise sets on the active span.
296
+ """
297
+ if value:
298
+ SessionManager.set_attribute_on_target_span(f"{Config.LIBRARY_NAME}.span.prompt", value, span_name)
299
+
287
300
 
288
301
  __all__ = ["Netra", "UsageModel", "ActionModel"]
@@ -0,0 +1,441 @@
1
+ """Netra decorator utilities.
2
+
3
+ This module provides decorators for common patterns in Netra SDK.
4
+ Decorators can be applied to both functions and classes.
5
+ """
6
+
7
+ import functools
8
+ import inspect
9
+ import json
10
+ import logging
11
+ from typing import (
12
+ Any,
13
+ AsyncGenerator,
14
+ Awaitable,
15
+ Callable,
16
+ Dict,
17
+ Generator,
18
+ Optional,
19
+ ParamSpec,
20
+ Tuple,
21
+ TypeVar,
22
+ Union,
23
+ cast,
24
+ )
25
+
26
+ try:
27
+ # Optional import: only present in FastAPI/Starlette environments
28
+ from starlette.responses import StreamingResponse
29
+ except Exception: # pragma: no cover - starlette may not be installed in some environments
30
+ StreamingResponse = None
31
+
32
+ from opentelemetry import trace
33
+
34
+ from .config import Config
35
+ from .session_manager import SessionManager
36
+
37
+ logger = logging.getLogger(__name__)
38
+
39
+ P = ParamSpec("P")
40
+ R = TypeVar("R")
41
+
42
+ F_Callable = TypeVar("F_Callable", bound=Callable[..., Any])
43
+ C = TypeVar("C", bound=type)
44
+
45
+
46
+ def _serialize_value(value: Any) -> str:
47
+ """Safely serialize a value to string for span attributes."""
48
+ try:
49
+ if isinstance(value, (str, int, float, bool, type(None))):
50
+ return str(value)
51
+ elif isinstance(value, (list, dict, tuple)):
52
+ return json.dumps(value, default=str)[:1000] # Limit size
53
+ else:
54
+ return str(value)[:1000] # Limit size
55
+ except Exception:
56
+ return str(type(value).__name__)
57
+
58
+
59
+ def _add_span_attributes(
60
+ span: trace.Span, func: Callable[..., Any], args: Tuple[Any, ...], kwargs: Dict[str, Any], entity_type: str
61
+ ) -> None:
62
+ """Helper function to add span attributes from function parameters."""
63
+ span.set_attribute(f"{Config.LIBRARY_NAME}.entity.type", entity_type)
64
+
65
+ try:
66
+ sig = inspect.signature(func)
67
+ param_names = list(sig.parameters.keys())
68
+ input_data = {}
69
+
70
+ for i, arg in enumerate(args):
71
+ if i < len(param_names):
72
+ param_name = param_names[i]
73
+ if param_name not in ("self", "cls"):
74
+ input_data[param_name] = _serialize_value(arg)
75
+
76
+ for key, value in kwargs.items():
77
+ input_data[key] = _serialize_value(value)
78
+
79
+ if input_data:
80
+ span.set_attribute(f"{Config.LIBRARY_NAME}.entity.input", json.dumps(input_data))
81
+
82
+ except Exception as e:
83
+ span.set_attribute(f"{Config.LIBRARY_NAME}.input_error", str(e))
84
+
85
+
86
+ def _add_output_attributes(span: trace.Span, result: Any) -> None:
87
+ """Helper function to add output attributes to span."""
88
+ try:
89
+ serialized_output = _serialize_value(result)
90
+ span.set_attribute(f"{Config.LIBRARY_NAME}.entity.output", serialized_output)
91
+ except Exception as e:
92
+ span.set_attribute(f"{Config.LIBRARY_NAME}.entity.output_error", str(e))
93
+
94
+
95
+ def _is_streaming_response(obj: Any) -> bool:
96
+ """Return True if obj is a Starlette StreamingResponse instance."""
97
+ if StreamingResponse is None:
98
+ return False
99
+ try:
100
+ return isinstance(obj, StreamingResponse)
101
+ except Exception:
102
+ return False
103
+
104
+
105
+ def _is_async_generator(obj: Any) -> bool:
106
+ return inspect.isasyncgen(obj)
107
+
108
+
109
+ def _is_sync_generator(obj: Any) -> bool:
110
+ return inspect.isgenerator(obj)
111
+
112
+
113
+ def _wrap_async_generator_with_span(
114
+ span: trace.Span,
115
+ agen: AsyncGenerator[Any, None],
116
+ span_name: str,
117
+ entity_type: str,
118
+ ) -> AsyncGenerator[Any, None]:
119
+ """Wrap an async generator so the span remains current for the full iteration and ends afterwards."""
120
+
121
+ async def _wrapped() -> AsyncGenerator[Any, None]:
122
+ # Activate span for the entire iteration
123
+ with trace.use_span(span, end_on_exit=False):
124
+ try:
125
+ async for item in agen:
126
+ yield item
127
+ except Exception as e:
128
+ try:
129
+ span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
130
+ span.record_exception(e)
131
+ finally:
132
+ span.end()
133
+ # De-register and pop entity at the very end for streaming lifecycle
134
+ try:
135
+ SessionManager.unregister_span(span_name, span)
136
+ except Exception:
137
+ logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
138
+ SessionManager.pop_entity(entity_type)
139
+ raise
140
+ else:
141
+ # Normal completion
142
+ span.end()
143
+ try:
144
+ SessionManager.unregister_span(span_name, span)
145
+ except Exception:
146
+ logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
147
+ SessionManager.pop_entity(entity_type)
148
+
149
+ return _wrapped()
150
+
151
+
152
+ def _wrap_sync_generator_with_span(
153
+ span: trace.Span,
154
+ gen: Generator[Any, None, None],
155
+ span_name: str,
156
+ entity_type: str,
157
+ ) -> Generator[Any, None, None]:
158
+ """Wrap a sync generator so the span remains current for the full iteration and ends afterwards."""
159
+
160
+ def _wrapped() -> Generator[Any, None, None]:
161
+ with trace.use_span(span, end_on_exit=False):
162
+ try:
163
+ for item in gen:
164
+ yield item
165
+ except Exception as e:
166
+ try:
167
+ span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
168
+ span.record_exception(e)
169
+ finally:
170
+ span.end()
171
+ try:
172
+ SessionManager.unregister_span(span_name, span)
173
+ except Exception:
174
+ logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
175
+ SessionManager.pop_entity(entity_type)
176
+ raise
177
+ else:
178
+ span.end()
179
+ try:
180
+ SessionManager.unregister_span(span_name, span)
181
+ except Exception:
182
+ logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
183
+ SessionManager.pop_entity(entity_type)
184
+
185
+ return _wrapped()
186
+
187
+
188
+ def _wrap_streaming_response_with_span(
189
+ span: trace.Span,
190
+ resp: Any,
191
+ span_name: str,
192
+ entity_type: str,
193
+ ) -> Any:
194
+ """Wrap StreamingResponse.body_iterator with a generator that keeps span current and ends it afterwards."""
195
+ try:
196
+ body_iter = getattr(resp, "body_iterator", None)
197
+ if body_iter is None:
198
+ return resp
199
+ # Async iterator
200
+ if inspect.isasyncgen(body_iter) or hasattr(body_iter, "__aiter__"):
201
+
202
+ async def _aiter_wrapper(): # type: ignore[no-untyped-def]
203
+ with trace.use_span(span, end_on_exit=False):
204
+ try:
205
+ async for chunk in body_iter:
206
+ yield chunk
207
+ except Exception as e:
208
+ try:
209
+ span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
210
+ span.record_exception(e)
211
+ finally:
212
+ span.end()
213
+ try:
214
+ SessionManager.unregister_span(span_name, span)
215
+ except Exception:
216
+ logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
217
+ SessionManager.pop_entity(entity_type)
218
+ raise
219
+ else:
220
+ span.end()
221
+ try:
222
+ SessionManager.unregister_span(span_name, span)
223
+ except Exception:
224
+ logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
225
+ SessionManager.pop_entity(entity_type)
226
+
227
+ resp.body_iterator = _aiter_wrapper() # type: ignore[no-untyped-call]
228
+ return resp
229
+
230
+ # Sync iterator
231
+ if inspect.isgenerator(body_iter) or hasattr(body_iter, "__iter__"):
232
+
233
+ def _iter_wrapper(): # type: ignore[no-untyped-def]
234
+ with trace.use_span(span, end_on_exit=False):
235
+ try:
236
+ for chunk in body_iter:
237
+ yield chunk
238
+ except Exception as e:
239
+ try:
240
+ span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
241
+ span.record_exception(e)
242
+ finally:
243
+ span.end()
244
+ try:
245
+ SessionManager.unregister_span(span_name, span)
246
+ except Exception:
247
+ logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
248
+ SessionManager.pop_entity(entity_type)
249
+ raise
250
+ else:
251
+ span.end()
252
+ try:
253
+ SessionManager.unregister_span(span_name, span)
254
+ except Exception:
255
+ logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
256
+ SessionManager.pop_entity(entity_type)
257
+
258
+ resp.body_iterator = _iter_wrapper() # type: ignore[no-untyped-call]
259
+ return resp
260
+ except Exception:
261
+ logger.exception("Failed to wrap StreamingResponse with span '%s'", span_name)
262
+ return resp
263
+
264
+
265
+ def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optional[str] = None) -> Callable[P, R]:
266
+ module_name = func.__name__
267
+ is_async = inspect.iscoroutinefunction(func)
268
+ span_name = name if name is not None else func.__name__
269
+
270
+ if is_async:
271
+
272
+ @functools.wraps(func)
273
+ async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
274
+ # Push entity before span starts so processors can capture it
275
+ SessionManager.push_entity(entity_type, span_name)
276
+
277
+ tracer = trace.get_tracer(module_name)
278
+ span = tracer.start_span(span_name)
279
+ # Register and activate span
280
+ try:
281
+ SessionManager.register_span(span_name, span)
282
+ SessionManager.set_current_span(span)
283
+ except Exception:
284
+ logger.exception("Failed to register span '%s' with SessionManager", span_name)
285
+
286
+ with trace.use_span(span, end_on_exit=False):
287
+ _add_span_attributes(span, func, args, kwargs, entity_type)
288
+ try:
289
+ result = await cast(Awaitable[Any], func(*args, **kwargs))
290
+ except Exception as e:
291
+ span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
292
+ span.record_exception(e)
293
+ span.end()
294
+ try:
295
+ SessionManager.unregister_span(span_name, span)
296
+ except Exception:
297
+ logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
298
+ SessionManager.pop_entity(entity_type)
299
+ raise
300
+
301
+ # If result is streaming, defer span end to when stream completes
302
+ if _is_streaming_response(result):
303
+ return _wrap_streaming_response_with_span(span, result, span_name, entity_type)
304
+ if _is_async_generator(result):
305
+ return _wrap_async_generator_with_span(span, result, span_name, entity_type)
306
+ if _is_sync_generator(result):
307
+ return _wrap_sync_generator_with_span(span, result, span_name, entity_type)
308
+
309
+ # Non-streaming: finalize now
310
+ _add_output_attributes(span, result)
311
+ span.end()
312
+ try:
313
+ SessionManager.unregister_span(span_name, span)
314
+ except Exception:
315
+ logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
316
+ SessionManager.pop_entity(entity_type)
317
+ return result
318
+
319
+ return cast(Callable[P, R], async_wrapper)
320
+
321
+ else:
322
+
323
+ @functools.wraps(func)
324
+ def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
325
+ # Push entity before span starts so processors can capture it
326
+ SessionManager.push_entity(entity_type, span_name)
327
+
328
+ tracer = trace.get_tracer(module_name)
329
+ span = tracer.start_span(span_name)
330
+ # Register and activate span
331
+ try:
332
+ SessionManager.register_span(span_name, span)
333
+ SessionManager.set_current_span(span)
334
+ except Exception:
335
+ logger.exception("Failed to register span '%s' with SessionManager", span_name)
336
+
337
+ with trace.use_span(span, end_on_exit=False):
338
+ _add_span_attributes(span, func, args, kwargs, entity_type)
339
+ try:
340
+ result = func(*args, **kwargs)
341
+ except Exception as e:
342
+ span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
343
+ span.record_exception(e)
344
+ span.end()
345
+ try:
346
+ SessionManager.unregister_span(span_name, span)
347
+ except Exception:
348
+ logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
349
+ SessionManager.pop_entity(entity_type)
350
+ raise
351
+
352
+ # If result is streaming, defer span end to when stream completes
353
+ if _is_streaming_response(result):
354
+ return _wrap_streaming_response_with_span(span, result, span_name, entity_type)
355
+ if _is_async_generator(result):
356
+ return _wrap_async_generator_with_span(span, result, span_name, entity_type) # type: ignore[arg-type]
357
+ if _is_sync_generator(result):
358
+ return _wrap_sync_generator_with_span(span, result, span_name, entity_type) # type: ignore[arg-type]
359
+
360
+ # Non-streaming: finalize now
361
+ _add_output_attributes(span, result)
362
+ span.end()
363
+ try:
364
+ SessionManager.unregister_span(span_name, span)
365
+ except Exception:
366
+ logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
367
+ SessionManager.pop_entity(entity_type)
368
+ return result
369
+
370
+ return cast(Callable[P, R], sync_wrapper)
371
+
372
+
373
+ def _wrap_class_methods(cls: C, entity_type: str, name: Optional[str] = None) -> C:
374
+ class_name = name if name is not None else cls.__name__
375
+ for attr_name in cls.__dict__:
376
+ attr = getattr(cls, attr_name)
377
+ if attr_name.startswith("_"):
378
+ continue
379
+ if callable(attr) and inspect.isfunction(attr):
380
+ method_span_name = f"{class_name}.{attr_name}"
381
+ wrapped_method = _create_function_wrapper(attr, entity_type, method_span_name)
382
+ setattr(cls, attr_name, wrapped_method)
383
+ return cls
384
+
385
+
386
+ def workflow(
387
+ target: Union[Callable[P, R], C, None] = None, *, name: Optional[str] = None
388
+ ) -> Union[Callable[P, R], C, Callable[[Callable[P, R]], Callable[P, R]]]:
389
+ def decorator(obj: Union[Callable[P, R], C]) -> Union[Callable[P, R], C]:
390
+ if inspect.isclass(obj):
391
+ return _wrap_class_methods(cast(C, obj), "workflow", name)
392
+ else:
393
+ return _create_function_wrapper(cast(Callable[P, R], obj), "workflow", name)
394
+
395
+ if target is not None:
396
+ return decorator(target)
397
+ return decorator
398
+
399
+
400
+ def agent(
401
+ target: Union[Callable[P, R], C, None] = None, *, name: Optional[str] = None
402
+ ) -> Union[Callable[P, R], C, Callable[[Callable[P, R]], Callable[P, R]]]:
403
+ def decorator(obj: Union[Callable[P, R], C]) -> Union[Callable[P, R], C]:
404
+ if inspect.isclass(obj):
405
+ return _wrap_class_methods(cast(C, obj), "agent", name)
406
+ else:
407
+ return _create_function_wrapper(cast(Callable[P, R], obj), "agent", name)
408
+
409
+ if target is not None:
410
+ return decorator(target)
411
+ return decorator
412
+
413
+
414
+ def task(
415
+ target: Union[Callable[P, R], C, None] = None, *, name: Optional[str] = None
416
+ ) -> Union[Callable[P, R], C, Callable[[Callable[P, R]], Callable[P, R]]]:
417
+ def decorator(obj: Union[Callable[P, R], C]) -> Union[Callable[P, R], C]:
418
+ if inspect.isclass(obj):
419
+ return _wrap_class_methods(cast(C, obj), "task", name)
420
+ else:
421
+ # When obj is a function, it should be type Callable[P, R]
422
+ return _create_function_wrapper(cast(Callable[P, R], obj), "task", name)
423
+
424
+ if target is not None:
425
+ return decorator(target)
426
+ return decorator
427
+
428
+
429
+ def span(
430
+ target: Union[Callable[P, R], C, None] = None, *, name: Optional[str] = None
431
+ ) -> Union[Callable[P, R], C, Callable[[Callable[P, R]], Callable[P, R]]]:
432
+ def decorator(obj: Union[Callable[P, R], C]) -> Union[Callable[P, R], C]:
433
+ if inspect.isclass(obj):
434
+ return _wrap_class_methods(cast(C, obj), "span", name)
435
+ else:
436
+ # When obj is a function, it should be type Callable[P, R]
437
+ return _create_function_wrapper(cast(Callable[P, R], obj), "span", name)
438
+
439
+ if target is not None:
440
+ return decorator(target)
441
+ return decorator
@@ -31,6 +31,10 @@ class SessionManager:
31
31
  # Span registry: name -> stack of spans (most-recent last)
32
32
  _spans_by_name: Dict[str, List[trace.Span]] = {}
33
33
 
34
+ # Global stack of active spans in creation order (oldest first, newest last)
35
+ # Maintained for spans registered via SessionManager (e.g., SpanWrapper)
36
+ _active_spans: List[trace.Span] = []
37
+
34
38
  @classmethod
35
39
  def set_current_span(cls, span: Optional[trace.Span]) -> None:
36
40
  """
@@ -62,6 +66,8 @@ class SessionManager:
62
66
  cls._spans_by_name[name] = [span]
63
67
  else:
64
68
  stack.append(span)
69
+ # Track globally as active
70
+ cls._active_spans.append(span)
65
71
  except Exception:
66
72
  logger.exception("Failed to register span '%s'", name)
67
73
 
@@ -81,6 +87,11 @@ class SessionManager:
81
87
  break
82
88
  if not stack:
83
89
  cls._spans_by_name.pop(name, None)
90
+ # Also remove from global active list (remove last matching instance)
91
+ for i in range(len(cls._active_spans) - 1, -1, -1):
92
+ if cls._active_spans[i] is span:
93
+ cls._active_spans.pop(i)
94
+ break
84
95
  except Exception:
85
96
  logger.exception("Failed to unregister span '%s'", name)
86
97
 
@@ -240,12 +251,13 @@ class SessionManager:
240
251
  @classmethod
241
252
  def set_attribute_on_target_span(cls, attr_key: str, attr_value: Any, span_name: Optional[str] = None) -> None:
242
253
  """
243
- Best-effort setter to annotate the active span with the provided attribute.
254
+ Best-effort setter to annotate a target span with the provided attribute.
244
255
 
245
- If span_name is provided, we look up that span via SessionManager's registry and set
246
- the attribute on that span explicitly. Otherwise, we annotate the active span.
247
- We first try the OpenTelemetry current span; if that's invalid, we fall back to
248
- the SDK-managed current span from `SessionManager`.
256
+ Behavior:
257
+ - If span_name is provided, set the attribute on the span registered with that name.
258
+ - If no span_name is provided, attempt to set the attribute on the SDK root span
259
+ (created when Netra.init(enable_root_span=True)). If the root span is unavailable,
260
+ fall back to the currently active span (OTel current span or SDK-managed current span).
249
261
  """
250
262
  try:
251
263
  # Convert attribute value to a JSON-safe string representation
@@ -268,10 +280,41 @@ class SessionManager:
268
280
  target.set_attribute(attr_key, attr_str)
269
281
  return
270
282
 
271
- # Otherwise annotate the active span
283
+ # Otherwise, attempt to set on the root-most span in the current trace
284
+ candidate = None
285
+
286
+ # Determine current trace_id from the active/current span
272
287
  current_span = trace.get_current_span()
273
288
  has_valid_current = getattr(current_span, "is_recording", None) is not None and current_span.is_recording()
274
- candidate = current_span if has_valid_current else cls.get_current_span()
289
+ base_span = current_span if has_valid_current else cls.get_current_span()
290
+ trace_id: Optional[int] = None
291
+ try:
292
+ if base_span is not None and hasattr(base_span, "get_span_context"):
293
+ sc = base_span.get_span_context()
294
+ trace_id = getattr(sc, "trace_id", None)
295
+ except Exception:
296
+ trace_id = None
297
+
298
+ # Find the earliest active span in this process that belongs to the same trace
299
+ if trace_id is not None:
300
+ try:
301
+ for s in cls._active_spans:
302
+ if s is None:
303
+ continue
304
+ if not getattr(s, "is_recording", lambda: False)():
305
+ continue
306
+ sc = getattr(s, "get_span_context", lambda: None)()
307
+ if sc is None:
308
+ continue
309
+ if getattr(sc, "trace_id", None) == trace_id:
310
+ candidate = s
311
+ break
312
+ except Exception:
313
+ candidate = None
314
+
315
+ # Fallback to the current active span if no root-most could be found
316
+ if candidate is None:
317
+ candidate = base_span
275
318
  if candidate is None:
276
319
  logger.debug("No active span found to set attribute %s", attr_key)
277
320
  return
@@ -0,0 +1 @@
1
+ __version__ = "0.1.36"
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
4
4
 
5
5
  [project]
6
6
  name = "netra-sdk"
7
- version = "0.1.34"
7
+ version = "0.1.36"
8
8
  description = "A Python SDK for AI application observability that provides OpenTelemetry-based monitoring, tracing, and PII protection for LLM and vector database applications. Enables easy instrumentation, session tracking, and privacy-focused data collection for AI systems in production environments."
9
9
  authors = [
10
10
  {name = "Sooraj Thomas",email = "sooraj@keyvalue.systems"}
@@ -1,218 +0,0 @@
1
- """Netra decorator utilities.
2
-
3
- This module provides decorators for common patterns in Netra SDK.
4
- Decorators can be applied to both functions and classes.
5
- """
6
-
7
- import functools
8
- import inspect
9
- import json
10
- import logging
11
- from typing import Any, Awaitable, Callable, Dict, Optional, ParamSpec, Tuple, TypeVar, Union, cast
12
-
13
- from opentelemetry import trace
14
-
15
- from .config import Config
16
- from .session_manager import SessionManager
17
-
18
- logger = logging.getLogger(__name__)
19
-
20
- P = ParamSpec("P")
21
- R = TypeVar("R")
22
-
23
- F_Callable = TypeVar("F_Callable", bound=Callable[..., Any])
24
- C = TypeVar("C", bound=type)
25
-
26
-
27
- def _serialize_value(value: Any) -> str:
28
- """Safely serialize a value to string for span attributes."""
29
- try:
30
- if isinstance(value, (str, int, float, bool, type(None))):
31
- return str(value)
32
- elif isinstance(value, (list, dict, tuple)):
33
- return json.dumps(value, default=str)[:1000] # Limit size
34
- else:
35
- return str(value)[:1000] # Limit size
36
- except Exception:
37
- return str(type(value).__name__)
38
-
39
-
40
- def _add_span_attributes(
41
- span: trace.Span, func: Callable[..., Any], args: Tuple[Any, ...], kwargs: Dict[str, Any], entity_type: str
42
- ) -> None:
43
- """Helper function to add span attributes from function parameters."""
44
- span.set_attribute(f"{Config.LIBRARY_NAME}.entity.type", entity_type)
45
-
46
- try:
47
- sig = inspect.signature(func)
48
- param_names = list(sig.parameters.keys())
49
- input_data = {}
50
-
51
- for i, arg in enumerate(args):
52
- if i < len(param_names):
53
- param_name = param_names[i]
54
- if param_name not in ("self", "cls"):
55
- input_data[param_name] = _serialize_value(arg)
56
-
57
- for key, value in kwargs.items():
58
- input_data[key] = _serialize_value(value)
59
-
60
- if input_data:
61
- span.set_attribute(f"{Config.LIBRARY_NAME}.entity.input", json.dumps(input_data))
62
-
63
- except Exception as e:
64
- span.set_attribute(f"{Config.LIBRARY_NAME}.input_error", str(e))
65
-
66
-
67
- def _add_output_attributes(span: trace.Span, result: Any) -> None:
68
- """Helper function to add output attributes to span."""
69
- try:
70
- serialized_output = _serialize_value(result)
71
- span.set_attribute(f"{Config.LIBRARY_NAME}.entity.output", serialized_output)
72
- except Exception as e:
73
- span.set_attribute(f"{Config.LIBRARY_NAME}.entity.output_error", str(e))
74
-
75
-
76
- def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optional[str] = None) -> Callable[P, R]:
77
- module_name = func.__name__
78
- is_async = inspect.iscoroutinefunction(func)
79
- span_name = name if name is not None else func.__name__
80
-
81
- if is_async:
82
-
83
- @functools.wraps(func)
84
- async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
85
- # Push entity to stack before span starts so SessionSpanProcessor can capture it
86
- SessionManager.push_entity(entity_type, span_name)
87
-
88
- tracer = trace.get_tracer(module_name)
89
- with tracer.start_as_current_span(span_name) as span:
90
- # Register the span by name for cross-context attribute setting
91
- try:
92
- SessionManager.register_span(span_name, span)
93
- SessionManager.set_current_span(span)
94
- except Exception:
95
- logger.exception("Failed to register span '%s' with SessionManager", span_name)
96
-
97
- _add_span_attributes(span, func, args, kwargs, entity_type)
98
- try:
99
- result = await cast(Awaitable[Any], func(*args, **kwargs))
100
- _add_output_attributes(span, result)
101
- return result
102
- except Exception as e:
103
- span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
104
- raise
105
- finally:
106
- # Unregister and pop entity from stack after function call is done
107
- try:
108
- SessionManager.unregister_span(span_name, span)
109
- except Exception:
110
- logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
111
- SessionManager.pop_entity(entity_type)
112
-
113
- return cast(Callable[P, R], async_wrapper)
114
-
115
- else:
116
-
117
- @functools.wraps(func)
118
- def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
119
- # Push entity to stack before span starts so SessionSpanProcessor can capture it
120
- SessionManager.push_entity(entity_type, span_name)
121
-
122
- tracer = trace.get_tracer(module_name)
123
- with tracer.start_as_current_span(span_name) as span:
124
- # Register the span by name for cross-context attribute setting
125
- try:
126
- SessionManager.register_span(span_name, span)
127
- SessionManager.set_current_span(span)
128
- except Exception:
129
- logger.exception("Failed to register span '%s' with SessionManager", span_name)
130
-
131
- _add_span_attributes(span, func, args, kwargs, entity_type)
132
- try:
133
- result = func(*args, **kwargs)
134
- _add_output_attributes(span, result)
135
- return result
136
- except Exception as e:
137
- span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
138
- raise
139
- finally:
140
- # Unregister and pop entity from stack after function call is done
141
- try:
142
- SessionManager.unregister_span(span_name, span)
143
- except Exception:
144
- logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
145
- SessionManager.pop_entity(entity_type)
146
-
147
- return cast(Callable[P, R], sync_wrapper)
148
-
149
-
150
- def _wrap_class_methods(cls: C, entity_type: str, name: Optional[str] = None) -> C:
151
- class_name = name if name is not None else cls.__name__
152
- for attr_name in cls.__dict__:
153
- attr = getattr(cls, attr_name)
154
- if attr_name.startswith("_"):
155
- continue
156
- if callable(attr) and inspect.isfunction(attr):
157
- method_span_name = f"{class_name}.{attr_name}"
158
- wrapped_method = _create_function_wrapper(attr, entity_type, method_span_name)
159
- setattr(cls, attr_name, wrapped_method)
160
- return cls
161
-
162
-
163
- def workflow(
164
- target: Union[Callable[P, R], C, None] = None, *, name: Optional[str] = None
165
- ) -> Union[Callable[P, R], C, Callable[[Callable[P, R]], Callable[P, R]]]:
166
- def decorator(obj: Union[Callable[P, R], C]) -> Union[Callable[P, R], C]:
167
- if inspect.isclass(obj):
168
- return _wrap_class_methods(cast(C, obj), "workflow", name)
169
- else:
170
- return _create_function_wrapper(cast(Callable[P, R], obj), "workflow", name)
171
-
172
- if target is not None:
173
- return decorator(target)
174
- return decorator
175
-
176
-
177
- def agent(
178
- target: Union[Callable[P, R], C, None] = None, *, name: Optional[str] = None
179
- ) -> Union[Callable[P, R], C, Callable[[Callable[P, R]], Callable[P, R]]]:
180
- def decorator(obj: Union[Callable[P, R], C]) -> Union[Callable[P, R], C]:
181
- if inspect.isclass(obj):
182
- return _wrap_class_methods(cast(C, obj), "agent", name)
183
- else:
184
- return _create_function_wrapper(cast(Callable[P, R], obj), "agent", name)
185
-
186
- if target is not None:
187
- return decorator(target)
188
- return decorator
189
-
190
-
191
- def task(
192
- target: Union[Callable[P, R], C, None] = None, *, name: Optional[str] = None
193
- ) -> Union[Callable[P, R], C, Callable[[Callable[P, R]], Callable[P, R]]]:
194
- def decorator(obj: Union[Callable[P, R], C]) -> Union[Callable[P, R], C]:
195
- if inspect.isclass(obj):
196
- return _wrap_class_methods(cast(C, obj), "task", name)
197
- else:
198
- # When obj is a function, it should be type Callable[P, R]
199
- return _create_function_wrapper(cast(Callable[P, R], obj), "task", name)
200
-
201
- if target is not None:
202
- return decorator(target)
203
- return decorator
204
-
205
-
206
- def span(
207
- target: Union[Callable[P, R], C, None] = None, *, name: Optional[str] = None
208
- ) -> Union[Callable[P, R], C, Callable[[Callable[P, R]], Callable[P, R]]]:
209
- def decorator(obj: Union[Callable[P, R], C]) -> Union[Callable[P, R], C]:
210
- if inspect.isclass(obj):
211
- return _wrap_class_methods(cast(C, obj), "span", name)
212
- else:
213
- # When obj is a function, it should be type Callable[P, R]
214
- return _create_function_wrapper(cast(Callable[P, R], obj), "span", name)
215
-
216
- if target is not None:
217
- return decorator(target)
218
- return decorator
@@ -1 +0,0 @@
1
- __version__ = "0.1.34"
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes