Merge feat/migrate-to-msgtrace-sdk into init

Complete migration from built-in OpenTelemetry to msgtrace-sdk.

Key changes:
- Integrated msgtrace-sdk v1.1.0 as core telemetry dependency
- Added new attributes to msgtrace-sdk for msgflux compatibility
- Removed all redundant telemetry abstractions
- Clean, direct integration with msgtrace-sdk
- Proper async/sync instrumentation

All tests passing.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

Author: vilsonrodrigues

CLAUDE.md

@@ -25,14 +25,16 @@ The project is being split into specialized libraries:
 - **Integration**: Already listed in pyproject.toml dependencies
 - **Note**: Not to be confused with msgtrace - this is msgspec-ext
 
-#### 2. msgtrace (In Development)
-- **Purpose**: Isolated observability service
+#### 2. msgtrace-sdk (Integrated ✓)
+- **Purpose**: OpenTelemetry-based tracing SDK for AI applications
+- **Status**: Integrated as core dependency (v1.0.0)
 - **Components**:
-  - Backend for trace collection
-  - Frontend dashboard
-  - SDK (msgtrace-sdk) with similar interface to msgflux's current implementation
-- **Status**: Currently being developed separately
-- **Migration Plan**: Will replace built-in OpenTelemetry observability in msgflux
+  - Spans API for creating traces (flows, modules, operations)
+  - MsgTraceAttributes for GenAI semantic conventions
+  - Zero-overhead when disabled
+  - Thread-safe with async support
+- **Integration**: msgflux now uses msgtrace-sdk for all telemetry
+- **Note**: msgflux keeps specialized decorators for tools and agents on top of msgtrace-sdk base
 
 #### 3. txml (Rust Implementation)
 - **Purpose**: XML parsing
@@ -45,13 +47,16 @@ The project is being split into specialized libraries:
 - **Language**: Rust
 - **Status**: New Rust-based implementation
 
-### Current Built-in Features (To Be Migrated)
+### Telemetry Architecture (Updated ✓)
 
-The msgflux codebase still includes:
-- OpenTelemetry-based observability (src/msgflux/telemetry/)
-  - Will be replaced by msgtrace integration
-- Python XML parsers (src/msgflux/dsl/typed_parsers/)
-  - Will be replaced by txml references
+msgflux now uses msgtrace-sdk for telemetry:
+- **Core**: msgtrace-sdk provides Spans API and MsgTraceAttributes
+- **Extensions**: msgflux adds specialized decorators for tools and agents
+- **Configuration**: Environment variables mapped from MSGFLUX_* to MSGTRACE_*
+- **Features Preserved**:
+  - Tool execution tracking (local/remote, protocols like MCP)
+  - Agent telemetry (name, ID, responses)
+  - Detailed argument and response capture
 
 ## Project Structure
 

pyproject.toml

@@ -11,18 +11,17 @@ readme = "README.md"
 authors = [
     { name = "Vilson Rodrigues", email = "vilson@msgflux.com" }
 ]
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 classifiers = [
   "Development Status :: 4 - Beta",
   "Intended Audience :: Developers",
-  "Intended Audience :: Information Technology",  
+  "Intended Audience :: Information Technology",
   "Intended Audience :: Science/Research",
   "License :: OSI Approved :: MIT License",
   "Operating System :: MacOS :: MacOS X",
   "Operating System :: POSIX :: Linux",
   "Operating System :: Microsoft :: Windows",
   "Programming Language :: Python",
-  "Programming Language :: Python :: 3.9",
   "Programming Language :: Python :: 3.10",
   "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
@@ -33,9 +32,7 @@ classifiers = [
 dependencies = [
     "jinja2>=3.1.6",
     "msgspec-ext>=0.1.0",
-    "opentelemetry-api>=1.35.0",
-    "opentelemetry-exporter-otlp>=1.35.0",
-    "opentelemetry-sdk>=1.35.0",
+    "msgtrace-sdk>=1.1.0",
     "tenacity>=8.2.3",
     "typing-extensions>=4.14.1",
     "uvloop>=0.21.0 ; sys_platform != 'win32'",

src/msgflux/__init__.py

@@ -10,7 +10,7 @@
 from .message import Message
 from .models.gateway import ModelGateway
 from .models import Model
-from .telemetry.span import instrument
+from .telemetry import Spans
 from .utils.chat import ChatBlock, ChatML
 from .utils.console import cprint
 from .utils.inspect import get_fn_name
@@ -39,7 +39,7 @@
     "dotdict",
     "get_fn_name",
     "inline",
-    "instrument",
+    "Spans",
     "load",
     "msgspec_dumps",
     "response_cache",

src/msgflux/envs.py

@@ -25,7 +25,7 @@ def set_envs(**kwargs: Any):
 
 class EnvironmentVariables(BaseSettings):
     model_config = SettingsConfigDict(
-        env_file=".msgflux_env",
+        env_file=".env",
         env_prefix="msgflux_",
     )
 
@@ -58,27 +58,16 @@ class EnvironmentVariables(BaseSettings):
     # If set to True, msgflux will trace function calls. Useful for debugging
     trace_function: bool = False
 
-    # if set, msgflux will track executions in nn modules using OTel
-    telemetry_requires_trace: bool = False
-
-    # OTLP endpoint
-    telemetry_otlp_endpoint: str = "http://localhost:8000/api/v1/traces/export"
-
-    # Span exporter type
-    telemetry_span_exporter_type: Literal["console", "otlp"] = "console"
-
-    # Capture state dict
-    telemetry_capture_state_dict: bool = False
-
-    # Capture platform details
-    telemetry_capture_platform: bool = False
-
+    # Telemetry configuration (msgflux-specific)
     # Capture tool call responses
     telemetry_capture_tool_call_responses: bool = True
 
     # Capture agent state, system prompt and tool schemas
     telemetry_capture_agent_prepare_model_execution: bool = False
 
+    # Capture state dict in module execution
+    telemetry_capture_state_dict: bool = False
+
     # State checkpoint, if True, if a module output is in message, skip process
     # if False, reprocess
     state_checkpoint: bool = False
@@ -106,3 +95,30 @@ class EnvironmentVariables(BaseSettings):
 
 
 envs = EnvironmentVariables()
+
+
+def configure_msgtrace_env(
+    enabled: bool = True,
+    otlp_endpoint: str = "http://localhost:8000/api/v1/traces/export",
+    exporter: Literal["console", "otlp"] = "otlp",
+    service_name: str = "msgflux",
+    capture_platform: bool = True,
+):
+    """Configure msgtrace-sdk environment variables.
+
+    This is a convenience function for users to quickly configure
+    msgtrace-sdk. Users can also configure msgtrace-sdk directly
+    using MSGTRACE_* environment variables.
+
+    Args:
+        enabled: Enable telemetry tracking
+        otlp_endpoint: OTLP endpoint URL
+        exporter: Exporter type ("otlp" or "console")
+        service_name: Service name for telemetry
+        capture_platform: Capture platform details (CPU, OS, Python version)
+    """
+    os.environ["MSGTRACE_TELEMETRY_ENABLED"] = "true" if enabled else "false"
+    os.environ["MSGTRACE_OTLP_ENDPOINT"] = otlp_endpoint
+    os.environ["MSGTRACE_EXPORTER"] = exporter
+    os.environ["MSGTRACE_SERVICE_NAME"] = service_name
+    os.environ["MSGTRACE_CAPTURE_PLATFORM"] = "true" if capture_platform else "false"

src/msgflux/nn/functional.py

@@ -8,7 +8,7 @@
 from msgflux.dotdict import dotdict
 from msgflux.logger import logger
 from msgflux.nn.modules.module import get_callable_name
-from msgflux.telemetry.span import instrument
+from msgflux.telemetry import Spans
 
 __all__ = [
     "abackground_task",
@@ -27,7 +27,7 @@
 ]
 
 
-@instrument("map_gather")
+@Spans.instrument()
 def map_gather(
     to_send: Callable,
     *,
@@ -107,7 +107,7 @@ def multiply(x, y=2): return x * y
     return tuple(responses)
 
 
-@instrument("scatter_gather")
+@Spans.instrument()
 def scatter_gather(
     to_send: List[Callable],
     args_list: Optional[List[Tuple[Any, ...]]] = None,
@@ -198,7 +198,7 @@ def farewell(person_name): return f"Goodbye, {person_name}"
     return tuple(responses)
 
 
-@instrument("msg_scatter_gather")
+@Spans.instrument()
 def msg_scatter_gather(
     to_send: List[Callable],
     messages: List[dotdict],
@@ -250,7 +250,7 @@ def msg_scatter_gather(
     return tuple(messages)
 
 
-@instrument("bcast_gather")
+@Spans.instrument()
 def bcast_gather(
     to_send: List[Callable], *args, timeout: Optional[float] = None, **kwargs
 ) -> Tuple[Any, ...]:
@@ -307,7 +307,7 @@ def fail(x): raise ValueError("Intentional error")
     return tuple(responses)
 
 
-@instrument("msg_bcast_gather")
+@Spans.instrument()
 def msg_bcast_gather(
     to_send: List[Callable],
     message: dotdict,
@@ -350,7 +350,7 @@ def msg_bcast_gather(
     return message
 
 
-@instrument("wait_for")
+@Spans.instrument()
 def wait_for(
     to_send: Callable, *args, timeout: Optional[float] = None, **kwargs
 ) -> Any:
@@ -394,7 +394,7 @@ async def f1(x):
         return None
 
 
-@instrument("wait_for_event")
+@Spans.instrument()
 def wait_for_event(event: asyncio.Event) -> None:
     """Waits synchronously for an asyncio.Event to be set.
 
@@ -417,7 +417,7 @@ def wait_for_event(event: asyncio.Event) -> None:
         logger.error(str(e))
 
 
-@instrument("background_task")
+@Spans.instrument()
 def background_task(to_send: Callable, *args, **kwargs) -> None:
     """Executes a task in the background asynchronously without blocking,
     using the AsyncExecutorPool. This function is "fire-and-forget".
@@ -468,7 +468,7 @@ def log_future(future: Future) -> None:
     future.add_done_callback(log_future)
 
 
-@instrument("abackground_task")
+@Spans.ainstrument()
 async def abackground_task(to_send: Callable, *args, **kwargs) -> None:
     """Executes an async task in the background without blocking.
     This is a truly async "fire-and-forget" function.
@@ -518,7 +518,7 @@ async def run_task():
     asyncio.create_task(run_task())
 
 
-@instrument("await_for_event")
+@Spans.ainstrument()
 async def await_for_event(event: asyncio.Event) -> None:
     """Waits asynchronously for an asyncio.Event to be set.
 
@@ -549,7 +549,7 @@ async def setter():
     await event.wait()
 
 
-@instrument("amap_gather")
+@Spans.ainstrument()
 async def amap_gather(
     to_send: Callable,
     *,
@@ -613,7 +613,7 @@ async def amap_gather(
     return tuple(results)
 
 
-@instrument("ascatter_gather")
+@Spans.ainstrument()
 async def ascatter_gather(
     to_send: List[Callable],
     args_list: Optional[List[Tuple[Any, ...]]] = None,
@@ -676,7 +676,7 @@ async def ascatter_gather(
     return tuple(results)
 
 
-@instrument("amsg_bcast_gather")
+@Spans.ainstrument()
 async def amsg_bcast_gather(
     to_send: List[Callable],
     message: dotdict,

src/msgflux/nn/modules/module.py

@@ -37,7 +37,8 @@
 from msgflux.models.model import Model
 from msgflux.models.response import ModelResponse, ModelStreamResponse
 from msgflux.nn.parameter import Parameter
-from msgflux.telemetry.span import spans
+from msgflux.telemetry import Spans
+from msgtrace.sdk import MsgTraceAttributes
 from msgflux.utils.convert import convert_camel_snake_to_title
 from msgflux.utils.encode import encode_data_to_base64
 from msgflux.utils.hooks import RemovableHandle
@@ -401,7 +402,6 @@ def __init__(self, *args, **kwargs) -> None:
         super().__setattr__("_load_state_dict_pre_hooks", OrderedDict())
         super().__setattr__("_load_state_dict_post_hooks", OrderedDict())
         super().__setattr__("_modules", {})
-        super().__setattr__("_spans", spans)
 
         if self.call_super_init:
             super().__init__(*args, **kwargs)
@@ -1382,8 +1382,10 @@ def _execute_with_span(self, module_name_title: str, module_type: str, *args, **
         Returns:
             Module output from forward method
         """
-        with self._spans.init_module(module_name_title, module_type) as span:
+        with Spans.init_module(module_name_title, module_type) as span:
             try:
+                MsgTraceAttributes.set_module_name(module_name_title)
+                MsgTraceAttributes.set_module_type(module_type)
                 result = self.forward(*args, **kwargs)
                 span.set_status(Status(StatusCode.OK))
                 return result
@@ -1406,10 +1408,12 @@ def _call(self, *args, **kwargs):
         current_span = trace.get_current_span()
         # If there is no active span or it is not recording, this is the root module
         if current_span is None or not current_span.is_recording():
-            with self._spans.init_flow(
+            with Spans.init_flow(
                 module_name_title, module_type, encoded_state_dict
             ) as span:
                 try:
+                    MsgTraceAttributes.set_module_name(module_name_title)
+                    MsgTraceAttributes.set_module_type(module_type)
                     module_output = self.forward(*args, **kwargs)
                     span.set_status(Status(StatusCode.OK))
                     return module_output
@@ -1458,8 +1462,10 @@ async def _aexecute_with_span(self, module_name_title: str, module_type: str, *a
         Returns:
             Module output from aforward method
         """
-        async with self._spans.ainit_module(module_name_title, module_type) as span:
+        async with Spans.ainit_module(module_name_title, module_type) as span:
             try:
+                MsgTraceAttributes.set_module_name(module_name_title)
+                MsgTraceAttributes.set_module_type(module_type)
                 result = await self.aforward(*args, **kwargs)
                 span.set_status(Status(StatusCode.OK))
                 return result
@@ -1482,10 +1488,12 @@ async def _acall(self, *args, **kwargs):
         current_span = trace.get_current_span()
         # If there is no active span or it is not recording, this is the root module
         if current_span is None or not current_span.is_recording():
-            async with self._spans.ainit_flow(
+            async with Spans.ainit_flow(
                 module_name_title, module_type, encoded_state_dict
             ) as span:
                 try:
+                    MsgTraceAttributes.set_module_name(module_name_title)
+                    MsgTraceAttributes.set_module_type(module_type)
                     module_output = await self.aforward(*args, **kwargs)
                     span.set_status(Status(StatusCode.OK))
                     return module_output