lmnr 0.5.0__py3-none-any.whl → 0.5.1a0__py3-none-any.whl

lmnr/openllmetry_sdk/__init__.py

@@ -1,5 +1,6 @@
 import sys
 
+from contextlib import contextmanager
 from typing import Optional, Set
 from opentelemetry.sdk.trace import SpanProcessor
 from opentelemetry.sdk.trace.export import SpanExporter
@@ -18,6 +19,7 @@ from typing import Dict
 
 class TracerManager:
     __tracer_wrapper: TracerWrapper
+    __initialized: bool = False
 
     @staticmethod
     def init(
@@ -51,9 +53,6 @@ class TracerManager:
 
         # Tracer init
         resource_attributes.update({SERVICE_NAME: app_name})
-        TracerWrapper.set_static_params(
-            resource_attributes, enable_content_tracing, api_endpoint, headers
-        )
         TracerManager.__tracer_wrapper = TracerWrapper(
             disable_batch=disable_batch,
             processor=processor,
@@ -64,9 +63,41 @@ class TracerManager:
             base_http_url=base_http_url,
             project_api_key=project_api_key,
             max_export_batch_size=max_export_batch_size,
+            resource_attributes=resource_attributes,
+            enable_content_tracing=enable_content_tracing,
+            endpoint=api_endpoint,
+            headers=headers,
         )
+        TracerManager.__initialized = True
+
+    @staticmethod
+    def flush() -> bool:
+        return TracerManager.__tracer_wrapper.flush()
+
+    @staticmethod
+    def shutdown() -> bool:
+        try:
+            res = TracerManager.__tracer_wrapper.shutdown()
+            TracerManager.__tracer_wrapper = None
+            TracerManager.__initialized = False
+            return res
+        except Exception:
+            return False
 
     @staticmethod
-    def flush():
-        if getattr(TracerManager, "__tracer_wrapper", None):
-            TracerManager.__tracer_wrapper.flush()
+    def is_initialized() -> bool:
+        return TracerManager.__initialized
+
+    @staticmethod
+    def get_tracer_wrapper() -> TracerWrapper:
+        return TracerManager.__tracer_wrapper
+
+
+@contextmanager
+def get_tracer(flush_on_exit: bool = False):
+    wrapper = TracerManager.get_tracer_wrapper()
+    try:
+        yield wrapper.get_tracer()
+    finally:
+        if flush_on_exit:
+            wrapper.flush()

lmnr/openllmetry_sdk/decorators/base.py

@@ -11,9 +11,9 @@ from opentelemetry import context as context_api
 from opentelemetry.trace import Span
 
 from lmnr.sdk.utils import get_input_from_func_args, is_method
-from lmnr.openllmetry_sdk.tracing import get_tracer
+from lmnr.openllmetry_sdk import get_tracer
 from lmnr.openllmetry_sdk.tracing.attributes import SPAN_INPUT, SPAN_OUTPUT, SPAN_TYPE
-from lmnr.openllmetry_sdk.tracing.tracing import TracerWrapper
+from lmnr.openllmetry_sdk import TracerManager
 from lmnr.openllmetry_sdk.utils.json_encoder import JSONEncoder
 from lmnr.openllmetry_sdk.config import MAX_MANUAL_SPAN_PAYLOAD_SIZE
 
@@ -46,7 +46,7 @@ def entity_method(
     def decorate(fn):
         @wraps(fn)
         def wrap(*args, **kwargs):
-            if not TracerWrapper.verify_initialized():
+            if not TracerManager.is_initialized():
                 return fn(*args, **kwargs)
 
             span_name = name or fn.__name__
@@ -114,7 +114,7 @@ def aentity_method(
     def decorate(fn):
         @wraps(fn)
         async def wrap(*args, **kwargs):
-            if not TracerWrapper.verify_initialized():
+            if not TracerManager.is_initialized():
                 return await fn(*args, **kwargs)
 
             span_name = name or fn.__name__

lmnr/openllmetry_sdk/tracing/__init__.py

@@ -1 +0,0 @@
-from lmnr.openllmetry_sdk.tracing.context_manager import get_tracer