lmnr 0.5.3__py3-none-any.whl → 0.6.0__py3-none-any.whl

lmnr/__init__.py

@@ -13,8 +13,10 @@ from .sdk.types import (
 )
 from .sdk.decorators import observe
 from .sdk.types import LaminarSpanContext
-from .opentelemetry_lib import Instruments
 from .opentelemetry_lib.tracing.attributes import Attributes
+from .opentelemetry_lib.tracing.instruments import Instruments
+from .opentelemetry_lib.tracing.processor import LaminarSpanProcessor
+from .opentelemetry_lib.tracing.tracer import get_laminar_tracer_provider, get_tracer
 from opentelemetry.trace import use_span
 
 __all__ = [
@@ -29,9 +31,12 @@ __all__ = [
     "LaminarClient",
     "LaminarDataset",
     "LaminarSpanContext",
+    "LaminarSpanProcessor",
     "RunAgentResponseChunk",
     "StepChunkContent",
     "TracingLevel",
+    "get_laminar_tracer_provider",
+    "get_tracer",
     "evaluate",
     "observe",
     "use_span",

lmnr/opentelemetry_lib/__init__.py

@@ -1,19 +1,13 @@
 import sys
 
 from typing import Optional, Set
-from opentelemetry.sdk.trace import SpanProcessor
 from opentelemetry.sdk.trace.export import SpanExporter
 from opentelemetry.sdk.resources import SERVICE_NAME
-from opentelemetry.propagators.textmap import TextMapPropagator
-from opentelemetry.util.re import parse_env_headers
 
-from lmnr.opentelemetry_lib.instruments import Instruments
-from lmnr.opentelemetry_lib.config import (
-    is_content_tracing_enabled,
-    is_tracing_enabled,
-)
-from lmnr.opentelemetry_lib.tracing.tracing import TracerWrapper
-from typing import Dict
+from lmnr.opentelemetry_lib.tracing.instruments import Instruments
+from lmnr.opentelemetry_lib.tracing import TracerWrapper
+
+MAX_MANUAL_SPAN_PAYLOAD_SIZE = 1024 * 1024  # 1MB
 
 
 class TracerManager:
@@ -22,48 +16,34 @@ class TracerManager:
     @staticmethod
     def init(
         app_name: Optional[str] = sys.argv[0],
-        api_endpoint: str = "https://api.lmnr.ai",
-        api_key: Optional[str] = None,
-        headers: Dict[str, str] = {},
         disable_batch=False,
         exporter: Optional[SpanExporter] = None,
-        processor: Optional[SpanProcessor] = None,
-        propagator: Optional[TextMapPropagator] = None,
-        should_enrich_metrics: bool = False,
         resource_attributes: dict = {},
         instruments: Optional[Set[Instruments]] = None,
-        base_http_url: Optional[str] = None,
+        base_url: str = "https://api.lmnr.ai",
+        port: int = 8443,
+        http_port: int = 443,
         project_api_key: Optional[str] = None,
         max_export_batch_size: Optional[int] = None,
+        force_http: bool = False,
+        timeout_seconds: int = 30,
     ) -> None:
-        if not is_tracing_enabled():
-            return
-
-        enable_content_tracing = is_content_tracing_enabled()
-
-        if isinstance(headers, str):
-            headers = parse_env_headers(headers)
-
-        if api_key and not exporter and not processor and not headers:
-            headers = {
-                "Authorization": f"Bearer {api_key}",
-            }
+        enable_content_tracing = True
 
         # Tracer init
         resource_attributes.update({SERVICE_NAME: app_name})
-        TracerWrapper.set_static_params(
-            resource_attributes, enable_content_tracing, api_endpoint, headers
-        )
+        TracerWrapper.set_static_params(resource_attributes, enable_content_tracing)
         TracerManager.__tracer_wrapper = TracerWrapper(
             disable_batch=disable_batch,
-            processor=processor,
-            propagator=propagator,
             exporter=exporter,
-            should_enrich_metrics=should_enrich_metrics,
             instruments=instruments,
-            base_http_url=base_http_url,
+            base_url=base_url,
+            port=port,
+            http_port=http_port,
             project_api_key=project_api_key,
             max_export_batch_size=max_export_batch_size,
+            force_http=force_http,
+            timeout_seconds=timeout_seconds,
         )
 
     @staticmethod

lmnr/opentelemetry_lib/decorators/__init__.py

@@ -0,0 +1,219 @@
+from functools import wraps
+import json
+import logging
+import pydantic
+import types
+from typing import Any, Literal, Optional, Union
+
+from opentelemetry import trace
+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.opentelemetry_lib import MAX_MANUAL_SPAN_PAYLOAD_SIZE
+from lmnr.opentelemetry_lib.tracing.tracer import get_tracer
+from lmnr.opentelemetry_lib.tracing.attributes import SPAN_INPUT, SPAN_OUTPUT, SPAN_TYPE
+from lmnr.opentelemetry_lib.tracing import TracerWrapper
+from lmnr.opentelemetry_lib.utils.json_encoder import JSONEncoder
+
+
+class CustomJSONEncoder(JSONEncoder):
+    def default(self, o: Any) -> Any:
+        if isinstance(o, pydantic.BaseModel):
+            return o.model_dump_json()
+        try:
+            return super().default(o)
+        except TypeError:
+            return str(o)  # Fallback to string representation for unsupported types
+
+
+def json_dumps(data: dict) -> str:
+    try:
+        return json.dumps(data, cls=CustomJSONEncoder)
+    except Exception:
+        # Log the exception and return a placeholder if serialization completely fails
+        logging.warning("Failed to serialize data to JSON, type: %s", type(data))
+        return "{}"  # Return an empty JSON object as a fallback
+
+
+def entity_method(
+    name: Optional[str] = None,
+    ignore_input: bool = False,
+    ignore_inputs: Optional[list[str]] = None,
+    ignore_output: bool = False,
+    span_type: Union[Literal["DEFAULT"], Literal["LLM"], Literal["TOOL"]] = "DEFAULT",
+):
+    def decorate(fn):
+        @wraps(fn)
+        def wrap(*args, **kwargs):
+            if not TracerWrapper.verify_initialized():
+                return fn(*args, **kwargs)
+
+            span_name = name or fn.__name__
+
+            with get_tracer() as tracer:
+                span = tracer.start_span(span_name, attributes={SPAN_TYPE: span_type})
+
+                ctx = trace.set_span_in_context(span, context_api.get_current())
+                ctx_token = context_api.attach(ctx)
+
+                try:
+                    if not ignore_input:
+                        inp = json_dumps(
+                            get_input_from_func_args(
+                                fn,
+                                is_method=is_method(fn),
+                                func_args=args,
+                                func_kwargs=kwargs,
+                                ignore_inputs=ignore_inputs,
+                            )
+                        )
+                        if len(inp) > MAX_MANUAL_SPAN_PAYLOAD_SIZE:
+                            span.set_attribute(
+                                SPAN_INPUT, "Laminar: input too large to record"
+                            )
+                        else:
+                            span.set_attribute(SPAN_INPUT, inp)
+                except TypeError:
+                    pass
+
+                try:
+                    res = fn(*args, **kwargs)
+                except Exception as e:
+                    _process_exception(span, e)
+                    span.end()
+                    raise e
+
+                # span will be ended in the generator
+                if isinstance(res, types.GeneratorType):
+                    return _handle_generator(span, ctx_token, res)
+                if isinstance(res, types.AsyncGeneratorType):
+                    # async def foo() -> AsyncGenerator[int, None]:
+                    # is not considered async in a classical sense in Python,
+                    # so we handle this inside the sync wrapper.
+                    # In particular, CO_COROUTINE is different from CO_ASYNC_GENERATOR.
+                    # Flags are listed from LSB here:
+                    # https://docs.python.org/3/library/inspect.html#inspect-module-co-flags
+                    # See also: https://groups.google.com/g/python-tulip/c/6rWweGXLutU?pli=1
+                    return _ahandle_generator(span, ctx_token, res)
+
+                try:
+                    if not ignore_output:
+                        output = json_dumps(res)
+                        if len(output) > MAX_MANUAL_SPAN_PAYLOAD_SIZE:
+                            span.set_attribute(
+                                SPAN_OUTPUT, "Laminar: output too large to record"
+                            )
+                        else:
+                            span.set_attribute(SPAN_OUTPUT, output)
+                except TypeError:
+                    pass
+
+                span.end()
+                context_api.detach(ctx_token)
+
+                return res
+
+        return wrap
+
+    return decorate
+
+
+# Async Decorators
+def aentity_method(
+    name: Optional[str] = None,
+    ignore_input: bool = False,
+    ignore_inputs: Optional[list[str]] = None,
+    ignore_output: bool = False,
+    span_type: Union[Literal["DEFAULT"], Literal["LLM"], Literal["TOOL"]] = "DEFAULT",
+):
+    def decorate(fn):
+        @wraps(fn)
+        async def wrap(*args, **kwargs):
+            if not TracerWrapper.verify_initialized():
+                return await fn(*args, **kwargs)
+
+            span_name = name or fn.__name__
+
+            with get_tracer() as tracer:
+                span = tracer.start_span(span_name, attributes={SPAN_TYPE: span_type})
+
+                ctx = trace.set_span_in_context(span, context_api.get_current())
+                ctx_token = context_api.attach(ctx)
+
+                try:
+                    if not ignore_input:
+                        inp = json_dumps(
+                            get_input_from_func_args(
+                                fn,
+                                is_method=is_method(fn),
+                                func_args=args,
+                                func_kwargs=kwargs,
+                                ignore_inputs=ignore_inputs,
+                            )
+                        )
+                        if len(inp) > MAX_MANUAL_SPAN_PAYLOAD_SIZE:
+                            span.set_attribute(
+                                SPAN_INPUT, "Laminar: input too large to record"
+                            )
+                        else:
+                            span.set_attribute(SPAN_INPUT, inp)
+                except TypeError:
+                    pass
+
+                try:
+                    res = await fn(*args, **kwargs)
+                except Exception as e:
+                    _process_exception(span, e)
+                    span.end()
+                    raise e
+
+                # span will be ended in the generator
+                if isinstance(res, types.AsyncGeneratorType):
+                    # probably unreachable, read the comment in the similar
+                    # part of the sync wrapper.
+                    return await _ahandle_generator(span, ctx_token, res)
+
+                try:
+                    if not ignore_output:
+                        output = json_dumps(res)
+                        if len(output) > MAX_MANUAL_SPAN_PAYLOAD_SIZE:
+                            span.set_attribute(
+                                SPAN_OUTPUT, "Laminar: output too large to record"
+                            )
+                        else:
+                            span.set_attribute(SPAN_OUTPUT, output)
+                except TypeError:
+                    pass
+
+                span.end()
+                context_api.detach(ctx_token)
+
+                return res
+
+        return wrap
+
+    return decorate
+
+
+def _handle_generator(span, ctx_token, res):
+    yield from res
+
+    span.end()
+    if ctx_token is not None:
+        context_api.detach(ctx_token)
+
+
+async def _ahandle_generator(span, ctx_token, res):
+    # async with contextlib.aclosing(res) as closing_gen:
+    async for part in res:
+        yield part
+
+    span.end()
+    if ctx_token is not None:
+        context_api.detach(ctx_token)
+
+
+def _process_exception(span: Span, e: Exception):
+    # Note that this `escaped` is sent as a StringValue("True"), not a boolean.
+    span.record_exception(e, escaped=True)

lmnr/opentelemetry_lib/tracing/__init__.py

@@ -1 +1,139 @@
-from lmnr.opentelemetry_lib.tracing.context_manager import get_tracer
+import atexit
+import logging
+
+from lmnr.opentelemetry_lib.tracing.processor import LaminarSpanProcessor
+from lmnr.sdk.client.asynchronous.async_client import AsyncLaminarClient
+from lmnr.sdk.client.synchronous.sync_client import LaminarClient
+from lmnr.sdk.log import VerboseColorfulFormatter
+from lmnr.opentelemetry_lib.tracing.instruments import (
+    Instruments,
+    init_instrumentations,
+)
+
+from opentelemetry import trace
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider, SpanProcessor
+from opentelemetry.sdk.trace.export import SpanExporter
+
+from typing import Optional, Set
+
+
+module_logger = logging.getLogger(__name__)
+console_log_handler = logging.StreamHandler()
+console_log_handler.setFormatter(VerboseColorfulFormatter())
+module_logger.addHandler(console_log_handler)
+
+
+TRACER_NAME = "lmnr.tracer"
+
+MAX_EVENTS_OR_ATTRIBUTES_PER_SPAN = 5000
+
+
+class TracerWrapper(object):
+    resource_attributes: dict = {}
+    enable_content_tracing: bool = True
+    __tracer_provider: Optional[TracerProvider] = None
+    __logger: logging.Logger
+    __client: LaminarClient
+    __async_client: AsyncLaminarClient
+    __resource: Resource
+    __span_processor: SpanProcessor
+
+    def __new__(
+        cls,
+        disable_batch=False,
+        exporter: Optional[SpanExporter] = None,
+        instruments: Optional[Set[Instruments]] = None,
+        block_instruments: Optional[Set[Instruments]] = None,
+        base_url: str = "https://api.lmnr.ai",
+        port: int = 8443,
+        http_port: int = 443,
+        project_api_key: Optional[str] = None,
+        max_export_batch_size: Optional[int] = None,
+        force_http: bool = False,
+        timeout_seconds: int = 10,
+    ) -> "TracerWrapper":
+        base_http_url = f"{base_url}:{http_port}"
+        cls._initialize_logger(cls)
+        if not hasattr(cls, "instance"):
+            obj = cls.instance = super(TracerWrapper, cls).__new__(cls)
+
+            obj.__client = LaminarClient(
+                base_url=base_http_url,
+                project_api_key=project_api_key,
+            )
+            obj.__async_client = AsyncLaminarClient(
+                base_url=base_http_url,
+                project_api_key=project_api_key,
+            )
+
+            obj.__resource = Resource(attributes=TracerWrapper.resource_attributes)
+            obj.__tracer_provider = TracerProvider(resource=obj.__resource)
+
+            obj.__span_processor = LaminarSpanProcessor(
+                base_url=base_url,
+                api_key=project_api_key,
+                port=http_port if force_http else port,
+                exporter=exporter,
+                max_export_batch_size=max_export_batch_size,
+                timeout_seconds=timeout_seconds,
+                force_http=force_http,
+                disable_batch=disable_batch,
+            )
+
+            obj.__tracer_provider.add_span_processor(obj.__span_processor)
+
+            init_instrumentations(
+                tracer_provider=obj.__tracer_provider,
+                instruments=instruments,
+                block_instruments=block_instruments,
+                client=obj.__client,
+                async_client=obj.__async_client,
+            )
+
+            # Force flushes for debug environments (e.g. local development)
+            atexit.register(obj.exit_handler)
+
+        return cls.instance
+
+    def exit_handler(self):
+        if isinstance(self.__span_processor, LaminarSpanProcessor):
+            self.__span_processor.clear()
+        self.flush()
+
+    def _initialize_logger(self):
+        self.__logger = logging.getLogger(__name__)
+        console_log_handler = logging.StreamHandler()
+        console_log_handler.setFormatter(VerboseColorfulFormatter())
+        self.__logger.addHandler(console_log_handler)
+
+    @staticmethod
+    def set_static_params(
+        resource_attributes: dict,
+        enable_content_tracing: bool,
+    ) -> None:
+        TracerWrapper.resource_attributes = resource_attributes
+        TracerWrapper.enable_content_tracing = enable_content_tracing
+
+    @classmethod
+    def verify_initialized(cls) -> bool:
+        return hasattr(cls, "instance")
+
+    @classmethod
+    def clear(cls):
+        # Any state cleanup. Now used in between tests
+        if isinstance(cls.instance.__span_processor, LaminarSpanProcessor):
+            cls.instance.__span_processor.clear()
+
+    def shutdown(self):
+        if self.__tracer_provider is None:
+            return
+        self.__tracer_provider.shutdown()
+
+    def flush(self):
+        return self.__span_processor.force_flush()
+
+    def get_tracer(self):
+        if self.__tracer_provider is None:
+            return trace.get_tracer_provider().get_tracer(TRACER_NAME)
+        return self.__tracer_provider.get_tracer(TRACER_NAME)