netra-sdk 0.1.4__tar.gz → 0.1.6__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 (43) hide show
  1. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/PKG-INFO +34 -59
  2. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/README.md +33 -58
  3. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/decorators.py +13 -2
  4. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/processors/error_detection_processor.py +20 -2
  5. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/processors/session_span_processor.py +6 -0
  6. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/session.py +29 -53
  7. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/session_manager.py +86 -1
  8. netra_sdk-0.1.6/netra/version.py +1 -0
  9. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/pyproject.toml +1 -1
  10. netra_sdk-0.1.4/netra/version.py +0 -1
  11. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/LICENCE +0 -0
  12. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/__init__.py +0 -0
  13. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/anonymizer/__init__.py +0 -0
  14. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/anonymizer/anonymizer.py +0 -0
  15. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/anonymizer/base.py +0 -0
  16. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/anonymizer/fp_anonymizer.py +0 -0
  17. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/config.py +0 -0
  18. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/exceptions/__init__.py +0 -0
  19. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/exceptions/injection.py +0 -0
  20. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/exceptions/pii.py +0 -0
  21. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/input_scanner.py +0 -0
  22. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/__init__.py +0 -0
  23. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/aiohttp/__init__.py +0 -0
  24. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/aiohttp/version.py +0 -0
  25. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/cohere/__init__.py +0 -0
  26. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/cohere/version.py +0 -0
  27. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/google_genai/__init__.py +0 -0
  28. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/google_genai/config.py +0 -0
  29. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/google_genai/utils.py +0 -0
  30. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/google_genai/version.py +0 -0
  31. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/httpx/__init__.py +0 -0
  32. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/httpx/version.py +0 -0
  33. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/instruments.py +0 -0
  34. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/mistralai/__init__.py +0 -0
  35. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/mistralai/config.py +0 -0
  36. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/mistralai/utils.py +0 -0
  37. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/mistralai/version.py +0 -0
  38. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/weaviate/__init__.py +0 -0
  39. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/instrumentation/weaviate/version.py +0 -0
  40. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/pii.py +0 -0
  41. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/processors/__init__.py +0 -0
  42. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/scanner.py +0 -0
  43. {netra_sdk-0.1.4 → netra_sdk-0.1.6}/netra/tracer.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: netra-sdk
3
- Version: 0.1.4
3
+ Version: 0.1.6
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
@@ -407,28 +407,38 @@ Netra.set_custom_event(event_name="conversion", attributes={
407
407
  Use the custom session tracking utility to track external API calls with detailed observability:
408
408
 
409
409
  ```python
410
- from netra import Netra
411
-
412
- # Initialize SDK
413
- Netra.init(app_name="My App")
414
-
415
- # Use session context manager for tracking API calls
416
- with Netra.start_session("video_generation_task") as session:
417
- # Set attributes before the API call
418
- session.set_prompt("A cat playing piano")
419
- session.set_image_height("1024")
420
- session.set_image_width("1024")
421
- session.set_model("stable-diffusion-xl")
422
-
423
- # Make your external API call
424
- result = external_api.generate_video(...)
425
-
426
- # Set post-call attributes
427
- session.set_tokens("1250")
428
- session.set_credits("30")
429
- session.set_cost("0.15")
430
-
431
- # Add events during session
410
+ from netra import Netra, Session
411
+ from netra.session import UsageModel
412
+
413
+ # Start a new session
414
+ with Netra.start_session("image_generation") as session:
415
+ # Set session attributes
416
+ session.set_prompt("A beautiful sunset over mountains")
417
+ session.set_negative_prompt("blurry, low quality")
418
+ session.set_height("1024")
419
+ session.set_width("1024")
420
+ session.set_model("dall-e-3")
421
+ session.set_llm_system("openai")
422
+
423
+ # Set usage data with UsageModel
424
+ usage_data = [
425
+ UsageModel(
426
+ model="dall-e-3",
427
+ type="image_generation",
428
+ unit_used=1,
429
+ cost_in_usd=0.02
430
+ )
431
+ ]
432
+ session.set_usage(usage_data)
433
+
434
+ # Your API calls here
435
+ # ...
436
+
437
+ # Set custom attributes
438
+ session.set_attribute("custom_key", "custom_value")
439
+
440
+ # Add events
441
+ session.add_event("generation_started", {"step": 1, "status": "processing"})
432
442
  session.add_event("processing_completed", {"step": "rendering"})
433
443
 
434
444
  # Session automatically captures duration, status, and any errors
@@ -662,42 +672,7 @@ pre-commit install --hook-type pre-push
662
672
 
663
673
  ## 🤝 Contributing
664
674
 
665
- We welcome contributions! Please follow these guidelines:
666
-
667
- ### Commit Message Format
668
-
669
- We use [Conventional Commits](https://www.conventionalcommits.org/) for commit messages:
670
-
671
- ```
672
- <type>[optional scope]: <description>
673
-
674
- [optional body]
675
-
676
- [optional footer(s)]
677
- ```
678
-
679
- **Types:**
680
- - **feat**: A new feature
681
- - **fix**: A bug fix
682
- - **docs**: Documentation only changes
683
- - **style**: Changes that do not affect the meaning of the code
684
- - **refactor**: A code change that neither fixes a bug nor adds a feature
685
- - **perf**: A code change that improves performance
686
- - **test**: Adding missing tests or correcting existing tests
687
- - **chore**: Changes to the build process or auxiliary tools
688
-
689
- **Examples:**
690
- ```
691
- feat: add support for Claude AI instrumentation
692
- fix(pii): resolve masking issue with nested objects
693
- docs: update installation instructions
694
- ```
695
-
696
- **Scope** can be used to specify the area of change (e.g., `pii`, `instrumentation`, `decorators`).
697
-
698
- **Body** should include the motivation for the change and contrast with previous behavior.
699
-
700
- **Footer** can be used for "BREAKING CHANGE:" or issue references.
675
+ We welcome contributions! Please see our [Contributing Guide](CONTRIBUTING.md) for detailed information on how to contribute to the project, including development setup, testing, and our commit message format.
701
676
 
702
677
  ---
703
678
 
@@ -331,28 +331,38 @@ Netra.set_custom_event(event_name="conversion", attributes={
331
331
  Use the custom session tracking utility to track external API calls with detailed observability:
332
332
 
333
333
  ```python
334
- from netra import Netra
335
-
336
- # Initialize SDK
337
- Netra.init(app_name="My App")
338
-
339
- # Use session context manager for tracking API calls
340
- with Netra.start_session("video_generation_task") as session:
341
- # Set attributes before the API call
342
- session.set_prompt("A cat playing piano")
343
- session.set_image_height("1024")
344
- session.set_image_width("1024")
345
- session.set_model("stable-diffusion-xl")
346
-
347
- # Make your external API call
348
- result = external_api.generate_video(...)
349
-
350
- # Set post-call attributes
351
- session.set_tokens("1250")
352
- session.set_credits("30")
353
- session.set_cost("0.15")
354
-
355
- # Add events during session
334
+ from netra import Netra, Session
335
+ from netra.session import UsageModel
336
+
337
+ # Start a new session
338
+ with Netra.start_session("image_generation") as session:
339
+ # Set session attributes
340
+ session.set_prompt("A beautiful sunset over mountains")
341
+ session.set_negative_prompt("blurry, low quality")
342
+ session.set_height("1024")
343
+ session.set_width("1024")
344
+ session.set_model("dall-e-3")
345
+ session.set_llm_system("openai")
346
+
347
+ # Set usage data with UsageModel
348
+ usage_data = [
349
+ UsageModel(
350
+ model="dall-e-3",
351
+ type="image_generation",
352
+ unit_used=1,
353
+ cost_in_usd=0.02
354
+ )
355
+ ]
356
+ session.set_usage(usage_data)
357
+
358
+ # Your API calls here
359
+ # ...
360
+
361
+ # Set custom attributes
362
+ session.set_attribute("custom_key", "custom_value")
363
+
364
+ # Add events
365
+ session.add_event("generation_started", {"step": 1, "status": "processing"})
356
366
  session.add_event("processing_completed", {"step": "rendering"})
357
367
 
358
368
  # Session automatically captures duration, status, and any errors
@@ -586,41 +596,6 @@ pre-commit install --hook-type pre-push
586
596
 
587
597
  ## 🤝 Contributing
588
598
 
589
- We welcome contributions! Please follow these guidelines:
590
-
591
- ### Commit Message Format
592
-
593
- We use [Conventional Commits](https://www.conventionalcommits.org/) for commit messages:
594
-
595
- ```
596
- <type>[optional scope]: <description>
597
-
598
- [optional body]
599
-
600
- [optional footer(s)]
601
- ```
602
-
603
- **Types:**
604
- - **feat**: A new feature
605
- - **fix**: A bug fix
606
- - **docs**: Documentation only changes
607
- - **style**: Changes that do not affect the meaning of the code
608
- - **refactor**: A code change that neither fixes a bug nor adds a feature
609
- - **perf**: A code change that improves performance
610
- - **test**: Adding missing tests or correcting existing tests
611
- - **chore**: Changes to the build process or auxiliary tools
612
-
613
- **Examples:**
614
- ```
615
- feat: add support for Claude AI instrumentation
616
- fix(pii): resolve masking issue with nested objects
617
- docs: update installation instructions
618
- ```
619
-
620
- **Scope** can be used to specify the area of change (e.g., `pii`, `instrumentation`, `decorators`).
621
-
622
- **Body** should include the motivation for the change and contrast with previous behavior.
623
-
624
- **Footer** can be used for "BREAKING CHANGE:" or issue references.
599
+ We welcome contributions! Please see our [Contributing Guide](CONTRIBUTING.md) for detailed information on how to contribute to the project, including development setup, testing, and our commit message format.
625
600
 
626
601
  ---
@@ -12,6 +12,7 @@ from typing import Any, Awaitable, Callable, Dict, Optional, ParamSpec, Tuple, T
12
12
  from opentelemetry import trace
13
13
 
14
14
  from .config import Config
15
+ from .session_manager import SessionManager
15
16
 
16
17
  P = ParamSpec("P")
17
18
  R = TypeVar("R")
@@ -78,6 +79,9 @@ def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optio
78
79
 
79
80
  @functools.wraps(func)
80
81
  async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
82
+ # Push entity to stack before span starts so SessionSpanProcessor can capture it
83
+ SessionManager.push_entity(entity_type, span_name)
84
+
81
85
  tracer = trace.get_tracer(module_name)
82
86
  with tracer.start_as_current_span(span_name) as span:
83
87
  _add_span_attributes(span, func, args, kwargs, entity_type)
@@ -87,8 +91,10 @@ def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optio
87
91
  return result
88
92
  except Exception as e:
89
93
  span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
90
- span.record_exception(e)
91
94
  raise
95
+ finally:
96
+ # Pop entity from stack after function call is done
97
+ SessionManager.pop_entity(entity_type)
92
98
 
93
99
  return cast(Callable[P, R], async_wrapper)
94
100
 
@@ -96,6 +102,9 @@ def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optio
96
102
 
97
103
  @functools.wraps(func)
98
104
  def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
105
+ # Push entity to stack before span starts so SessionSpanProcessor can capture it
106
+ SessionManager.push_entity(entity_type, span_name)
107
+
99
108
  tracer = trace.get_tracer(module_name)
100
109
  with tracer.start_as_current_span(span_name) as span:
101
110
  _add_span_attributes(span, func, args, kwargs, entity_type)
@@ -105,8 +114,10 @@ def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optio
105
114
  return result
106
115
  except Exception as e:
107
116
  span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
108
- span.record_exception(e)
109
117
  raise
118
+ finally:
119
+ # Pop entity from stack after function call is done
120
+ SessionManager.pop_entity(entity_type)
110
121
 
111
122
  return cast(Callable[P, R], sync_wrapper)
112
123
 
@@ -1,9 +1,9 @@
1
1
  import logging
2
- from typing import Any, Optional
2
+ from typing import Any, Optional, Union
3
3
 
4
4
  import httpx
5
5
  from opentelemetry.sdk.trace import SpanProcessor
6
- from opentelemetry.trace import Context, Span
6
+ from opentelemetry.trace import Context, Span, Status, StatusCode
7
7
 
8
8
  from netra import Netra
9
9
 
@@ -63,4 +63,22 @@ class ErrorDetectionProcessor(SpanProcessor): # type: ignore[misc]
63
63
 
64
64
  return original_set_attribute(key, value)
65
65
 
66
+ # Wrap set_status
67
+ original_set_status = span.set_status
68
+
69
+ def wrapped_set_status(status: Union[Status, StatusCode]) -> Any:
70
+ # Check if status code is ERROR
71
+ if isinstance(status, Status):
72
+ status_code = status.status_code
73
+ elif isinstance(status, StatusCode):
74
+ status_code = status
75
+ if status_code == StatusCode.ERROR:
76
+ event_attributes = {
77
+ "has_error": True,
78
+ }
79
+ Netra.set_custom_event(event_name="error_detected", attributes=event_attributes)
80
+
81
+ return original_set_status(status)
82
+
66
83
  span.set_attribute = wrapped_set_attribute
84
+ span.set_status = wrapped_set_status
@@ -42,6 +42,12 @@ class SessionSpanProcessor(SpanProcessor): # type: ignore[misc]
42
42
  value = baggage.get_baggage(f"custom.{key}", ctx)
43
43
  if value:
44
44
  span.set_attribute(f"{Config.LIBRARY_NAME}.custom.{key}", value)
45
+
46
+ # Add entity attributes from SessionManager
47
+ entity_attributes = SessionManager.get_current_entity_attributes()
48
+ for attr_key, attr_value in entity_attributes.items():
49
+ span.set_attribute(attr_key, attr_value)
50
+
45
51
  except Exception as e:
46
52
  logger.exception(f"Error setting span attributes: {e}")
47
53
 
@@ -1,11 +1,13 @@
1
+ import json
1
2
  import logging
2
3
  import time
3
- from typing import Any, Dict, Literal, Optional
4
+ from typing import Any, Dict, List, Literal, Optional
4
5
 
5
6
  from opentelemetry import context as context_api
6
7
  from opentelemetry import trace
7
8
  from opentelemetry.trace import SpanKind, Status, StatusCode
8
9
  from opentelemetry.trace.propagation import set_span_in_context
10
+ from pydantic import BaseModel
9
11
 
10
12
  from netra.config import Config
11
13
 
@@ -14,22 +16,22 @@ logging.basicConfig(level=logging.INFO)
14
16
  logger = logging.getLogger(__name__)
15
17
 
16
18
 
19
+ class UsageModel(BaseModel): # type: ignore[misc]
20
+ model: str
21
+ type: str
22
+ unit_used: Optional[int] = None
23
+ cost_in_usd: Optional[float] = None
24
+
25
+
17
26
  class ATTRIBUTE:
18
27
  LLM_SYSTEM = "llm_system"
19
28
  MODEL = "model"
20
29
  PROMPT = "prompt"
21
30
  NEGATIVE_PROMPT = "negative_prompt"
22
- IMAGE_HEIGHT = "image_height"
23
- IMAGE_WIDTH = "image_width"
24
- CREDITS = "credits"
25
- TOTAL_COST = "total_cost"
26
- TOTAL_TOKENS = "total_tokens"
27
- PROMPT_TOKENS_COST = "prompt_tokens_cost"
28
- PROMPT_TOKENS = "prompt_tokens"
29
- COMPLETION_TOKENS_COST = "completion_tokens_cost"
30
- COMPLETION_TOKENS = "completion_tokens"
31
- CACHED_TOKENS_COST = "cached_tokens_cost"
32
- CACHED_TOKENS = "cached_tokens"
31
+ HEIGHT = "height"
32
+ WIDTH = "width"
33
+ OUTPUT_TYPE = "output_type"
34
+ USAGE = "usage"
33
35
  STATUS = "status"
34
36
  DURATION_MS = "duration_ms"
35
37
  ERROR_MESSAGE = "error_message"
@@ -46,7 +48,7 @@ class Session:
46
48
  # External API call
47
49
  result = external_api.generate_video(...)
48
50
 
49
- session.set_tokens("20").set_credits("30")
51
+ session.set_usage(usage_data)
50
52
  """
51
53
 
52
54
  def __init__(self, name: str, attributes: Optional[Dict[str, str]] = None, module_name: str = "combat_sdk"):
@@ -139,49 +141,23 @@ class Session:
139
141
  """Set the negative prompt."""
140
142
  return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.NEGATIVE_PROMPT}", negative_prompt)
141
143
 
142
- def set_image_height(self, height: str) -> "Session":
143
- """Set the image height."""
144
- return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.IMAGE_HEIGHT}", height)
145
-
146
- def set_image_width(self, width: str) -> "Session":
147
- """Set the image width."""
148
- return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.IMAGE_WIDTH}", width)
149
-
150
- def set_total_tokens(self, tokens: str) -> "Session":
151
- """Set the number of tokens used."""
152
- return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.TOTAL_TOKENS}", tokens)
153
-
154
- def set_prompt_tokens_cost(self, cost: str) -> "Session":
155
- """Set the number of tokens used."""
156
- return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.PROMPT_TOKENS_COST}", cost)
157
-
158
- def set_prompt_tokens(self, tokens: str) -> "Session":
159
- """Set the number of tokens used."""
160
- return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.PROMPT_TOKENS}", tokens)
161
-
162
- def set_completion_tokens_cost(self, cost: str) -> "Session":
163
- """Set the number of tokens used."""
164
- return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.COMPLETION_TOKENS_COST}", cost)
165
-
166
- def set_completion_tokens(self, tokens: str) -> "Session":
167
- """Set the number of tokens used."""
168
- return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.COMPLETION_TOKENS}", tokens)
169
-
170
- def set_cached_tokens_cost(self, cost: str) -> "Session":
171
- """Set the number of tokens used."""
172
- return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.CACHED_TOKENS_COST}", cost)
144
+ def set_height(self, height: str) -> "Session":
145
+ """Set the height."""
146
+ return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.HEIGHT}", height)
173
147
 
174
- def set_cached_tokens(self, tokens: str) -> "Session":
175
- """Set the number of tokens used."""
176
- return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.CACHED_TOKENS}", tokens)
148
+ def set_width(self, width: str) -> "Session":
149
+ """Set the width."""
150
+ return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.WIDTH}", width)
177
151
 
178
- def set_credits(self, credits: str) -> "Session":
179
- """Set the number of credits used."""
180
- return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.CREDITS}", credits)
152
+ def set_output_type(self, output_type: str) -> "Session":
153
+ """Set the output type."""
154
+ return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.OUTPUT_TYPE}", output_type)
181
155
 
182
- def set_total_cost(self, cost: str) -> "Session":
183
- """Set the cost of the operation."""
184
- return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.TOTAL_COST}", cost)
156
+ def set_usage(self, usage: List[UsageModel]) -> "Session":
157
+ """Set the usage data as a JSON string."""
158
+ usage_dict = [u.model_dump() for u in usage]
159
+ usage_json = json.dumps(usage_dict)
160
+ return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.USAGE}", usage_json)
185
161
 
186
162
  def set_model(self, model: str) -> "Session":
187
163
  """Set the model used."""
@@ -5,7 +5,7 @@ Handles automatic session and user ID management for applications.
5
5
 
6
6
  import logging
7
7
  from datetime import datetime
8
- from typing import Any, Dict, Optional, Union
8
+ from typing import Any, Dict, List, Optional, Union
9
9
 
10
10
  from opentelemetry import baggage
11
11
  from opentelemetry import context as otel_context
@@ -22,6 +22,11 @@ class SessionManager:
22
22
  # Class variable to track the current span
23
23
  _current_span: Optional[trace.Span] = None
24
24
 
25
+ # Class variables to track separate entity stacks
26
+ _workflow_stack: List[str] = []
27
+ _task_stack: List[str] = []
28
+ _agent_stack: List[str] = []
29
+
25
30
  @classmethod
26
31
  def set_current_span(cls, span: Optional[trace.Span]) -> None:
27
32
  """
@@ -42,6 +47,86 @@ class SessionManager:
42
47
  """
43
48
  return cls._current_span
44
49
 
50
+ @classmethod
51
+ def push_entity(cls, entity_type: str, entity_name: str) -> None:
52
+ """
53
+ Push an entity onto the appropriate entity stack.
54
+
55
+ Args:
56
+ entity_type: Type of entity (workflow, task, agent)
57
+ entity_name: Name of the entity
58
+ """
59
+ if entity_type == "workflow":
60
+ cls._workflow_stack.append(entity_name)
61
+ elif entity_type == "task":
62
+ cls._task_stack.append(entity_name)
63
+ elif entity_type == "agent":
64
+ cls._agent_stack.append(entity_name)
65
+
66
+ @classmethod
67
+ def pop_entity(cls, entity_type: str) -> Optional[str]:
68
+ """
69
+ Pop the most recent entity from the specified entity stack.
70
+
71
+ Args:
72
+ entity_type: Type of entity (workflow, task, agent)
73
+
74
+ Returns:
75
+ Entity name or None if stack is empty
76
+ """
77
+ if entity_type == "workflow" and cls._workflow_stack:
78
+ return cls._workflow_stack.pop()
79
+ elif entity_type == "task" and cls._task_stack:
80
+ return cls._task_stack.pop()
81
+ elif entity_type == "agent" and cls._agent_stack:
82
+ return cls._agent_stack.pop()
83
+ return None
84
+
85
+ @classmethod
86
+ def get_current_entity_attributes(cls) -> Dict[str, str]:
87
+ """
88
+ Get current entity attributes for span annotation.
89
+
90
+ Returns:
91
+ Dictionary of entity attributes to add to spans
92
+ """
93
+ attributes = {}
94
+
95
+ # Add current workflow if exists
96
+ if cls._workflow_stack:
97
+ attributes[f"{Config.LIBRARY_NAME}.workflow.name"] = cls._workflow_stack[-1]
98
+
99
+ # Add current task if exists
100
+ if cls._task_stack:
101
+ attributes[f"{Config.LIBRARY_NAME}.task.name"] = cls._task_stack[-1]
102
+
103
+ # Add current agent if exists
104
+ if cls._agent_stack:
105
+ attributes[f"{Config.LIBRARY_NAME}.agent.name"] = cls._agent_stack[-1]
106
+
107
+ return attributes
108
+
109
+ @classmethod
110
+ def clear_entity_stacks(cls) -> None:
111
+ """Clear all entity stacks."""
112
+ cls._workflow_stack.clear()
113
+ cls._task_stack.clear()
114
+ cls._agent_stack.clear()
115
+
116
+ @classmethod
117
+ def get_stack_info(cls) -> Dict[str, List[str]]:
118
+ """
119
+ Get information about all current stacks.
120
+
121
+ Returns:
122
+ Dictionary containing all stack contents
123
+ """
124
+ return {
125
+ "workflows": cls._workflow_stack.copy(),
126
+ "tasks": cls._task_stack.copy(),
127
+ "agents": cls._agent_stack.copy(),
128
+ }
129
+
45
130
  @staticmethod
46
131
  def set_session_context(session_key: str, value: Union[str, Dict[str, str]]) -> None:
47
132
  """
@@ -0,0 +1 @@
1
+ __version__ = "0.1.6"
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
4
4
 
5
5
  [project]
6
6
  name = "netra-sdk"
7
- version = "0.1.4"
7
+ version = "0.1.6"
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 +0,0 @@
1
- __version__ = "0.1.4"
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes