netra-sdk 0.1.3__py3-none-any.whl → 0.1.4__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.
Potentially problematic release.
This version of netra-sdk might be problematic. Click here for more details.
- netra/instrumentation/__init__.py +5 -16
- netra/processors/__init__.py +2 -2
- netra/processors/error_detection_processor.py +66 -0
- netra/tracer.py +2 -2
- netra/version.py +1 -1
- {netra_sdk-0.1.3.dist-info → netra_sdk-0.1.4.dist-info}/METADATA +1 -1
- {netra_sdk-0.1.3.dist-info → netra_sdk-0.1.4.dist-info}/RECORD +9 -9
- netra/processors/span_aggregation_processor.py +0 -365
- {netra_sdk-0.1.3.dist-info → netra_sdk-0.1.4.dist-info}/LICENCE +0 -0
- {netra_sdk-0.1.3.dist-info → netra_sdk-0.1.4.dist-info}/WHEEL +0 -0
|
@@ -281,23 +281,12 @@ def init_fastapi_instrumentation() -> bool:
|
|
|
281
281
|
bool: True if initialization was successful, False otherwise.
|
|
282
282
|
"""
|
|
283
283
|
try:
|
|
284
|
-
if
|
|
285
|
-
|
|
286
|
-
from fastapi import FastAPI
|
|
284
|
+
if is_package_installed("fastapi"):
|
|
285
|
+
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
|
|
287
286
|
|
|
288
|
-
|
|
289
|
-
|
|
290
|
-
|
|
291
|
-
original_init(self, *args, **kwargs)
|
|
292
|
-
|
|
293
|
-
try:
|
|
294
|
-
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
|
|
295
|
-
|
|
296
|
-
FastAPIInstrumentor().instrument_app(self)
|
|
297
|
-
except Exception as e:
|
|
298
|
-
logging.warning(f"Failed to auto-instrument FastAPI: {e}")
|
|
299
|
-
|
|
300
|
-
FastAPI.__init__ = _patched_init
|
|
287
|
+
instrumentor = FastAPIInstrumentor()
|
|
288
|
+
if not instrumentor.is_instrumented_by_opentelemetry:
|
|
289
|
+
instrumentor.instrument()
|
|
301
290
|
return True
|
|
302
291
|
except Exception as e:
|
|
303
292
|
logging.error(f"Error initializing FastAPI instrumentor: {e}")
|
netra/processors/__init__.py
CHANGED
|
@@ -1,4 +1,4 @@
|
|
|
1
|
+
from netra.processors.error_detection_processor import ErrorDetectionProcessor
|
|
1
2
|
from netra.processors.session_span_processor import SessionSpanProcessor
|
|
2
|
-
from netra.processors.span_aggregation_processor import SpanAggregationProcessor
|
|
3
3
|
|
|
4
|
-
__all__ = ["
|
|
4
|
+
__all__ = ["ErrorDetectionProcessor", "SessionSpanProcessor"]
|
|
@@ -0,0 +1,66 @@
|
|
|
1
|
+
import logging
|
|
2
|
+
from typing import Any, Optional
|
|
3
|
+
|
|
4
|
+
import httpx
|
|
5
|
+
from opentelemetry.sdk.trace import SpanProcessor
|
|
6
|
+
from opentelemetry.trace import Context, Span
|
|
7
|
+
|
|
8
|
+
from netra import Netra
|
|
9
|
+
|
|
10
|
+
logger = logging.getLogger(__name__)
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class ErrorDetectionProcessor(SpanProcessor): # type: ignore[misc]
|
|
14
|
+
"""
|
|
15
|
+
OpenTelemetry span processor that monitors for error attributes in spans and creates custom events.
|
|
16
|
+
"""
|
|
17
|
+
|
|
18
|
+
def __init__(self) -> None:
|
|
19
|
+
pass
|
|
20
|
+
|
|
21
|
+
def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
|
|
22
|
+
"""Called when a span starts."""
|
|
23
|
+
span_id = self._get_span_id(span)
|
|
24
|
+
if not span_id:
|
|
25
|
+
return
|
|
26
|
+
|
|
27
|
+
# Wrap span methods to capture data
|
|
28
|
+
self._wrap_span_methods(span, span_id)
|
|
29
|
+
|
|
30
|
+
def on_end(self, span: Span) -> None:
|
|
31
|
+
"""Called when a span ends."""
|
|
32
|
+
|
|
33
|
+
def force_flush(self, timeout_millis: int = 30000) -> bool:
|
|
34
|
+
"""Force flush any pending data."""
|
|
35
|
+
return True
|
|
36
|
+
|
|
37
|
+
def shutdown(self) -> bool:
|
|
38
|
+
"""Shutdown the processor."""
|
|
39
|
+
return True
|
|
40
|
+
|
|
41
|
+
def _get_span_id(self, span: Span) -> Optional[str]:
|
|
42
|
+
"""Get a unique identifier for the span."""
|
|
43
|
+
try:
|
|
44
|
+
span_context = span.get_span_context()
|
|
45
|
+
return f"{span_context.trace_id:032x}-{span_context.span_id:016x}"
|
|
46
|
+
except Exception:
|
|
47
|
+
return None
|
|
48
|
+
|
|
49
|
+
def _status_code_processing(self, status_code: int) -> None:
|
|
50
|
+
if httpx.codes.is_error(status_code):
|
|
51
|
+
event_attributes = {"has_error": True, "status_code": status_code}
|
|
52
|
+
Netra.set_custom_event(event_name="error_detected", attributes=event_attributes)
|
|
53
|
+
|
|
54
|
+
def _wrap_span_methods(self, span: Span, span_id: str) -> Any:
|
|
55
|
+
"""Wrap span methods to capture attributes and events."""
|
|
56
|
+
# Wrap set_attribute
|
|
57
|
+
original_set_attribute = span.set_attribute
|
|
58
|
+
|
|
59
|
+
def wrapped_set_attribute(key: str, value: Any) -> Any:
|
|
60
|
+
# Status code processing
|
|
61
|
+
if key == "http.status_code":
|
|
62
|
+
self._status_code_processing(value)
|
|
63
|
+
|
|
64
|
+
return original_set_attribute(key, value)
|
|
65
|
+
|
|
66
|
+
span.set_attribute = wrapped_set_attribute
|
netra/tracer.py
CHANGED
|
@@ -66,10 +66,10 @@ class Tracer:
|
|
|
66
66
|
headers=self.cfg.headers,
|
|
67
67
|
)
|
|
68
68
|
# Add span processors for session span processing and data aggregation processing
|
|
69
|
-
from netra.processors import
|
|
69
|
+
from netra.processors import ErrorDetectionProcessor, SessionSpanProcessor
|
|
70
70
|
|
|
71
71
|
provider.add_span_processor(SessionSpanProcessor())
|
|
72
|
-
provider.add_span_processor(
|
|
72
|
+
provider.add_span_processor(ErrorDetectionProcessor())
|
|
73
73
|
|
|
74
74
|
# Install appropriate span processor
|
|
75
75
|
if self.cfg.disable_batch:
|
netra/version.py
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
__version__ = "0.1.
|
|
1
|
+
__version__ = "0.1.4"
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.3
|
|
2
2
|
Name: netra-sdk
|
|
3
|
-
Version: 0.1.
|
|
3
|
+
Version: 0.1.4
|
|
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
|
|
@@ -9,7 +9,7 @@ netra/exceptions/__init__.py,sha256=uDgcBxmC4WhdS7HRYQk_TtJyxH1s1o6wZmcsnSHLAcM,
|
|
|
9
9
|
netra/exceptions/injection.py,sha256=ke4eUXRYUFJkMZgdSyPPkPt5PdxToTI6xLEBI0hTWUQ,1332
|
|
10
10
|
netra/exceptions/pii.py,sha256=MT4p_x-zH3VtYudTSxw1Z9qQZADJDspq64WrYqSWlZc,2438
|
|
11
11
|
netra/input_scanner.py,sha256=bzP3s7YudGHQrIbUgQGrcIBEJ6CmOewzuYNSu75cVXM,4988
|
|
12
|
-
netra/instrumentation/__init__.py,sha256=
|
|
12
|
+
netra/instrumentation/__init__.py,sha256=s-sXykQZ4CKUHLqHRR7buOrkN9hXGTZpNALRZkdIHB0,38757
|
|
13
13
|
netra/instrumentation/aiohttp/__init__.py,sha256=M1kuF0R3gKY5rlbhEC1AR13UWHelmfokluL2yFysKWc,14398
|
|
14
14
|
netra/instrumentation/aiohttp/version.py,sha256=Zy-0Aukx-HS_Mo3NKPWg-hlUoWKDzS0w58gLoVtJec8,24
|
|
15
15
|
netra/instrumentation/cohere/__init__.py,sha256=3XwmCAZwZiMkHdNN3YvcBOLsNCx80ymbU31TyMzv1IY,17685
|
|
@@ -28,15 +28,15 @@ netra/instrumentation/mistralai/version.py,sha256=d6593s-XBNvVxri9lr2qLUDZQ3Zk3-
|
|
|
28
28
|
netra/instrumentation/weaviate/__init__.py,sha256=EOlpWxobOLHYKqo_kMct_7nu26x1hr8qkeG5_h99wtg,4330
|
|
29
29
|
netra/instrumentation/weaviate/version.py,sha256=PiCZHjonujPbnIn0KmD3Yl68hrjPRG_oKe5vJF3mmG8,24
|
|
30
30
|
netra/pii.py,sha256=S7GnVzoNJEzKiUWnqN9bOCKPeNLsriztgB2E6Rx-yJU,27023
|
|
31
|
-
netra/processors/__init__.py,sha256=
|
|
31
|
+
netra/processors/__init__.py,sha256=G16VumYTpgV4jsWrKNFSgm6xMQAsZ2Rrux25UVeo5YQ,215
|
|
32
|
+
netra/processors/error_detection_processor.py,sha256=2E0X5tVOfk0fu21dkPvQhIZ8MWPZNpxaV0FtX63njCY,2169
|
|
32
33
|
netra/processors/session_span_processor.py,sha256=uv-KL5lwil3C3wQGdYWiYQMHLBsXrt8hTy_ql6kUWXE,2171
|
|
33
|
-
netra/processors/span_aggregation_processor.py,sha256=eOBAHTi5pgBCrnNthWGby0IY6b7iSUPAIFkDDmwLKOY,15579
|
|
34
34
|
netra/scanner.py,sha256=wqjMZnEbVvrGMiUSI352grUyHpkk94oBfHfMiXPhpGU,3866
|
|
35
35
|
netra/session.py,sha256=Lsd7yps2YtjN7P10HsGN3bRZxufoeUnGxAH6Us8_l-Y,8734
|
|
36
36
|
netra/session_manager.py,sha256=UusP3MRZlLeU4NtBVlXQ_sCgRg-LGleVdYPq5MwLvi8,3555
|
|
37
|
-
netra/tracer.py,sha256=
|
|
38
|
-
netra/version.py,sha256=
|
|
39
|
-
netra_sdk-0.1.
|
|
40
|
-
netra_sdk-0.1.
|
|
41
|
-
netra_sdk-0.1.
|
|
42
|
-
netra_sdk-0.1.
|
|
37
|
+
netra/tracer.py,sha256=9jAKdIHXbaZ6WV_p8I1syQiMdqXVCXMhpEhCBsbbci8,3538
|
|
38
|
+
netra/version.py,sha256=Wzf5T3NBDfhQoTnhnRNHSlAsE0XMqbclXG-M81Vas70,22
|
|
39
|
+
netra_sdk-0.1.4.dist-info/LICENCE,sha256=8B_UoZ-BAl0AqiHAHUETCgd3I2B9yYJ1WEQtVb_qFMA,11359
|
|
40
|
+
netra_sdk-0.1.4.dist-info/METADATA,sha256=jn4RUWgJxhgdMdfmperQ903uiabwOZdQsAbG8JVIU3g,24349
|
|
41
|
+
netra_sdk-0.1.4.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
|
|
42
|
+
netra_sdk-0.1.4.dist-info/RECORD,,
|
|
@@ -1,365 +0,0 @@
|
|
|
1
|
-
"""
|
|
2
|
-
Span aggregation utilities for Combat SDK.
|
|
3
|
-
Handles aggregation of child span data into parent spans.
|
|
4
|
-
"""
|
|
5
|
-
|
|
6
|
-
import json
|
|
7
|
-
import logging
|
|
8
|
-
from collections import defaultdict
|
|
9
|
-
from typing import Any, Dict, Optional, Set
|
|
10
|
-
|
|
11
|
-
import httpx
|
|
12
|
-
from opentelemetry import trace
|
|
13
|
-
from opentelemetry.sdk.trace import SpanProcessor
|
|
14
|
-
from opentelemetry.trace import Context, Span
|
|
15
|
-
|
|
16
|
-
from netra import Netra
|
|
17
|
-
from netra.config import Config
|
|
18
|
-
|
|
19
|
-
logger = logging.getLogger(__name__)
|
|
20
|
-
|
|
21
|
-
|
|
22
|
-
class SpanAggregationData:
|
|
23
|
-
"""Holds aggregated data for a span."""
|
|
24
|
-
|
|
25
|
-
def __init__(self) -> None:
|
|
26
|
-
self.tokens: Dict[str, Dict[str, int]] = defaultdict(lambda: defaultdict(int))
|
|
27
|
-
self.models: Set[str] = set()
|
|
28
|
-
self.has_pii: bool = False
|
|
29
|
-
self.pii_entities: Set[str] = set()
|
|
30
|
-
self.pii_actions: Dict[str, Set[str]] = defaultdict(set)
|
|
31
|
-
self.has_violation: bool = False
|
|
32
|
-
self.violations: Set[str] = set()
|
|
33
|
-
self.violation_actions: Dict[str, Set[str]] = defaultdict(set)
|
|
34
|
-
self.has_error: bool = False
|
|
35
|
-
self.status_codes: Set[int] = set()
|
|
36
|
-
|
|
37
|
-
def merge_from_other(self, other: "SpanAggregationData") -> None:
|
|
38
|
-
"""Merge data from another SpanAggregationData instance."""
|
|
39
|
-
# Merge error data
|
|
40
|
-
if other.has_error:
|
|
41
|
-
self.has_error = True
|
|
42
|
-
self.status_codes.update(other.status_codes)
|
|
43
|
-
|
|
44
|
-
# Merge tokens - take the maximum values for each model
|
|
45
|
-
for model, token_data in other.tokens.items():
|
|
46
|
-
if model not in self.tokens:
|
|
47
|
-
self.tokens[model] = {}
|
|
48
|
-
for token_type, count in token_data.items():
|
|
49
|
-
self.tokens[model][token_type] = max(self.tokens[model].get(token_type, 0), count)
|
|
50
|
-
|
|
51
|
-
# Merge models
|
|
52
|
-
self.models.update(other.models)
|
|
53
|
-
|
|
54
|
-
# Merge PII data
|
|
55
|
-
if other.has_pii:
|
|
56
|
-
self.has_pii = True
|
|
57
|
-
self.pii_entities.update(other.pii_entities)
|
|
58
|
-
for action, entities in other.pii_actions.items():
|
|
59
|
-
self.pii_actions[action].update(entities)
|
|
60
|
-
|
|
61
|
-
# Merge violation data
|
|
62
|
-
if other.has_violation:
|
|
63
|
-
self.has_violation = True
|
|
64
|
-
self.violations.update(other.violations)
|
|
65
|
-
for action, violations in other.violation_actions.items():
|
|
66
|
-
self.violation_actions[action].update(violations)
|
|
67
|
-
|
|
68
|
-
def to_attributes(self) -> Dict[str, str]:
|
|
69
|
-
"""Convert aggregated data to span attributes."""
|
|
70
|
-
attributes = {}
|
|
71
|
-
|
|
72
|
-
# Error Data
|
|
73
|
-
attributes["has_error"] = str(self.has_error).lower()
|
|
74
|
-
if self.has_error:
|
|
75
|
-
attributes["status_codes"] = json.dumps(list(self.status_codes))
|
|
76
|
-
|
|
77
|
-
# Token usage by model
|
|
78
|
-
if self.tokens:
|
|
79
|
-
tokens_dict = {}
|
|
80
|
-
for model, usage in self.tokens.items():
|
|
81
|
-
tokens_dict[model] = dict(usage)
|
|
82
|
-
attributes["tokens"] = json.dumps(tokens_dict)
|
|
83
|
-
|
|
84
|
-
# Models used
|
|
85
|
-
if self.models:
|
|
86
|
-
attributes["models"] = json.dumps(sorted(list(self.models)))
|
|
87
|
-
|
|
88
|
-
# PII information
|
|
89
|
-
attributes["has_pii"] = str(self.has_pii).lower()
|
|
90
|
-
if self.pii_entities:
|
|
91
|
-
attributes["pii_entities"] = json.dumps(sorted(list(self.pii_entities)))
|
|
92
|
-
if self.pii_actions:
|
|
93
|
-
pii_actions_dict = {}
|
|
94
|
-
for action, entities in self.pii_actions.items():
|
|
95
|
-
pii_actions_dict[action] = sorted(list(entities))
|
|
96
|
-
attributes["pii_actions"] = json.dumps(pii_actions_dict)
|
|
97
|
-
|
|
98
|
-
# Violation information
|
|
99
|
-
attributes["has_violation"] = str(self.has_violation).lower()
|
|
100
|
-
if self.violations:
|
|
101
|
-
attributes["violations"] = json.dumps(sorted(list(self.violations)))
|
|
102
|
-
if self.violation_actions:
|
|
103
|
-
violation_actions_dict = {}
|
|
104
|
-
for action, violations in self.violation_actions.items():
|
|
105
|
-
violation_actions_dict[action] = sorted(list(violations))
|
|
106
|
-
attributes["violation_actions"] = json.dumps(violation_actions_dict)
|
|
107
|
-
|
|
108
|
-
return attributes
|
|
109
|
-
|
|
110
|
-
|
|
111
|
-
class SpanAggregationProcessor(SpanProcessor): # type: ignore[misc]
|
|
112
|
-
"""
|
|
113
|
-
OpenTelemetry span processor that aggregates data from child spans into parent spans.
|
|
114
|
-
"""
|
|
115
|
-
|
|
116
|
-
def __init__(self) -> None:
|
|
117
|
-
self._span_data: Dict[str, SpanAggregationData] = {}
|
|
118
|
-
self._span_hierarchy: Dict[str, Optional[str]] = {} # child_id -> parent_id
|
|
119
|
-
self._root_spans: Set[str] = set()
|
|
120
|
-
self._captured_data: Dict[str, Dict[str, Any]] = {} # span_id -> {attributes, events}
|
|
121
|
-
self._active_spans: Dict[str, Span] = {} # span_id -> original span reference
|
|
122
|
-
|
|
123
|
-
def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
|
|
124
|
-
"""Called when a span starts."""
|
|
125
|
-
span_id = self._get_span_id(span)
|
|
126
|
-
if not span_id:
|
|
127
|
-
return
|
|
128
|
-
|
|
129
|
-
# Store the original span for later use
|
|
130
|
-
self._active_spans[span_id] = span
|
|
131
|
-
|
|
132
|
-
# Initialize aggregation data
|
|
133
|
-
self._span_data[span_id] = SpanAggregationData()
|
|
134
|
-
self._captured_data[span_id] = {"attributes": {}, "events": []}
|
|
135
|
-
|
|
136
|
-
# Check if this is a root span (no parent)
|
|
137
|
-
if span.parent is None:
|
|
138
|
-
self._root_spans.add(span_id)
|
|
139
|
-
else:
|
|
140
|
-
# Track parent-child relationship - span.parent is a SpanContext, not a Span
|
|
141
|
-
try:
|
|
142
|
-
parent_span_context = span.parent
|
|
143
|
-
if parent_span_context and parent_span_context.span_id:
|
|
144
|
-
parent_span_id = f"{parent_span_context.trace_id:032x}-{parent_span_context.span_id:016x}"
|
|
145
|
-
self._span_hierarchy[span_id] = parent_span_id
|
|
146
|
-
else:
|
|
147
|
-
logger.warning(f"DEBUG: Parent span context is invalid for child {span_id}")
|
|
148
|
-
except Exception as e:
|
|
149
|
-
logger.warning(f"DEBUG: Could not get parent span ID for child {span_id}: {e}")
|
|
150
|
-
|
|
151
|
-
# Wrap span methods to capture data
|
|
152
|
-
self._wrap_span_methods(span, span_id)
|
|
153
|
-
|
|
154
|
-
def on_end(self, span: Span) -> None:
|
|
155
|
-
"""Called when a span ends."""
|
|
156
|
-
span_id = self._get_span_id(span)
|
|
157
|
-
if not span_id or span_id not in self._span_data:
|
|
158
|
-
return
|
|
159
|
-
|
|
160
|
-
try:
|
|
161
|
-
# Process this span's captured data
|
|
162
|
-
captured = self._captured_data.get(span_id, {})
|
|
163
|
-
self._process_attributes(self._span_data[span_id], captured.get("attributes", {}))
|
|
164
|
-
|
|
165
|
-
# Set aggregated attributes on this span
|
|
166
|
-
original_span = self._active_spans.get(span_id)
|
|
167
|
-
if original_span and original_span.is_recording():
|
|
168
|
-
self._set_span_attributes(original_span, self._span_data[span_id])
|
|
169
|
-
|
|
170
|
-
# Handle parent-child aggregation for any remaining data
|
|
171
|
-
self._aggregate_to_all_parents(span_id)
|
|
172
|
-
|
|
173
|
-
except Exception as e:
|
|
174
|
-
logger.error(f"Error during span aggregation for span {span_id}: {e}")
|
|
175
|
-
# Even if there's an error, try to do basic aggregation
|
|
176
|
-
try:
|
|
177
|
-
original_span = self._active_spans.get(span_id)
|
|
178
|
-
if original_span and original_span.is_recording():
|
|
179
|
-
self._set_span_attributes(original_span, self._span_data[span_id])
|
|
180
|
-
except Exception as inner_e:
|
|
181
|
-
logger.error(f"Failed to set basic aggregation attributes: {inner_e}")
|
|
182
|
-
|
|
183
|
-
# Clean up
|
|
184
|
-
self._span_data.pop(span_id, None)
|
|
185
|
-
self._captured_data.pop(span_id, None)
|
|
186
|
-
self._active_spans.pop(span_id, None)
|
|
187
|
-
self._root_spans.discard(span_id)
|
|
188
|
-
self._span_hierarchy.pop(span_id, None)
|
|
189
|
-
|
|
190
|
-
def _wrap_span_methods(self, span: Span, span_id: str) -> Any:
|
|
191
|
-
"""Wrap span methods to capture attributes and events."""
|
|
192
|
-
# Wrap set_attribute
|
|
193
|
-
original_set_attribute = span.set_attribute
|
|
194
|
-
|
|
195
|
-
def wrapped_set_attribute(key: str, value: Any) -> Any:
|
|
196
|
-
# Status code processing
|
|
197
|
-
if key == "http.status_code":
|
|
198
|
-
self._status_code_processing(value)
|
|
199
|
-
|
|
200
|
-
# Capture the all the attribute data
|
|
201
|
-
self._captured_data[span_id]["attributes"][key] = value
|
|
202
|
-
return original_set_attribute(key, value)
|
|
203
|
-
|
|
204
|
-
span.set_attribute = wrapped_set_attribute
|
|
205
|
-
|
|
206
|
-
# Wrap add_event
|
|
207
|
-
original_add_event = span.add_event
|
|
208
|
-
|
|
209
|
-
def wrapped_add_event(name: str, attributes: Dict[str, Any] = {}, timestamp: int = 0) -> Any:
|
|
210
|
-
# Only process PII and violation events
|
|
211
|
-
if name == "pii_detected" and attributes:
|
|
212
|
-
self._process_pii_event(self._span_data[span_id], attributes)
|
|
213
|
-
if span.is_recording():
|
|
214
|
-
self._set_span_attributes(span, self._span_data[span_id])
|
|
215
|
-
# Immediately aggregate to parent spans
|
|
216
|
-
self._aggregate_to_all_parents(span_id)
|
|
217
|
-
elif name == "violation_detected" and attributes:
|
|
218
|
-
self._process_violation_event(self._span_data[span_id], attributes)
|
|
219
|
-
if span.is_recording():
|
|
220
|
-
self._set_span_attributes(span, self._span_data[span_id])
|
|
221
|
-
# Immediately aggregate to parent spans
|
|
222
|
-
self._aggregate_to_all_parents(span_id)
|
|
223
|
-
|
|
224
|
-
# Check if span is still recording before adding event
|
|
225
|
-
if not span.is_recording():
|
|
226
|
-
logger.debug(f"Attempted to add event to ended span {span_id}")
|
|
227
|
-
return None
|
|
228
|
-
return original_add_event(name, attributes, timestamp)
|
|
229
|
-
|
|
230
|
-
span.add_event = wrapped_add_event
|
|
231
|
-
|
|
232
|
-
def _process_attributes(self, data: SpanAggregationData, attributes: Dict[str, Any]) -> None:
|
|
233
|
-
"""Process span attributes for aggregation."""
|
|
234
|
-
# Extract status code for error identification
|
|
235
|
-
status_code = attributes.get("http.status_code", 200)
|
|
236
|
-
if httpx.codes.is_error(status_code):
|
|
237
|
-
data.has_error = True
|
|
238
|
-
data.status_codes.update([status_code])
|
|
239
|
-
|
|
240
|
-
# Extract model information
|
|
241
|
-
model = attributes.get("gen_ai.request.model") or attributes.get("gen_ai.response.model")
|
|
242
|
-
if model:
|
|
243
|
-
data.models.add(model)
|
|
244
|
-
# Extract token usage
|
|
245
|
-
token_fields = {
|
|
246
|
-
"prompt_tokens": attributes.get("gen_ai.usage.prompt_tokens", 0),
|
|
247
|
-
"completion_tokens": attributes.get("gen_ai.usage.completion_tokens", 0),
|
|
248
|
-
"total_tokens": attributes.get("llm.usage.total_tokens", 0),
|
|
249
|
-
"cache_read_input_tokens": attributes.get("gen_ai.usage.cache_read_input_tokens", 0),
|
|
250
|
-
}
|
|
251
|
-
|
|
252
|
-
# Initialize token fields if they don't exist
|
|
253
|
-
if model not in data.tokens:
|
|
254
|
-
data.tokens[model] = {}
|
|
255
|
-
|
|
256
|
-
# Add token values
|
|
257
|
-
for field, value in token_fields.items():
|
|
258
|
-
if isinstance(value, (int, str)):
|
|
259
|
-
current_value = data.tokens[model].get(field, 0)
|
|
260
|
-
data.tokens[model][field] = current_value + int(value)
|
|
261
|
-
|
|
262
|
-
def _process_pii_event(self, data: SpanAggregationData, attrs: Dict[str, Any]) -> None:
|
|
263
|
-
"""Process pii_detected event."""
|
|
264
|
-
if attrs.get("has_pii"):
|
|
265
|
-
data.has_pii = True
|
|
266
|
-
|
|
267
|
-
# Extract entities from pii_entities field
|
|
268
|
-
entity_counts_str = attrs.get("pii_entities")
|
|
269
|
-
if entity_counts_str:
|
|
270
|
-
try:
|
|
271
|
-
entity_counts = (
|
|
272
|
-
json.loads(entity_counts_str) if isinstance(entity_counts_str, str) else entity_counts_str
|
|
273
|
-
)
|
|
274
|
-
if isinstance(entity_counts, dict):
|
|
275
|
-
entities = set(entity_counts.keys())
|
|
276
|
-
data.pii_entities.update(entities)
|
|
277
|
-
|
|
278
|
-
# Determine action
|
|
279
|
-
if attrs.get("is_blocked"):
|
|
280
|
-
data.pii_actions["BLOCK"].update(entities)
|
|
281
|
-
elif attrs.get("is_masked"):
|
|
282
|
-
data.pii_actions["MASK"].update(entities)
|
|
283
|
-
else:
|
|
284
|
-
data.pii_actions["FLAG"].update(entities)
|
|
285
|
-
except (json.JSONDecodeError, TypeError):
|
|
286
|
-
logger.error(f"Failed to parse pii_entities: {entity_counts_str}")
|
|
287
|
-
|
|
288
|
-
def _process_violation_event(self, data: SpanAggregationData, attrs: Dict[str, Any]) -> None:
|
|
289
|
-
"""Process violation_detected event."""
|
|
290
|
-
if attrs.get("has_violation"):
|
|
291
|
-
data.has_violation = True
|
|
292
|
-
violations = attrs.get("violations", [])
|
|
293
|
-
if violations:
|
|
294
|
-
data.violations.update(violations)
|
|
295
|
-
# Set action based on is_blocked flag
|
|
296
|
-
action = "BLOCK" if attrs.get("is_blocked") else "FLAG"
|
|
297
|
-
data.violation_actions[action].update(violations)
|
|
298
|
-
|
|
299
|
-
def _aggregate_to_all_parents(self, child_span_id: str) -> None:
|
|
300
|
-
"""Aggregate data from child span to all its parent spans in the hierarchy."""
|
|
301
|
-
if child_span_id not in self._span_data:
|
|
302
|
-
return
|
|
303
|
-
|
|
304
|
-
child_data = self._span_data[child_span_id]
|
|
305
|
-
current_span_id = child_span_id
|
|
306
|
-
|
|
307
|
-
# Traverse up the parent hierarchy
|
|
308
|
-
while True:
|
|
309
|
-
parent_id = self._span_hierarchy.get(current_span_id)
|
|
310
|
-
if not parent_id or parent_id not in self._span_data:
|
|
311
|
-
break
|
|
312
|
-
|
|
313
|
-
# Merge child data into parent
|
|
314
|
-
self._span_data[parent_id].merge_from_other(child_data)
|
|
315
|
-
|
|
316
|
-
# Update parent span attributes if it's still active and recording
|
|
317
|
-
parent_span = self._active_spans.get(parent_id)
|
|
318
|
-
if parent_span and parent_span.is_recording():
|
|
319
|
-
self._set_span_attributes(parent_span, self._span_data[parent_id])
|
|
320
|
-
|
|
321
|
-
# Move up to the next parent
|
|
322
|
-
current_span_id = parent_id
|
|
323
|
-
|
|
324
|
-
def _set_span_attributes(self, span: Span, data: SpanAggregationData) -> None:
|
|
325
|
-
"""Set aggregated attributes on the given span."""
|
|
326
|
-
try:
|
|
327
|
-
aggregated_attrs = data.to_attributes()
|
|
328
|
-
# Set all aggregated attributes under a single 'aggregator' key as a JSON object
|
|
329
|
-
span.set_attribute(f"{Config.LIBRARY_NAME}.aggregated_attributes", json.dumps(aggregated_attrs))
|
|
330
|
-
except Exception as e:
|
|
331
|
-
logger.error(f"Failed to set aggregated attributes: {e}")
|
|
332
|
-
|
|
333
|
-
def _get_span_id(self, span: Span) -> Optional[str]:
|
|
334
|
-
"""Get a unique identifier for the span."""
|
|
335
|
-
try:
|
|
336
|
-
span_context = span.get_span_context()
|
|
337
|
-
return f"{span_context.trace_id:032x}-{span_context.span_id:016x}"
|
|
338
|
-
except Exception:
|
|
339
|
-
return None
|
|
340
|
-
|
|
341
|
-
def _get_span_id_from_context(self, context: Context) -> Optional[str]:
|
|
342
|
-
"""Extract span ID from context."""
|
|
343
|
-
if context:
|
|
344
|
-
span_context = trace.get_current_span(context).get_span_context()
|
|
345
|
-
if span_context and span_context.span_id:
|
|
346
|
-
return f"{span_context.trace_id:032x}-{span_context.span_id:016x}"
|
|
347
|
-
return None
|
|
348
|
-
|
|
349
|
-
def _status_code_processing(self, status_code: int) -> None:
|
|
350
|
-
if httpx.codes.is_error(status_code):
|
|
351
|
-
event_attributes = {"has_error": True, "status_code": status_code}
|
|
352
|
-
Netra.set_custom_event(event_name="error_detected", attributes=event_attributes)
|
|
353
|
-
|
|
354
|
-
def force_flush(self, timeout_millis: int = 30000) -> bool:
|
|
355
|
-
"""Force flush any pending data."""
|
|
356
|
-
return True
|
|
357
|
-
|
|
358
|
-
def shutdown(self) -> bool:
|
|
359
|
-
"""Shutdown the processor."""
|
|
360
|
-
self._span_data.clear()
|
|
361
|
-
self._span_hierarchy.clear()
|
|
362
|
-
self._root_spans.clear()
|
|
363
|
-
self._captured_data.clear()
|
|
364
|
-
self._active_spans.clear()
|
|
365
|
-
return True
|
|
File without changes
|
|
File without changes
|