novastack-telemetry 1.0.0__tar.gz

novastack_telemetry-1.0.0/.gitignore

@@ -0,0 +1,85 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+#IDE
+.DS_Store
+.idea
+.vscode
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Mkdocs documentation
+docs/_build/
+docs/api_reference/site/
+
+# Ruff
+.ruff_cache/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+.python-version
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/

novastack_telemetry-1.0.0/PKG-INFO

@@ -0,0 +1,31 @@
+Metadata-Version: 2.4
+Name: novastack-telemetry
+Version: 1.0.0
+Summary: Provides the telemetry hooks for observability in Novastack.
+Project-URL: Repository, https://github.com/novastack-project/novastack/tree/main/novastack-telemetry
+Author-email: Leonardo Furnielis <leonardofurnielis@outlook.com>
+License: Apache-2.0
+Keywords: AI,LLM,QA,RAG,data,observability,retrieval,semantic-search
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Application Frameworks
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: <3.14,>=3.11
+Requires-Dist: pydantic<3.0.0,>=2.12.5
+Requires-Dist: treelib<2.0.0,>=1.8.0
+Requires-Dist: wrapt<3.0.0,>=2.2.1
+Provides-Extra: dev
+Requires-Dist: pytest-asyncio<2.0.0,>=1.3.0; extra == 'dev'
+Requires-Dist: pytest-mock<4.0.0,>=3.15.1; extra == 'dev'
+Requires-Dist: pytest<10.0.0,>=9.0.3; extra == 'dev'
+Requires-Dist: ruff<1.0.0,>=0.15.8; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# Novastack telemetry
+
+This project provides the telemetry hooks for observability in Novastack.
+
+## Installation 
+
+```bash
+pip install novastack-telemetry
+```

novastack_telemetry-1.0.0/README.md

@@ -0,0 +1,9 @@
+# Novastack telemetry
+
+This project provides the telemetry hooks for observability in Novastack.
+
+## Installation 
+
+```bash
+pip install novastack-telemetry
+```

novastack_telemetry-1.0.0/novastack_telemetry/__init__.py

@@ -0,0 +1,9 @@
+from novastack_telemetry.dispatcher import get_dispatcher
+from novastack_telemetry.events import SpanExceptionEvent
+from novastack_telemetry.mixin import DispatcherSpanMixin
+
+__all__ = [
+    "DispatcherSpanMixin",
+    "get_dispatcher",
+    "SpanExceptionEvent",
+]

novastack_telemetry-1.0.0/novastack_telemetry/_dispatcher_core.py

@@ -0,0 +1,370 @@
+import asyncio
+import inspect
+import logging
+import uuid
+from contextlib import contextmanager
+from contextvars import Context, ContextVar, Token, copy_context
+from functools import partial
+from typing import Any, Callable, Generator, Optional, TypeVar
+
+import wrapt
+from pydantic import BaseModel, Field
+
+from novastack_telemetry.events import BaseEvent, SpanExceptionEvent
+from novastack_telemetry.observability import BaseObservability
+from novastack_telemetry.span import _active_span_id
+
+_CONTEXT_METADATA_KEY = "context_metadata"
+_DISPATCHER_SPAN_DECORATED_ATTR = "__dispatcher_span_decorated__"
+
+_logger = logging.getLogger(__name__)
+
+# ContextVar for managing active context metadata
+_active_context_metadata: ContextVar[dict[str, Any]] = ContextVar(
+    "_context_metadata", default={}
+)
+_R = TypeVar("_R")
+
+
+@contextmanager
+def _context_metadata(new_metadata: dict[str, Any]) -> Generator[None, None, None]:
+    token = _active_context_metadata.set(new_metadata)
+    try:
+        yield
+    finally:
+        _active_context_metadata.reset(token)
+
+
+class Dispatcher(BaseModel):
+    """
+    Orchestrates telemetry events and span lifecycle management.
+
+    Routes events and span signals to registered handlers through a hierarchical
+    propagation chain. Provides a decorator-based API (@dispatcher.span) for
+    automatic span tracking in both sync and async contexts. Supports thread-safe
+    and async-safe operations using ContextVars, with automatic parent-child span
+    relationships, trace context management, and silent exception handling to
+    prevent telemetry failures from affecting application code.
+    """
+
+    model_config = {"arbitrary_types_allowed": True}
+    name: str = Field(default_factory=str, description="The name of dispatcher")
+    handlers: list[BaseObservability] = Field(
+        default=[], description="List of unified handlers"
+    )
+    parent_name: str = Field(
+        default_factory=str, description="The name of parent Dispatcher."
+    )
+    dispatcher_manager: Optional["_DispatcherManager"] = Field(
+        default=None, description="Dispatcher manager."
+    )
+    root_name: str = Field(
+        default="root", description="The name of Dispatcher root tree."
+    )
+    propagate: bool = Field(
+        default=True,
+        description="Whether to propagate the event to parent dispatchers and their handlers",
+    )
+
+    @property
+    def parent(self) -> "Dispatcher":
+        assert self.dispatcher_manager is not None
+        return self.dispatcher_manager.dispatchers[self.parent_name]
+
+    @property
+    def root(self) -> "Dispatcher":
+        assert self.dispatcher_manager is not None
+        return self.dispatcher_manager.dispatchers[self.root_name]
+
+    def _get_handler_hierarchy(self) -> Generator[BaseObservability, None, None]:
+        """Retrieve every handlers reachable via the propagation chain."""
+        h: Dispatcher | None = self
+        while h:
+            yield from h.handlers
+            if not h.propagate:
+                break
+            h = h.parent
+
+    def _dispatch_to_handlers(
+        self, handler_method: str, *args: Any, **kwargs: Any
+    ) -> None:
+        """
+        Invoke handler method across the propagation chain with error isolation.
+
+        Calls the specified method on all handlers in the chain. Handler exceptions
+        are silently caught to ensure telemetry failures don't break application code.
+
+        Args:
+            handler_method: Name of the handler method to call
+            *args: Positional arguments to pass to the handler method
+            **kwargs: Keyword arguments to pass to the handler method
+        """
+        for h in self._get_handler_hierarchy():
+            try:
+                getattr(h, handler_method)(*args, **kwargs)
+            except BaseException:
+                pass
+
+    def add_handler(self, handler: BaseObservability) -> None:
+        """Add handler to set of handlers."""
+        self.handlers += [handler]
+
+    def event(self, event: BaseEvent, **kwargs: Any) -> None:
+        """Dispatch event to all registered handlers."""
+        event.metadata.update(_active_context_metadata.get())
+        self._dispatch_to_handlers("on_event", event, **kwargs)
+
+    def _span_start(
+        self,
+        id_: str,
+        bound_args: inspect.BoundArguments,
+        instance: Any | None = None,
+        parent_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Internal: Send notice to handlers that a span with id_ has started."""
+        self._dispatch_to_handlers(
+            "on_span_start",
+            id_=id_,
+            bound_args=bound_args,
+            instance=instance,
+            parent_id=parent_id,
+            metadata=metadata,
+            **kwargs,
+        )
+
+    def _span_end(
+        self,
+        id_: str,
+        bound_args: inspect.BoundArguments,
+        instance: Any | None = None,
+        result: Any | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Internal: Send notice to handlers that a span with id_ is exiting."""
+        self._dispatch_to_handlers(
+            "on_span_end",
+            id_=id_,
+            bound_args=bound_args,
+            instance=instance,
+            result=result,
+            **kwargs,
+        )
+
+    def _span_exception(
+        self,
+        id_: str,
+        bound_args: inspect.BoundArguments,
+        instance: Any | None = None,
+        err: BaseException | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Internal: Send notice to handlers that a span with id_ is being exited due an exception."""
+        self._dispatch_to_handlers(
+            "on_span_exception",
+            id_=id_,
+            bound_args=bound_args,
+            instance=instance,
+            err=err,
+            **kwargs,
+        )
+
+    def capture_propagation_context(self) -> dict[str, Any]:
+        """
+        Capture trace propagation context from all handlers and active metadata.
+
+        Returns a serializable dictionary with namespaced handler data and context
+        metadata, suitable for cross-process propagation via restore_propagation_context().
+        """
+        result: dict[str, Any] = {}
+        for h in self._get_handler_hierarchy():
+            try:
+                result.update(h.capture_propagation_context())
+            except BaseException:
+                _logger.warning("Error capturing propagation context", exc_info=True)
+        metadata = _active_context_metadata.get()
+        if metadata:
+            result[_CONTEXT_METADATA_KEY] = dict(metadata)
+        return result
+
+    def restore_propagation_context(self, context: dict[str, Any]) -> None:
+        """
+        Restore trace propagation context across all handlers and metadata.
+
+        Applies the context to all handlers in the chain and updates active
+        context metadata for subsequent span operations.
+        """
+        for h in self._get_handler_hierarchy():
+            try:
+                h.restore_propagation_context(context)
+            except BaseException:
+                _logger.warning("Error restoring propagation context", exc_info=True)
+        metadata = context.get(_CONTEXT_METADATA_KEY)
+        if metadata:
+            _active_context_metadata.set(dict(metadata))
+
+    def shutdown(self) -> None:
+        """
+        Gracefully shutdown all handlers in the propagation chain.
+
+        Invokes shutdown() on each handler while suppressing exceptions to ensure
+        complete cleanup even if individual handler fail.
+        """
+        for h in self._get_handler_hierarchy():
+            try:
+                h.shutdown()
+            except BaseException:
+                _logger.warning("Error closing handler %s", h, exc_info=True)
+
+    def span(self, func: Callable[..., _R]) -> Callable[..., _R]:
+        try:
+            if hasattr(func, _DISPATCHER_SPAN_DECORATED_ATTR):
+                return func
+            setattr(func, _DISPATCHER_SPAN_DECORATED_ATTR, True)
+        except AttributeError:
+            pass
+
+        @wrapt.decorator
+        def wrapper(func: Callable, instance: Any, args: list, kwargs: dict) -> Any:
+            bound_args = inspect.signature(func).bind(*args, **kwargs)
+            if instance is not None:
+                actual_class = type(instance).__name__
+                method_name = func.__name__
+                id_ = f"{actual_class}.{method_name}-{uuid.uuid4()}"
+            else:
+                id_ = f"{func.__qualname__}-{uuid.uuid4()}"
+            metadata = _active_context_metadata.get()
+            result = None
+
+            # Copy the current context
+            context = copy_context()
+
+            token = _active_span_id.set(id_)
+            parent_id = None if token.old_value is Token.MISSING else token.old_value
+            self._span_start(
+                id_=id_,
+                bound_args=bound_args,
+                instance=instance,
+                parent_id=parent_id,
+                metadata=metadata,
+            )
+
+            def handle_future_result(
+                future: asyncio.Future,
+                span_id: str,
+                bound_args: inspect.BoundArguments,
+                instance: Any,
+                context: Context,
+            ) -> None:
+                try:
+                    result = None if future.exception() else future.result()
+
+                    self._span_end(
+                        id_=span_id,
+                        bound_args=bound_args,
+                        instance=instance,
+                        result=result,
+                    )
+                    return result
+                except BaseException as e:
+                    self.event(SpanExceptionEvent(span_id=span_id, err_str=str(e)))
+                    self._span_exception(
+                        id_=span_id, bound_args=bound_args, instance=instance, err=e
+                    )
+                    raise
+                finally:
+                    try:
+                        context.run(_active_span_id.reset, token)
+                    except ValueError as e:
+                        _logger.debug(f"Failed to reset _active_span_id: {e}")
+
+            try:
+                result = func(*args, **kwargs)
+                if isinstance(result, asyncio.Future):
+                    new_future = asyncio.ensure_future(result)
+                    new_future.add_done_callback(
+                        partial(
+                            handle_future_result,
+                            span_id=id_,
+                            bound_args=bound_args,
+                            instance=instance,
+                            context=context,
+                        )
+                    )
+                    return new_future
+                else:
+                    self._span_end(
+                        id_=id_, bound_args=bound_args, instance=instance, result=result
+                    )
+                    return result
+            except BaseException as e:
+                self.event(SpanExceptionEvent(span_id=id_, err_str=str(e)))
+                self._span_exception(
+                    id_=id_, bound_args=bound_args, instance=instance, err=e
+                )
+                raise
+            finally:
+                if not isinstance(result, asyncio.Future):
+                    _active_span_id.reset(token)
+
+        @wrapt.decorator
+        async def async_wrapper(
+            func: Callable, instance: Any, args: list, kwargs: dict
+        ) -> Any:
+            bound_args = inspect.signature(func).bind(*args, **kwargs)
+            if instance is not None:
+                actual_class = type(instance).__name__
+                method_name = func.__name__
+                id_ = f"{actual_class}.{method_name}-{uuid.uuid4()}"
+            else:
+                id_ = f"{func.__qualname__}-{uuid.uuid4()}"
+            metadata = _active_context_metadata.get()
+
+            token = _active_span_id.set(id_)
+            parent_id = None if token.old_value is Token.MISSING else token.old_value
+            self._span_start(
+                id_=id_,
+                bound_args=bound_args,
+                instance=instance,
+                parent_id=parent_id,
+                metadata=metadata,
+            )
+            try:
+                result = await func(*args, **kwargs)
+            except BaseException as e:
+                self.event(SpanExceptionEvent(span_id=id_, err_str=str(e)))
+                self._span_exception(
+                    id_=id_, bound_args=bound_args, instance=instance, err=e
+                )
+                raise
+            else:
+                self._span_end(
+                    id_=id_, bound_args=bound_args, instance=instance, result=result
+                )
+                return result
+            finally:
+                _active_span_id.reset(token)
+
+        if inspect.iscoroutinefunction(func):
+            return async_wrapper(func)  # type: ignore
+        else:
+            return wrapper(func)  # type: ignore
+
+    @property
+    def log_name(self) -> str:
+        if self.parent:
+            return f"{self.parent.name}.{self.name}"
+        else:
+            return self.name
+
+
+class _DispatcherManager:
+    def __init__(self, root: Dispatcher) -> None:
+        self.dispatchers: dict[str, Dispatcher] = {root.name: root}
+
+    def add_dispatcher(self, d: Dispatcher) -> None:
+        self.dispatchers.setdefault(d.name, d)
+
+
+Dispatcher.model_rebuild()

novastack_telemetry-1.0.0/novastack_telemetry/dispatcher.py

@@ -0,0 +1,43 @@
+from novastack_telemetry._dispatcher_core import Dispatcher, _DispatcherManager
+
+root_dispatcher: Dispatcher = Dispatcher(
+    name="root",
+    handlers=[],
+    propagate=False,
+)
+
+root_manager: _DispatcherManager = _DispatcherManager(root_dispatcher)
+
+
+def get_dispatcher(name: str = "root") -> Dispatcher:
+    """
+    Get or create a Dispatcher by name with hierarchical parent resolution.
+
+    Returns existing dispatcher or creates a new one. Parent is determined by
+    dot-notation hierarchy (e.g., "a.b.c" → parent "a.b"), falling back to "root".
+
+    Args:
+        name: The name of the dispatcher. Defaults to "root".
+    """
+    # Return existing dispatcher if found
+    if existing := root_manager.dispatchers.get(name):
+        return existing
+
+    # Determine parent: try hierarchical parent first, fallback to root
+    candidate_parent_name = ".".join(name.split(".")[:-1])
+    parent_name = (
+        candidate_parent_name
+        if candidate_parent_name in root_manager.dispatchers
+        else "root"
+    )
+
+    # Create and register new dispatcher
+    new_dispatcher = Dispatcher(
+        name=name,
+        root_name=root_dispatcher.name,
+        parent_name=parent_name,
+        dispatcher_manager=root_manager,
+    )
+    root_manager.add_dispatcher(new_dispatcher)
+
+    return new_dispatcher

novastack_telemetry-1.0.0/novastack_telemetry/events/__init__.py

@@ -0,0 +1,6 @@
+from novastack_telemetry.events.base import BaseEvent, SpanExceptionEvent
+
+__all__ = [
+    "BaseEvent",
+    "SpanExceptionEvent",
+]

novastack_telemetry-1.0.0/novastack_telemetry/events/base.py

@@ -0,0 +1,30 @@
+from datetime import datetime
+from typing import Any
+from uuid import uuid4
+
+from pydantic import BaseModel, Field
+
+from novastack_telemetry.span import _active_span_id
+
+
+class BaseEvent(BaseModel):
+    model_config = {"arbitrary_types_allowed": True}
+
+    timestamp: datetime = Field(default_factory=lambda: datetime.now())
+    id_: str = Field(default_factory=lambda: str(uuid4()))
+    span_id: str | None = Field(default_factory=_active_span_id.get)
+    metadata: dict[str, Any] = Field(default={})
+
+    @classmethod
+    def class_name(cls) -> str:
+        return "BaseEvent"
+
+
+class SpanExceptionEvent(BaseEvent):
+    """SpanExceptionEvent."""
+
+    err_str: str
+
+    @classmethod
+    def class_name(cls) -> str:
+        return "SpanExceptionEvent"

novastack_telemetry-1.0.0/novastack_telemetry/mixin.py

@@ -0,0 +1,66 @@
+import inspect
+from abc import ABC
+from typing import Any
+
+from novastack_telemetry._dispatcher_core import _DISPATCHER_SPAN_DECORATED_ATTR
+from novastack_telemetry.dispatcher import get_dispatcher
+
+
+class DispatcherSpanMixin(ABC):
+    """
+    Automatically applies `dispatcher.span` to methods that override decorated methods
+    from base classes. This ensures that when you override a method that was decorated
+    with `@dispatcher.span` in a parent class, the override will also be automatically
+    decorated.
+    """
+
+    @staticmethod
+    def _is_abstract_method(method: Any) -> bool:
+        """Check if a method is abstract."""
+        return getattr(method, "__isabstractmethod__", False)
+
+    @staticmethod
+    def _is_decorated_method(method: Any) -> bool:
+        """Check if a method is decorated with dispatcher.span."""
+        return hasattr(method, _DISPATCHER_SPAN_DECORATED_ATTR)
+
+    @classmethod
+    def _collect_methods_to_decorate(cls, target_cls: type) -> set[str]:
+        """
+        Collect method names from base classes that should be decorated
+        when overridden in subclasses.
+
+        Only collects methods that were explicitly decorated with @dispatcher.span.
+        """
+        decorated_methods: set[str] = set()
+
+        # Iterate through base classes (excluding the target class itself)
+        for base_cls in inspect.getmro(target_cls)[1:]:
+            for attr, method in base_cls.__dict__.items():
+                if not callable(method):
+                    continue
+
+                # Only collect methods that are already decorated
+                if cls._is_decorated_method(method):
+                    decorated_methods.add(attr)
+
+        return decorated_methods
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        super().__init_subclass__(**kwargs)
+
+        # Collect methods that need decoration from base classes
+        methods_to_decorate = cls._collect_methods_to_decorate(cls)
+
+        # Get dispatcher for this module
+        dispatcher = get_dispatcher(cls.__module__)
+
+        # Decorate overridden methods in the current class
+        for attr, method in cls.__dict__.items():
+            # Skip non-callable or abstract methods
+            if not callable(method) or cls._is_abstract_method(method):
+                continue
+
+            # Decorate if this method overrides a decorated method
+            if attr in methods_to_decorate:
+                setattr(cls, attr, dispatcher.span(method))

novastack_telemetry-1.0.0/novastack_telemetry/observability/__init__.py

@@ -0,0 +1,9 @@
+from novastack_telemetry.observability.base import BaseObservability
+from novastack_telemetry.observability.novastack_debug import (
+    NovastackDebugObservability,
+)
+
+__all__ = [
+    "BaseObservability",
+    "NovastackDebugObservability",
+]

novastack_telemetry-1.0.0/novastack_telemetry/observability/base.py

@@ -0,0 +1,59 @@
+import inspect
+from typing import Any
+
+from pydantic import BaseModel
+
+from novastack_telemetry.events.base import BaseEvent
+from novastack_telemetry.span import Span
+
+
+class BaseObservability(BaseModel):
+    """
+    Unified observability that can handle events and spans.
+
+    This is the foundation for moving toward to observability.
+    BaseObservability can implement event handling and span handling.
+    """
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    @classmethod
+    def class_name(cls) -> str:
+        return "BaseObservability"
+
+    def on_event(self, event: BaseEvent, **kwargs: Any) -> Any:
+        """Handle an event."""
+
+    def on_span_start(
+        self,
+        id_: str,
+        bound_args: inspect.BoundArguments,
+        instance: Any | None = None,
+        parent_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Handle span start."""
+
+    def on_span_end(
+        self,
+        id_: str,
+        bound_args: inspect.BoundArguments,
+        instance: Any | None = None,
+        result: Any | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Handle span end."""
+
+    def on_span_exception(
+        self,
+        id_: str,
+        bound_args: inspect.BoundArguments,
+        instance: Any | None = None,
+        err: BaseException | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Handle span exception."""
+
+    def shutdown(self) -> None:
+        """Optional cleanup hook called during dispatcher shutdown."""

novastack_telemetry-1.0.0/novastack_telemetry/observability/novastack_debug.py

@@ -0,0 +1,233 @@
+import inspect
+import threading
+from collections import defaultdict
+from datetime import datetime
+from typing import Any, Optional
+
+from pydantic import Field, PrivateAttr
+from treelib.tree import Tree
+
+from novastack_telemetry.events import BaseEvent
+from novastack_telemetry.observability.base import BaseObservability
+from novastack_telemetry.span import Span
+
+
+class NovastackDebugObservability(BaseObservability):
+    """Novastack observability to track debug information (in-memory)."""
+
+    model_config = {"arbitrary_types_allowed": True}
+    open_spans: dict[str, Span] = Field(
+        default_factory=dict, description="Dictionary of open spans."
+    )
+    completed_spans: list[Span] = Field(
+        default_factory=list, description="List of completed spans."
+    )
+    dropped_spans: list[Span] = Field(
+        default_factory=list, description="List of dropped spans."
+    )
+    events: list[BaseEvent] = Field(default_factory=list, description="List of events.")
+    _lock: Optional[threading.Lock] = PrivateAttr(default=None)
+
+    print_span_on_end: bool = Field(
+        default=True,
+        description="Automatically print trace tree when a root span completes.",
+    )
+
+    @property
+    def lock(self) -> threading.Lock:
+        if self._lock is None:
+            self._lock = threading.Lock()
+        return self._lock
+
+    @classmethod
+    def class_name(cls) -> str:
+        return "NovastackDebugObservability"
+
+    def on_event(self, event: BaseEvent, **kwargs: Any) -> None:
+        """Handle an event."""
+        with self.lock:
+            self.events.append(event)
+
+    def on_span_start(
+        self,
+        id_: str,
+        bound_args: inspect.BoundArguments,
+        instance: Any | None = None,
+        parent_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Handle span start."""
+        span = Span(id_=id_, parent_id=parent_id, metadata=metadata or {})
+        with self.lock:
+            self.open_spans[id_] = span
+
+    def on_span_end(
+        self,
+        id_: str,
+        bound_args: inspect.BoundArguments,
+        instance: Any | None = None,
+        result: Any | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Handle span end."""
+        with self.lock:
+            span = self.open_spans.pop(id_, None)
+            if span:
+                span.end_time = datetime.now()
+                span.duration = (span.end_time - span.start_time).total_seconds()
+                self.completed_spans.append(span)
+
+                # Auto-print trace tree if enabled and this is a root span
+                if self.print_span_on_end and span.parent_id is None:
+                    self.print_trace_trees(include_events=True)
+
+    def on_span_exception(
+        self,
+        id_: str,
+        bound_args: inspect.BoundArguments,
+        instance: Any | None = None,
+        err: BaseException | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Handle span exception."""
+        with self.lock:
+            span = self.open_spans.pop(id_, None)
+            if span:
+                span.metadata["error"] = str(err)
+                self.dropped_spans.append(span)
+
+    def _get_parents(self) -> list[Span]:
+        """Extract root spans from completed and dropped span collections."""
+        all_spans = self.completed_spans + self.dropped_spans
+        return [s for s in all_spans if s.parent_id is None]
+
+    def _build_spans_by_parent_index(
+        self, spans: list[Span]
+    ) -> dict[str | None, list[Span]]:
+        """Build optimized parent-to-children span index for O(1) hierarchy traversal."""
+        spans_by_parent: dict[str | None, list[Span]] = defaultdict(list)
+        for span in spans:
+            spans_by_parent[span.parent_id].append(span)
+        return spans_by_parent
+
+    def _build_tree_by_parent(
+        self, parent: Span, spans_by_parent: dict[str | None, list[Span]]
+    ) -> list[Span]:
+        """Recursively construct flattened span hierarchy using depth-first traversal."""
+        result = [parent]
+        children = spans_by_parent.get(parent.id_, [])
+
+        for child in children:
+            result.extend(self._build_tree_by_parent(child, spans_by_parent))
+
+        return result
+
+    def _get_trace_trees(self, include_events: bool = True) -> list[Tree]:
+        """Generate hierarchical tree structures representing execution traces with span durations."""
+        all_spans = self.completed_spans + self.dropped_spans
+        for s in all_spans:
+            if s.parent_id is None:
+                continue
+            if not any(ns.id_ == s.parent_id for ns in all_spans):
+                s.parent_id += "-missing"
+                all_spans.append(Span(id_=s.parent_id, parent_id=None))
+
+        # Build index once for O(n) tree building
+        spans_by_parent = self._build_spans_by_parent_index(all_spans)
+
+        parents = self._get_parents()
+        span_groups = []
+        for p in parents:
+            this_span_group = self._build_tree_by_parent(
+                parent=p, spans_by_parent=spans_by_parent
+            )
+            sorted_span_group = sorted(this_span_group, key=lambda x: x.start_time)
+            span_groups.append(sorted_span_group)
+
+        trees = []
+        tree = Tree()
+        for grp in span_groups:
+            for span in grp:
+                if span.parent_id is None:
+                    if tree.all_nodes():
+                        trees.append(tree)
+                        tree = Tree()
+
+                duration_str = f"{span.duration:.6f}s" if span.duration else ""
+                tree.create_node(
+                    tag=f"{span.id_} - {duration_str}",
+                    identifier=span.id_,
+                    parent=span.parent_id,
+                    data=span.start_time,
+                )
+
+                # Add events that belong to this span if requested
+                if include_events:
+                    span_events = [e for e in self.events if e.span_id == span.id_]
+                    for event in sorted(span_events, key=lambda e: e.timestamp):
+                        event_id = f"event-{event.id_}"
+                        tree.create_node(
+                            tag=event.class_name(),
+                            identifier=event_id,
+                            parent=span.id_,
+                            data=event.timestamp,
+                        )
+
+        trees.append(tree)
+        return trees
+
+    def _get_event_trees(self) -> list["Tree"]:
+        """Generate event-centric tree structures with spans as roots and events as children."""
+        try:
+            from treelib.tree import Tree
+        except ImportError as e:
+            raise ImportError(
+                "`treelib` package is missing. Please install it by using "
+                "`pip install treelib`."
+            )
+
+        # Group events by span_id
+        events_by_span: dict[str, list[BaseEvent]] = defaultdict(list)
+        for event in self.events:
+            if event.span_id:
+                events_by_span[event.span_id].append(event)
+
+        trees = []
+        for span_id, span_events in events_by_span.items():
+            tree = Tree()
+            # Create root node for the span
+            tree.create_node(
+                tag=f"{span_id} (SPAN)",
+                identifier=span_id,
+                parent=None,
+                data=min(e.timestamp for e in span_events) if span_events else None,
+            )
+
+            # Add events as children
+            for event in sorted(span_events, key=lambda e: e.timestamp):
+                event_id = f"event-{event.id_}"
+                tree.create_node(
+                    tag=f"{event.class_name()}: {event.id_}",
+                    identifier=event_id,
+                    parent=span_id,
+                    data=event.timestamp,
+                )
+
+            trees.append(tree)
+
+        return trees
+
+    def print_trace_trees(self, include_events: bool = True) -> None:
+        """Display formatted execution traces with chronologically sorted span hierarchies."""
+        trees = self._get_trace_trees(include_events=include_events)
+        for tree in trees:
+            print(tree.show(stdout=False, sorting=True, key=lambda node: node.data))
+            print("")
+
+    def print_event_trees(self) -> None:
+        """Display event sequences grouped by parent span for focused event analysis."""
+        trees = self._get_event_trees()
+        for tree in trees:
+            print(tree.show(stdout=False, sorting=True, key=lambda node: node.data))
+            print("")

novastack_telemetry-1.0.0/novastack_telemetry/span/__init__.py

@@ -0,0 +1,10 @@
+from contextvars import ContextVar
+from typing import Optional
+
+from novastack_telemetry.span.base import BaseSpan, Span
+
+# ContextVar for managing active spans
+_active_span_id: ContextVar[Optional[str]] = ContextVar("_active_span_id", default=None)
+_active_span_id.set(None)
+
+__all__ = ["BaseSpan", "Span", "_active_span_id"]

novastack_telemetry-1.0.0/novastack_telemetry/span/base.py

@@ -0,0 +1,25 @@
+from datetime import datetime
+from typing import Any
+from uuid import uuid4
+
+from pydantic import BaseModel, Field
+
+
+class BaseSpan(BaseModel):
+    """Abstract base class defining the interface for span models."""
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    id_: str = Field(
+        default_factory=lambda: str(uuid4()), description="The Id of span."
+    )
+    parent_id: str | None = Field(default=None, description="The Id of parent span.")
+    metadata: dict[str, Any] = Field(default={})
+
+
+class Span(BaseSpan):
+    """Simple span class."""
+
+    start_time: datetime = Field(default_factory=lambda: datetime.now())
+    end_time: datetime | None = Field(default=None)
+    duration: float = Field(default=0.0, description="Duration of span in seconds.")

novastack_telemetry-1.0.0/pyproject.toml

@@ -0,0 +1,40 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "novastack-telemetry"
+version = "1.0.0"
+description = "Provides the telemetry hooks for observability in Novastack."
+authors = [{ name = "Leonardo Furnielis", email = "leonardofurnielis@outlook.com" }]
+license = { text = "Apache-2.0" }
+readme = "README.md"
+requires-python = ">=3.11,<3.14"
+keywords = ["AI", "LLM", "QA", "RAG", "data", "observability", "retrieval", "semantic-search"]
+classifiers = [
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+    "Topic :: Software Development :: Libraries :: Application Frameworks",
+]
+dependencies = [
+    "pydantic>=2.12.5,<3.0.0",
+    "treelib>=1.8.0,<2.0.0",
+    "wrapt>=2.2.1,<3.0.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=9.0.3,<10.0.0",
+    "pytest-asyncio>=1.3.0,<2.0.0",
+    "pytest-mock>=3.15.1,<4.0.0",
+    "ruff>=0.15.8,<1.0.0",
+]
+
+[project.urls]
+Repository = "https://github.com/novastack-project/novastack/tree/main/novastack-telemetry"
+
+[tool.hatch.build.targets.sdist]
+include = ["novastack_telemetry/"]
+
+[tool.hatch.build.targets.wheel]
+include = ["novastack_telemetry/"]