netra-sdk 0.1.2__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.

@@ -86,10 +86,6 @@ def init_instrumentations(
86
86
  if CustomInstruments.AIO_PIKA in netra_custom_instruments:
87
87
  init_aio_pika_instrumentation()
88
88
 
89
- # Initialize aiohttp_server instrumentation.
90
- if CustomInstruments.AIOHTTP_SERVER in netra_custom_instruments:
91
- init_aiohttp_server_instrumentation()
92
-
93
89
  # Initialize aiokafka instrumentation.
94
90
  if CustomInstruments.AIOKAFKA in netra_custom_instruments:
95
91
  init_aiokafka_instrumentation()
@@ -114,10 +110,6 @@ def init_instrumentations(
114
110
  if CustomInstruments.AWS_LAMBDA in netra_custom_instruments:
115
111
  init_aws_lambda_instrumentation()
116
112
 
117
- # Initialize boto instrumentation.
118
- if CustomInstruments.BOTO in netra_custom_instruments:
119
- init_boto_instrumentation()
120
-
121
113
  # Initialize boto3sqs instrumentation.
122
114
  if CustomInstruments.BOTO3SQS in netra_custom_instruments:
123
115
  init_boto3sqs_instrumentation()
@@ -210,10 +202,6 @@ def init_instrumentations(
210
202
  if CustomInstruments.PYMYSQL in netra_custom_instruments:
211
203
  init_pymysql_instrumentation()
212
204
 
213
- # Initialize pyramid instrumentation.
214
- if CustomInstruments.PYRAMID in netra_custom_instruments:
215
- init_pyramid_instrumentation()
216
-
217
205
  # Initialize redis instrumentation.
218
206
  if CustomInstruments.REDIS in netra_custom_instruments:
219
207
  init_redis_instrumentation()
@@ -293,23 +281,12 @@ def init_fastapi_instrumentation() -> bool:
293
281
  bool: True if initialization was successful, False otherwise.
294
282
  """
295
283
  try:
296
- if not is_package_installed("fastapi"):
297
- return True
298
- from fastapi import FastAPI
299
-
300
- original_init = FastAPI.__init__
301
-
302
- def _patched_init(self: FastAPI, *args: Any, **kwargs: Any) -> None:
303
- original_init(self, *args, **kwargs)
284
+ if is_package_installed("fastapi"):
285
+ from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
304
286
 
305
- try:
306
- from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
307
-
308
- FastAPIInstrumentor().instrument_app(self)
309
- except Exception as e:
310
- logging.warning(f"Failed to auto-instrument FastAPI: {e}")
311
-
312
- FastAPI.__init__ = _patched_init
287
+ instrumentor = FastAPIInstrumentor()
288
+ if not instrumentor.is_instrumented_by_opentelemetry:
289
+ instrumentor.instrument()
313
290
  return True
314
291
  except Exception as e:
315
292
  logging.error(f"Error initializing FastAPI instrumentor: {e}")
@@ -455,22 +432,6 @@ def init_aio_pika_instrumentation() -> bool:
455
432
  return False
456
433
 
457
434
 
458
- def init_aiohttp_server_instrumentation() -> bool:
459
- """Initialize aiohttp_server instrumentation."""
460
- try:
461
- if is_package_installed("aiohttp"):
462
- from opentelemetry.instrumentation.aiohttp_server import AioHttpServerInstrumentor
463
-
464
- instrumentor = AioHttpServerInstrumentor()
465
- if not instrumentor.is_instrumented_by_opentelemetry:
466
- instrumentor.instrument()
467
- return True
468
- except Exception as e:
469
- logging.error(f"Error initializing aiohttp_server instrumentor: {e}")
470
- Telemetry().log_exception(e)
471
- return False
472
-
473
-
474
435
  def init_aiokafka_instrumentation() -> bool:
475
436
  """Initialize aiokafka instrumentation."""
476
437
  try:
@@ -567,22 +528,6 @@ def init_aws_lambda_instrumentation() -> bool:
567
528
  return False
568
529
 
569
530
 
570
- def init_boto_instrumentation() -> bool:
571
- """Initialize boto instrumentation."""
572
- try:
573
- if is_package_installed("boto"):
574
- from opentelemetry.instrumentation.boto import BotoInstrumentor
575
-
576
- instrumentor = BotoInstrumentor()
577
- if not instrumentor.is_instrumented_by_opentelemetry:
578
- instrumentor.instrument()
579
- return True
580
- except Exception as e:
581
- logging.error(f"Error initializing boto instrumentor: {e}")
582
- Telemetry().log_exception(e)
583
- return False
584
-
585
-
586
531
  def init_boto3sqs_instrumentation() -> bool:
587
532
  """Initialize boto3sqs instrumentation."""
588
533
  try:
@@ -618,7 +563,7 @@ def init_botocore_instrumentation() -> bool:
618
563
  def init_cassandra_instrumentation() -> bool:
619
564
  """Initialize cassandra instrumentation."""
620
565
  try:
621
- if is_package_installed("cassandra-driver"):
566
+ if is_package_installed("cassandra-driver") and is_package_installed("scylla-driver"):
622
567
  from opentelemetry.instrumentation.cassandra import CassandraInstrumentor
623
568
 
624
569
  instrumentor = CassandraInstrumentor()
@@ -951,22 +896,6 @@ def init_pymysql_instrumentation() -> bool:
951
896
  return False
952
897
 
953
898
 
954
- def init_pyramid_instrumentation() -> bool:
955
- """Initialize pyramid instrumentation."""
956
- try:
957
- if is_package_installed("pyramid"):
958
- from opentelemetry.instrumentation.pyramid import PyramidInstrumentor
959
-
960
- instrumentor = PyramidInstrumentor()
961
- if not instrumentor.is_instrumented_by_opentelemetry:
962
- instrumentor.instrument()
963
- return True
964
- except Exception as e:
965
- logging.error(f"Error initializing pyramid instrumentor: {e}")
966
- Telemetry().log_exception(e)
967
- return False
968
-
969
-
970
899
  def init_redis_instrumentation() -> bool:
971
900
  """Initialize redis instrumentation."""
972
901
  try:
@@ -90,13 +90,11 @@ NetraInstruments follows the given structure. Refer this for usage within Netra
90
90
 
91
91
  class InstrumentSet(Enum):
92
92
  AIOHTTP = "aiohttp"
93
- AIOHTTP_SERVER = "aiohttp_server"
94
93
  AIO_PIKA = "aio_pika"
95
94
  AIOKAFKA = "aiokafka"
96
95
  AIOPG = "aiopg"
97
96
  ALEPHALPHA = "alephalpha"
98
97
  ANTHROPIC = "anthropic"
99
- ASGI = "asgi"
100
98
  ASYNCCLICK = "asyncclick"
101
99
  ASYNCIO = "asyncio"
102
100
  ASYNCPG = "asyncpg"
@@ -145,7 +143,6 @@ class InstrumentSet(Enum):
145
143
  PYMONGO = "pymongo"
146
144
  PYMSSQL = "pymssql"
147
145
  PYMYSQL = "pymysql"
148
- PYRAMID = "pyramid"
149
146
  QDRANTDB = "qdrant_db"
150
147
  REDIS = "redis"
151
148
  REMOULADE = "remoulade"
@@ -166,5 +163,4 @@ class InstrumentSet(Enum):
166
163
  VERTEXAI = "vertexai"
167
164
  WATSONX = "watsonx"
168
165
  WEAVIATEDB = "weaviate_db"
169
- WSGI = "wsgi"
170
166
  """
@@ -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__ = ["SpanAggregationProcessor", "SessionSpanProcessor"]
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 SessionSpanProcessor, SpanAggregationProcessor
69
+ from netra.processors import ErrorDetectionProcessor, SessionSpanProcessor
70
70
 
71
71
  provider.add_span_processor(SessionSpanProcessor())
72
- provider.add_span_processor(SpanAggregationProcessor())
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.2"
1
+ __version__ = "0.1.4"
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: netra-sdk
3
- Version: 0.1.2
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
@@ -22,14 +22,12 @@ Requires-Dist: llm-guard (>=0.3.16,<0.4.0) ; extra == "llm-guard"
22
22
  Requires-Dist: opentelemetry-api (>=1.34.0,<2.0.0)
23
23
  Requires-Dist: opentelemetry-instrumentation-aio-pika (>=0.55b1,<1.0.0)
24
24
  Requires-Dist: opentelemetry-instrumentation-aiohttp-client (>=0.55b1,<1.0.0)
25
- Requires-Dist: opentelemetry-instrumentation-aiohttp-server (>=0.55b1,<1.0.0)
26
25
  Requires-Dist: opentelemetry-instrumentation-aiokafka (>=0.55b1,<1.0.0)
27
26
  Requires-Dist: opentelemetry-instrumentation-aiopg (>=0.55b1,<1.0.0)
28
27
  Requires-Dist: opentelemetry-instrumentation-asyncclick (>=0.55b1,<1.0.0)
29
28
  Requires-Dist: opentelemetry-instrumentation-asyncio (>=0.55b1,<1.0.0)
30
29
  Requires-Dist: opentelemetry-instrumentation-asyncpg (>=0.55b1,<1.0.0)
31
30
  Requires-Dist: opentelemetry-instrumentation-aws-lambda (>=0.55b1,<1.0.0)
32
- Requires-Dist: opentelemetry-instrumentation-boto (>=0.55b1,<1.0.0)
33
31
  Requires-Dist: opentelemetry-instrumentation-boto3sqs (>=0.55b1,<1.0.0)
34
32
  Requires-Dist: opentelemetry-instrumentation-botocore (>=0.55b1,<1.0.0)
35
33
  Requires-Dist: opentelemetry-instrumentation-cassandra (>=0.55b1,<1.0.0)
@@ -55,7 +53,6 @@ Requires-Dist: opentelemetry-instrumentation-pymemcache (>=0.55b1,<1.0.0)
55
53
  Requires-Dist: opentelemetry-instrumentation-pymongo (>=0.55b1,<1.0.0)
56
54
  Requires-Dist: opentelemetry-instrumentation-pymssql (>=0.55b1,<1.0.0)
57
55
  Requires-Dist: opentelemetry-instrumentation-pymysql (>=0.55b1,<1.0.0)
58
- Requires-Dist: opentelemetry-instrumentation-pyramid (>=0.55b1,<1.0.0)
59
56
  Requires-Dist: opentelemetry-instrumentation-redis (>=0.55b1,<1.0.0)
60
57
  Requires-Dist: opentelemetry-instrumentation-remoulade (>=0.55b1,<1.0.0)
61
58
  Requires-Dist: opentelemetry-instrumentation-requests (>=0.55b1,<1.0.0)
@@ -228,7 +225,6 @@ class CustomerSupportAgent:
228
225
  - **Django** - High-level Python web framework
229
226
  - **Flask** - Lightweight WSGI web application framework
230
227
  - **Falcon** - High-performance Python web framework
231
- - **Pyramid** - Small, fast, down-to-earth Python web framework
232
228
  - **Starlette** - Lightweight ASGI framework/toolkit
233
229
  - **Tornado** - Asynchronous networking library and web framework
234
230
  - **gRPC** - High-performance, open-source universal RPC framework
@@ -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=CvoAQF_tHzDph-CAypPL-eBvaM53JDtJwDTHZ7XPPeE,41170
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
@@ -20,7 +20,7 @@ netra/instrumentation/google_genai/utils.py,sha256=2OeSN5jUaMKF4x5zWiW65R1LB_a44
20
20
  netra/instrumentation/google_genai/version.py,sha256=Hww1duZrC8kYK7ThBSQVyz0HNOb0ys_o8Pln-wVQ1hI,23
21
21
  netra/instrumentation/httpx/__init__.py,sha256=w1su_eQP_w5ZJHq0Lf-4miF5zM4OOW0ItmRp0wi85Ew,19388
22
22
  netra/instrumentation/httpx/version.py,sha256=ZRQKbgDaGz_yuLk-cUKuk6ZBKCSRKZC8nQd041NRNXk,23
23
- netra/instrumentation/instruments.py,sha256=5U4r9XMMfjccd9BWHAbHMo1WcCUClnbU2oWTLAAqMw4,4345
23
+ netra/instrumentation/instruments.py,sha256=M_-4N1YML-Lc1Jb2dIzoHF9sCmsFz4euUvm8a-VKROM,4247
24
24
  netra/instrumentation/mistralai/__init__.py,sha256=RE0b-rS6iXdoynJMFKHL9s97eYo5HghrJa013fR4ZhI,18910
25
25
  netra/instrumentation/mistralai/config.py,sha256=XCyo3mk30qkvqyCqeTrKwROahu0gcOEwmbDLOo53J5k,121
26
26
  netra/instrumentation/mistralai/utils.py,sha256=nhdIer5gJFxuGwg8FCT222hggDHeMQDhJctnDSwLqcc,894
@@ -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=f8Ck-uCdDwNgEvm3hm6I9wFSVlnRYjmSfrEIQkdBaZQ,218
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=WZDBIpPp3oL8MkKBhlw0gxAXXSd1FY3bicLUKt2GNxc,3540
38
- netra/version.py,sha256=YvuYzWnKtqBb-IqG8HAu-nhIYAsgj9Vmc_b9o7vO-js,22
39
- netra_sdk-0.1.2.dist-info/LICENCE,sha256=8B_UoZ-BAl0AqiHAHUETCgd3I2B9yYJ1WEQtVb_qFMA,11359
40
- netra_sdk-0.1.2.dist-info/METADATA,sha256=0gBniDy_OiHTa27cSm0sZxWO68qF4ieLt5TKugq1Pbk,24630
41
- netra_sdk-0.1.2.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
42
- netra_sdk-0.1.2.dist-info/RECORD,,
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