Upload 11 files

aworld/trace/instrumentation/__init__.py

@@ -0,0 +1,85 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+from abc import ABC, abstractmethod
+from typing import Any, Collection
+from packaging.requirements import Requirement, InvalidRequirement
+from importlib_metadata import version, PackageNotFoundError
+from aworld.logs.util import logger
+
+
+class Instrumentor(ABC):
+    _instance = None
+    _has_instrumented = False
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = object.__new__(cls)
+
+        return cls._instance
+
+    def instrument(self, **kwargs: Any):
+        """
+        Instrument the library.
+        """
+        if self._has_instrumented:
+            logger.warning(
+                f"Instrumentor[{self.__class__.__name__}] has already instrumented, skip")
+            return
+
+        if not self._check_dependency_conflicts():
+            return
+
+        result = self._instrument(**kwargs)
+        self._has_instrumented = True
+        return result
+
+    def uninstrument(self, **kwargs: Any):
+        """
+        Uninstrument the library.
+        """
+        if not self._has_instrumented:
+            logger.warning("Instrumentor has not instrumented, skip")
+            return
+        self._uninstrument(**kwargs)
+        self._has_instrumented = False
+
+    @abstractmethod
+    def _uninstrument(self, **kwargs: Any):
+        """
+        Uninstrument the library.
+        """
+
+    @abstractmethod
+    def _instrument(self, **kwargs: Any):
+        """
+        Instrument the library.
+        """
+
+    def _check_dependency_conflicts(self):
+        dependencies = self.instrumentation_dependencies()
+        for dependence in dependencies:
+            try:
+                requirement = Requirement(dependence)
+            except InvalidRequirement as exc:
+                logger.warning(
+                    f'error parsing dependency, reporting as a conflict: "{dependence}" - {exc}')
+                return False
+            try:
+                dist_version = version(requirement.name)
+            except PackageNotFoundError as exc:
+                logger.warning(
+                    f'dependency not found, reporting as a conflict: "{dependence}" - {exc}')
+                return False
+
+            if requirement.specifier and not requirement.specifier.contains(dist_version):
+                logger.warning(
+                    f'dependency version conflict, reporting as a conflict: requested: "{self.required}" but found: "{self.found}"')
+                return False
+
+        return True
+
+    @abstractmethod
+    def instrumentation_dependencies(self) -> Collection[str]:
+        """
+        Return a list of dependencies that the instrumentation requires.
+        """

aworld/trace/instrumentation/asgi.py

@@ -0,0 +1,203 @@
+from timeit import default_timer
+from typing import Any, Awaitable, Callable
+from functools import wraps
+from aworld.metrics.context_manager import MetricContext
+from aworld.trace.instrumentation.http_util import (
+    collect_request_attributes_asgi,
+    url_disabled,
+    parser_host_port_url_from_asgi
+)
+from aworld.trace.base import Span, TraceProvider, TraceContext, Tracer, SpanType
+from aworld.trace.propagator import get_global_trace_propagator
+from aworld.trace.propagator.carrier import DictCarrier, ListTupleCarrier
+from aworld.metrics.metric import MetricType
+from aworld.metrics.template import MetricTemplate
+from aworld.logs.util import logger
+
+
+def _wrapped_receive(
+    server_span: Span,
+    server_span_name: str,
+    scope: dict[str, Any],
+    receive: Callable[[], Awaitable[dict[str, Any]]],
+    attributes: dict[str],
+    client_request_hook: Callable = None
+):
+
+    @wraps(receive)
+    async def otel_receive():
+        message = await receive()
+        if client_request_hook and callable(client_request_hook):
+            client_request_hook(scope, message)
+
+        server_span.set_attribute("asgi.event.type", message.get("type", ""))
+        return message
+
+    return otel_receive
+
+
+def _wrapped_send(
+    server_span: Span,
+    server_span_name: str,
+    scope: dict[str, Any],
+    send: Callable[[dict[str, Any]], Awaitable[None]],
+    attributes: dict[str],
+    client_response_hook: Callable = None
+):
+    expecting_trailers = False
+
+    @wraps(send)
+    async def otel_send(message: dict[str, Any]):
+        nonlocal expecting_trailers
+
+        status_code = None
+        if message["type"] == "http.response.start":
+            status_code = message["status"]
+        elif message["type"] == "websocket.send":
+            status_code = 200
+
+        # raw_headers = message.get("headers")
+        # if raw_headers:
+        if status_code:
+            server_span.set_attribute(
+                "http.response.status_code", status_code)
+
+        if callable(client_response_hook):
+            client_response_hook(scope, message)
+
+        if message["type"] == "http.response.start":
+            expecting_trailers = message.get("trailers", False)
+
+        propagator = get_global_trace_propagator()
+        if propagator:
+            trace_context = TraceContext(
+                trace_id=server_span.get_trace_id(),
+                span_id=server_span.get_span_id()
+            )
+            propagator.inject(
+                trace_context, DictCarrier(message))
+
+        await send(message)
+
+        if (
+            not expecting_trailers
+            and message["type"] == "http.response.body"
+            and not message.get("more_body", False)
+        ) or (
+            expecting_trailers
+            and message["type"] == "http.response.trailers"
+            and not message.get("more_trailers", False)
+        ):
+            server_span.end()
+
+    return otel_send
+
+
+class TraceMiddleware:
+    """
+    A ASGI Middleware for tracing requests and responses.
+    """
+
+    def __init__(
+            self,
+            app,
+            excluded_urls=None,
+            tracer_provider: TraceProvider = None,
+            tracer: Tracer = None,
+            server_request_hook: Callable = None,
+            client_request_hook: Callable = None,
+            client_response_hook: Callable = None,):
+        self.app = app
+        self.excluded_urls = excluded_urls
+        self.tracer_provider = tracer_provider
+        self.server_request_hook = server_request_hook
+        self.client_request_hook = client_request_hook
+        self.client_response_hook = client_response_hook
+
+        self.tracer: Tracer = (self.tracer_provider.get_tracer(
+            "aworld.trace.instrumentation.asgi"
+        ) if tracer is None else tracer)
+
+        self.duration_histogram = MetricTemplate(
+            type=MetricType.HISTOGRAM,
+            name="asgi_request_duration_histogram",
+            description="Duration of flask HTTP server requests."
+        )
+
+        self.active_requests_counter = MetricTemplate(
+            type=MetricType.UPDOWNCOUNTER,
+            name="asgi_active_request_counter",
+            unit="1",
+            description="Number of active HTTP server requests.",
+        )
+
+    async def __call__(
+        self,
+        scope: dict[str, Any],
+        receive: Callable[[], Awaitable[dict[str, Any]]],
+        send: Callable[[dict[str, Any]], Awaitable[None]],
+    ):
+        start = default_timer()
+        if scope["type"] not in ("http", "websocket"):
+            return await self.app(scope, receive, send)
+
+        _, _, url = parser_host_port_url_from_asgi(scope)
+        if self.excluded_urls and url_disabled(url, self.excluded_urls):
+            return await self.app(scope, receive, send)
+
+        span_name = scope.get("method", "HTTP").strip(
+        ).upper() + "_" + scope.get("path", "").strip()
+
+        attributes = collect_request_attributes_asgi(scope)
+
+        if scope["type"] == "http" and MetricContext.metric_initialized():
+            MetricContext.inc(self.active_requests_counter, 1, attributes)
+
+        trace_context = None
+        propagator = get_global_trace_propagator()
+        if propagator:
+            trace_context = propagator.extract(
+                ListTupleCarrier(scope.get("headers", [])))
+            logger.info(
+                f"asgi extract trace_context: {trace_context}, scope: {scope}")
+        try:
+            with self.tracer.start_as_current_span(
+                span_name, span_type=SpanType.SERVER, trace_context=trace_context, attributes=attributes
+            ) as span:
+
+                if callable(self.server_request_hook):
+                    self.server_request_hook(scope)
+
+                wrappered_receive = _wrapped_receive(
+                    span,
+                    span_name,
+                    scope,
+                    receive,
+                    attributes,
+                    self.client_request_hook
+                )
+                wrappered_send = _wrapped_send(
+                    span,
+                    span_name,
+                    scope,
+                    send,
+                    attributes,
+                    self.client_response_hook
+                )
+
+                await self.app(scope, wrappered_receive, wrappered_send)
+        finally:
+            if scope["type"] == "http":
+                duration_s = default_timer() - start
+
+                if MetricContext.metric_initialized():
+                    MetricContext.histogram_record(
+                        self.duration_histogram,
+                        duration_s,
+                        attributes
+                    )
+                    MetricContext.inc(
+                        self.active_requests_counter, -1, attributes)
+
+            if span.is_recording():
+                span.end()

aworld/trace/instrumentation/eventbus.py

@@ -0,0 +1,115 @@
+import wrapt
+from typing import Any, Collection
+from aworld.trace.instrumentation import Instrumentor
+from aworld.trace.base import Tracer, get_tracer_provider_silent, TraceContext
+from aworld.trace.propagator import get_global_trace_propagator, get_global_trace_context
+from aworld.trace.propagator.carrier import DictCarrier
+from aworld.logs.util import logger
+
+
+def _emit_message_class_wrapper(tracer: Tracer):
+    async def awrapper(wrapped, instance, args, kwargs):
+        from aworld.core.event.base import Message
+        try:
+            event = args[0] if len(args) > 0 else kwargs.get("event")
+            propagator = get_global_trace_propagator()
+            trace_provider = get_tracer_provider_silent()
+            if trace_provider and propagator and event and isinstance(event, Message):
+                if not event.headers:
+                    event.headers = {}
+                current_span = trace_provider.get_current_span()
+                if current_span:
+                    trace_context = TraceContext(
+                        trace_id=current_span.get_trace_id(), span_id=current_span.get_span_id())
+                    propagator.inject(trace_context=trace_context,
+                                      carrier=DictCarrier(event.headers))
+                logger.info(
+                    f"EventManager emit_message trace propagate, event.headers={event.headers}")
+        except Exception as e:
+            logger.error(
+                f"EventManager emit_message trace propagate exception: {e}")
+        return await wrapped(*args, **kwargs)
+    return awrapper
+
+
+def _emit_message_instance_wrapper(tracer: Tracer):
+
+    @wrapt.decorator
+    async def awrapper(wrapped, instance, args, kwargs):
+        wrapper = _emit_message_class_wrapper(tracer)
+        return await wrapper(wrapped, instance, args, kwargs)
+
+    return awrapper
+
+
+def _consume_class_wrapper(tracer: Tracer):
+    async def awrapper(wrapped, instance, args, kwargs):
+        from aworld.core.event.base import Message
+        event = await wrapped(*args, **kwargs)
+        try:
+            propagator = get_global_trace_propagator()
+            if propagator and event and isinstance(event, Message) and event.headers:
+                trace_context = propagator.extract(DictCarrier(event.headers))
+                logger.info(
+                    f"extract trace_context from event: {trace_context}")
+                if trace_context:
+                    get_global_trace_context().set(trace_context)
+        except Exception as e:
+            logger.error(
+                f"EventManager consume trace propagate exception: {e}")
+        return event
+    return awrapper
+
+
+def _consume_instance_wrapper(tracer: Tracer):
+
+    @wrapt.decorator
+    async def awrapper(wrapped, instance, args, kwargs):
+        wrapper = _consume_class_wrapper(tracer)
+        return await wrapper(wrapped, instance, args, kwargs)
+
+    return awrapper
+
+
+class EventBusInstrumentor(Instrumentor):
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ()
+
+    def _uninstrument(self, **kwargs: Any):
+        pass
+
+    def _instrument(self, **kwargs: Any):
+        tracer_provider = get_tracer_provider_silent()
+        if not tracer_provider:
+            return
+        tracer = tracer_provider.get_tracer(
+            "aworld.trace.instrumentation.eventbus")
+
+        wrapt.wrap_function_wrapper(
+            "aworld.events.manager",
+            "EventManager.emit_message",
+            _emit_message_class_wrapper(tracer=tracer)
+        )
+
+        wrapt.wrap_function_wrapper(
+            "aworld.events.manager",
+            "EventManager.consume",
+            _consume_class_wrapper(tracer=tracer)
+        )
+
+
+def wrap_event_manager(manager: 'aworld.events.manager.EventManager'):
+    tracer_provider = get_tracer_provider_silent()
+    if not tracer_provider:
+        return manager
+    tracer = tracer_provider.get_tracer(
+        "aworld.trace.instrumentation.eventbus")
+
+    emit_wrapper = _emit_message_instance_wrapper(tracer)
+    consume_wrapper = _consume_instance_wrapper(tracer)
+
+    manager.emit_message = emit_wrapper(manager.emit_message)
+    manager.consume = consume_wrapper(manager.consume)
+
+    return manager

aworld/trace/instrumentation/fastapi.py

@@ -0,0 +1,107 @@
+from typing import Any, Callable
+from .asgi import TraceMiddleware
+from aworld.trace.instrumentation import Instrumentor
+from aworld.trace.base import TraceProvider, get_tracer_provider
+from aworld.trace.instrumentation.http_util import (
+    get_excluded_urls,
+    parse_excluded_urls,
+)
+from aworld.utils.import_package import import_packages
+from aworld.logs.util import logger
+
+import_packages(['fastapi'])  # noqa
+import fastapi  # noqa
+
+
+class _InstrumentedFastAPI(fastapi.FastAPI):
+    """Instrumented FastAPI class."""
+    _tracer_provider: TraceProvider = None
+    _excluded_urls: list[str] = None
+    _server_request_hook: Callable = None
+    _client_request_hook: Callable = None
+    _client_response_hook: Callable = None
+    _instrumented_fastapi_apps = set()
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        tracer = self._tracer_provider.get_tracer(
+            "aworld.trace.instrumentation.fastapi")
+
+        self.add_middleware(
+            TraceMiddleware,
+            tracer=tracer,
+            excluded_urls=self._excluded_urls,
+            server_request_hook=self._server_request_hook,
+            client_request_hook=self._client_request_hook,
+            client_response_hook=self._client_response_hook
+        )
+
+        self._is_instrumented_by_trace = True
+        self._instrumented_fastapi_apps.add(self)
+
+    def __del__(self):
+        if self in self._instrumented_fastapi_apps:
+            self._instrumented_fastapi_apps.remove(self)
+
+
+class FastAPIInstrumentor(Instrumentor):
+    """FastAPI Instrumentor."""
+    _original_fastapi = None
+
+    @staticmethod
+    def uninstrument_app(app: fastapi.FastAPI):
+        app.user_middleware = [
+            x
+            for x in app.user_middleware
+            if x.cls is not TraceMiddleware
+        ]
+        app.middleware_stack = app.build_middleware_stack()
+        app._is_instrumented_by_trace = False
+
+    def instrumentation_dependencies(self) -> dict[str, Any]:
+        return {"fastapi": fastapi}
+
+    def _instrument(self, **kwargs):
+        self._original_fastapi = fastapi.FastAPI
+        _InstrumentedFastAPI._tracer_provider = kwargs.get("tracer_provider")
+        _InstrumentedFastAPI._server_request_hook = kwargs.get(
+            "server_request_hook"
+        )
+        _InstrumentedFastAPI._client_request_hook = kwargs.get(
+            "client_request_hook"
+        )
+        _InstrumentedFastAPI._client_response_hook = kwargs.get(
+            "client_response_hook"
+        )
+        excluded_urls = kwargs.get("excluded_urls")
+        _InstrumentedFastAPI._excluded_urls = (
+            get_excluded_urls("FASTAPI")
+            if excluded_urls is None
+            else parse_excluded_urls(excluded_urls)
+        )
+        fastapi.FastAPI = _InstrumentedFastAPI
+
+    def _uninstrument(self, **kwargs):
+        for app in _InstrumentedFastAPI._instrumented_fastapi_apps:
+            self.uninstrument_app(app)
+        _InstrumentedFastAPI._instrumented_fastapi_apps.clear()
+        fastapi.FastAPI = self._original_fastapi
+
+
+def instrument_fastapi(excluded_urls: str = None,
+                       server_request_hook: Callable = None,
+                       client_request_hook: Callable = None,
+                       client_response_hook: Callable = None,
+                       tracer_provider: TraceProvider = None,
+                       **kwargs: Any,
+                       ):
+    kwargs.update({
+        "excluded_urls": excluded_urls,
+        "server_request_hook": server_request_hook,
+        "client_request_hook": client_request_hook,
+        "client_response_hook": client_response_hook,
+        "tracer_provider": tracer_provider or get_tracer_provider(),
+    })
+    FastAPIInstrumentor().instrument(**kwargs)
+    logger.info("FastAPI instrumented.")

aworld/trace/instrumentation/flask.py

@@ -0,0 +1,279 @@
+import flask
+import weakref
+from typing import Any, Callable, Collection
+from time import time_ns
+from timeit import default_timer
+from importlib_metadata import version
+from packaging import version as package_version
+from aworld.trace.instrumentation import Instrumentor
+from aworld.trace.base import Span, TraceProvider, TraceContext, Tracer, SpanType, get_tracer_provider
+from aworld.metrics.metric import MetricType
+from aworld.metrics.template import MetricTemplate
+from aworld.logs.util import logger
+from aworld.trace.instrumentation.http_util import (
+    collect_request_attributes,
+    url_disabled,
+    get_excluded_urls,
+    parse_excluded_urls,
+    HTTP_ROUTE
+)
+from aworld.trace.propagator import get_global_trace_propagator
+from aworld.metrics.context_manager import MetricContext
+from aworld.trace.propagator.carrier import ListTupleCarrier, DictCarrier
+
+_ENVIRON_STARTTIME_KEY = "aworld-flask.starttime_key"
+_ENVIRON_SPAN_KEY = "aworld-flask.span_key"
+_ENVIRON_REQCTX_REF_KEY = "aworld-flask.reqctx_ref_key"
+
+flask_version = version("flask")
+if package_version.parse(flask_version) >= package_version.parse("2.2.0"):
+
+    def _request_ctx_ref() -> weakref.ReferenceType:
+        return weakref.ref(flask.globals.request_ctx._get_current_object())
+
+else:
+
+    def _request_ctx_ref() -> weakref.ReferenceType:
+        return weakref.ref(flask._request_ctx_stack.top)
+
+
+def _rewrapped_app(
+    wsgi_app,
+    active_requests_counter,
+    duration_histogram,
+    response_hook=None,
+    excluded_urls=None,
+):
+    def _wrapped_app(wrapped_app_environ, start_response):
+        # We want to measure the time for route matching, etc.
+        # In theory, we could start the span here and use
+        # update_name later but that API is "highly discouraged" so
+        # we better avoid it.
+        wrapped_app_environ[_ENVIRON_STARTTIME_KEY] = time_ns()
+        start = default_timer()
+        attributes = collect_request_attributes(wrapped_app_environ)
+
+        if MetricContext.metric_initialized():
+            MetricContext.inc(active_requests_counter, 1, attributes)
+
+        request_route = None
+
+        def _start_response(status, response_headers, *args, **kwargs):
+            if flask.request and (
+                excluded_urls is None
+                or not url_disabled(flask.request.url, excluded_urls)
+            ):
+                nonlocal request_route
+                request_route = flask.request.url_rule
+
+                span: Span = flask.request.environ.get(_ENVIRON_SPAN_KEY)
+
+                propagator = get_global_trace_propagator()
+                if propagator and span:
+                    trace_context = TraceContext(
+                        trace_id=span.get_trace_id(),
+                        span_id=span.get_span_id()
+                    )
+                    propagator.inject(
+                        trace_context, ListTupleCarrier(response_headers))
+
+                if span and span.is_recording():
+                    status_code_str, _ = status.split(" ", 1)
+                    try:
+                        status_code = int(status_code_str)
+                    except ValueError:
+                        status_code = -1
+
+                    span.set_attribute(
+                        "http.response.status_code", status_code)
+                    span.set_attributes(attributes)
+
+                if response_hook is not None:
+                    response_hook(span, status, response_headers)
+            return start_response(status, response_headers, *args, **kwargs)
+
+        result = wsgi_app(wrapped_app_environ, _start_response)
+        duration_s = default_timer() - start
+
+        if MetricContext.metric_initialized():
+            MetricContext.histogram_record(
+                duration_histogram,
+                duration_s,
+                attributes
+            )
+            MetricContext.dec(active_requests_counter, 1, attributes)
+        return result
+
+    return _wrapped_app
+
+
+def _wrapped_before_request(
+    request_hook=None,
+    tracer: Tracer = None,
+    excluded_urls=None
+):
+    def _before_request():
+        if excluded_urls and url_disabled(flask.request.url, excluded_urls):
+            return
+        flask_request_environ = flask.request.environ
+        logger.info(
+            f"_wrapped_before_request flask_request_environ={flask_request_environ}")
+
+        attributes = collect_request_attributes(flask_request_environ)
+
+        if flask.request.url_rule:
+            # For 404 that result from no route found, etc, we
+            # don't have a url_rule.
+            attributes[HTTP_ROUTE] = flask.request.url_rule.rule
+            span_name = f"HTTP {flask.request.url_rule.rule}"
+        else:
+            span_name = f"HTTP {flask.request.url}"
+
+        propagator = get_global_trace_propagator()
+        trace_context = None
+        if propagator:
+            trace_context = propagator.extract(
+                DictCarrier(flask_request_environ))
+
+        logger.info(f"_wrapped_before_request trace_context={trace_context}")
+
+        span = tracer.start_span(
+            span_name,
+            SpanType.SERVER,
+            attributes=attributes,
+            start_time=flask_request_environ.get(_ENVIRON_STARTTIME_KEY),
+            trace_context=trace_context
+        )
+
+        if request_hook:
+            request_hook(span, flask_request_environ)
+
+        flask_request_environ[_ENVIRON_SPAN_KEY] = span
+        flask_request_environ[_ENVIRON_REQCTX_REF_KEY] = _request_ctx_ref()
+
+    return _before_request
+
+
+def _wrapped_teardown_request(
+    excluded_urls=None,
+):
+    def _teardown_request(exc):
+        if excluded_urls and url_disabled(flask.request.url, excluded_urls):
+            return
+
+        span: Span = flask.request.environ.get(_ENVIRON_SPAN_KEY)
+
+        original_reqctx_ref = flask.request.environ.get(
+            _ENVIRON_REQCTX_REF_KEY
+        )
+        current_reqctx_ref = _request_ctx_ref()
+        if not span or original_reqctx_ref != current_reqctx_ref:
+            return
+        if exc is None:
+            span.end()
+        else:
+            span.record_exception(exc)
+            span.end()
+
+    return _teardown_request
+
+
+class _InstrumentedFlask(flask.Flask):
+    _excluded_urls = None
+    _tracer_provider: TraceProvider = None
+    _request_hook = None
+    _response_hook = None
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        tracer = self._tracer_provider.get_tracer(
+            "aworld.trace.instrumentation.flask")
+
+        duration_histogram = MetricTemplate(
+            type=MetricType.HISTOGRAM,
+            name="flask_request_duration_histogram",
+            description="Duration of flask HTTP server requests."
+        )
+
+        active_requests_counter = MetricTemplate(
+            type=MetricType.UPDOWNCOUNTER,
+            name="flask_active_request_counter",
+            unit="1",
+            description="Number of active HTTP server requests.",
+        )
+
+        self.wsgi_app = _rewrapped_app(
+            self.wsgi_app,
+            active_requests_counter,
+            duration_histogram,
+            _InstrumentedFlask._response_hook,
+            excluded_urls=_InstrumentedFlask._excluded_urls
+        )
+
+        _before_request = _wrapped_before_request(
+            _InstrumentedFlask._request_hook,
+            tracer,
+            excluded_urls=_InstrumentedFlask._excluded_urls
+        )
+        self._before_request = _before_request
+        self.before_request(_before_request)
+
+        _teardown_request = _wrapped_teardown_request(
+            excluded_urls=_InstrumentedFlask._excluded_urls,
+        )
+        self.teardown_request(_teardown_request)
+
+
+class FlaskInstrumentor(Instrumentor):
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ("flask >= 1.0",)
+
+    def _instrument(self, **kwargs: Any):
+        logger.info("Flask _instrument entered.")
+        self._original_flask = flask.Flask
+        request_hook = kwargs.get("request_hook")
+        response_hook = kwargs.get("response_hook")
+        if callable(request_hook):
+            _InstrumentedFlask._request_hook = request_hook
+        if callable(response_hook):
+            _InstrumentedFlask._response_hook = response_hook
+        tracer_provider = kwargs.get("tracer_provider")
+        _InstrumentedFlask._tracer_provider = tracer_provider
+        excluded_urls = kwargs.get("excluded_urls")
+        _InstrumentedFlask._excluded_urls = (
+            get_excluded_urls("FLASK")
+            if excluded_urls is None
+            else parse_excluded_urls(excluded_urls)
+        )
+        flask.Flask = _InstrumentedFlask
+        logger.info("Flask _instrument exited.")
+
+    def _uninstrument(self, **kwargs):
+        flask.Flask = self._original_flask
+
+
+def instrument_flask(excluded_urls: str = None,
+                     request_hook: Callable = None,
+                     response_hook: Callable = None,
+                     tracer_provider: TraceProvider = None,
+                     **kwargs: Any,
+                     ):
+    """
+    Instrument the Flask application.
+    Args:
+        excluded_urls (str): A comma separated list of URLs to be excluded from instrumentation.
+        request_hook (Callable): A function to be called before a request is processed.
+        response_hook (Callable): A function to be called after a request is processed.
+        tracer_provider (TraceProvider): The trace provider to use.
+    """
+    all_kwargs = {
+        "excluded_urls": excluded_urls,
+        "request_hook": request_hook,
+        "response_hook": response_hook,
+        "tracer_provider": tracer_provider or get_tracer_provider(),
+        **kwargs
+    }
+    FlaskInstrumentor().instrument(**all_kwargs)
+    logger.info("Flask instrumented.")

aworld/trace/instrumentation/http_util.py

@@ -0,0 +1,193 @@
+import os
+from re import compile as re_compile
+from re import search
+from typing import Final, Iterable, Any
+from urllib.parse import urlparse, urlunparse, unquote
+from wsgiref.types import WSGIEnvironment
+from requests.models import PreparedRequest
+
+HTTP_REQUEST_METHOD: Final = "http.request.method"
+HTTP_FLAVOR: Final = "http.flavor"
+HTTP_HOST: Final = "http.host"
+HTTP_SCHEME: Final = "http.scheme"
+HTTP_USER_AGENT: Final = "http.user_agent"
+HTTP_SERVER_NAME: Final = "http.server_name"
+SERVER_ADDRESS: Final = "server.address"
+SERVER_PORT: Final = "server.port"
+URL_PATH: Final = "url.path"
+URL_QUERY: Final = "url.query"
+CLIENT_ADDRESS: Final = "client.address"
+CLIENT_PORT: Final = "client.port"
+URL_FULL: Final = "url.full"
+
+HTTP_REQUEST_BODY_SIZE: Final = "http.request.body.size"
+HTTP_REQUEST_HEADER: Final = "http.request.header"
+HTTP_REQUEST_SIZE: Final = "http.request.size"
+HTTP_RESPONSE_BODY_SIZE: Final = "http.response.body.size"
+HTTP_RESPONSE_HEADER: Final = "http.response.header"
+HTTP_RESPONSE_SIZE: Final = "http.response.size"
+HTTP_RESPONSE_STATUS_CODE: Final = "http.response.status_code"
+HTTP_ROUTE = "http.route"
+
+
+def collect_request_attributes(environ: WSGIEnvironment):
+
+    attributes: dict[str] = {}
+
+    request_method = environ.get("REQUEST_METHOD", "")
+    request_method = request_method.upper()
+    attributes[HTTP_REQUEST_METHOD] = request_method
+    attributes[HTTP_FLAVOR] = environ.get("SERVER_PROTOCOL", "")
+    attributes[HTTP_SCHEME] = environ.get("wsgi.url_scheme", "")
+    attributes[HTTP_SERVER_NAME] = environ.get("SERVER_NAME", "")
+    attributes[HTTP_HOST] = environ.get("HTTP_HOST", "")
+    host_port = environ.get("SERVER_PORT")
+    if host_port:
+        attributes[SERVER_PORT] = host_port
+    target = environ.get("RAW_URI")
+    if target is None:
+        target = environ.get("REQUEST_URI")
+    if target:
+        path, query = _parse_url_query(target)
+        attributes[URL_PATH] = path
+        attributes[URL_QUERY] = query
+    remote_addr = environ.get("REMOTE_ADDR", "")
+    attributes[CLIENT_ADDRESS] = remote_addr
+    attributes[CLIENT_PORT] = environ.get("REMOTE_PORT", "")
+    remote_host = environ.get("REMOTE_HOST")
+    if remote_host and remote_host != remote_addr:
+        attributes[CLIENT_ADDRESS] = remote_host
+    attributes[HTTP_USER_AGENT] = environ.get("HTTP_USER_AGENT", "")
+    return attributes
+
+
+def collect_attributes_from_request(request: PreparedRequest) -> dict[str]:
+    attributes: dict[str] = {}
+
+    url = remove_url_credentials(request.url)
+    attributes[HTTP_REQUEST_METHOD] = request.method
+    attributes[URL_FULL] = url
+    parsed_url = urlparse(url)
+    if parsed_url.scheme:
+        attributes[HTTP_SCHEME] = parsed_url.scheme
+    if parsed_url.hostname:
+        attributes[HTTP_HOST] = parsed_url.hostname
+    if parsed_url.port:
+        attributes[SERVER_PORT] = parsed_url.port
+    return attributes
+
+
+def url_disabled(url: str, excluded_urls: Iterable[str]) -> bool:
+    """
+    Check if the url is disabled.
+    Args:
+        url: The url to check.
+        excluded_urls: The excluded urls.
+    Returns:
+        True if the url is disabled, False otherwise.
+    """
+    if excluded_urls is None:
+        return False
+    regex = re_compile("|".join(excluded_urls))
+    return search(regex, url)
+
+
+def get_excluded_urls(instrumentation: str) -> list[str]:
+    """
+    Get the excluded urls.
+    Args:
+        instrumentation: The instrumentation to get the excluded urls for.
+    Returns:
+        The excluded urls.
+    """
+
+    excluded_urls = os.environ.get(f"{instrumentation}_EXCLUDED_URLS")
+
+    return parse_excluded_urls(excluded_urls)
+
+
+def parse_excluded_urls(excluded_urls: str) -> list[str]:
+    """
+    Parse the excluded urls.
+    Args:
+        excluded_urls: The excluded urls.
+    Returns:
+        The excluded urls.
+    """
+    if excluded_urls:
+        excluded_url_list = [
+            excluded_url.strip() for excluded_url in excluded_urls.split(",")
+        ]
+    else:
+        excluded_url_list = []
+
+    return excluded_url_list
+
+
+def remove_url_credentials(url: str) -> str:
+    """Given a string url, remove the username and password only if it is a valid url"""
+
+    try:
+        parsed = urlparse(url)
+        if all([parsed.scheme, parsed.netloc]):  # checks for valid url
+            parsed_url = urlparse(url)
+            _, _, netloc = parsed.netloc.rpartition("@")
+            return urlunparse(
+                (
+                    parsed_url.scheme,
+                    netloc,
+                    parsed_url.path,
+                    parsed_url.params,
+                    parsed_url.query,
+                    parsed_url.fragment,
+                )
+            )
+    except ValueError:  # an unparsable url was passed
+        pass
+    return url
+
+
+def parser_host_port_url_from_asgi(scope: dict[str, Any]):
+    """Returns (host, port, full_url) tuple."""
+    server = scope.get("server") or ["0.0.0.0", 80]
+    port = server[1]
+    server_host = server[0] + (":" + str(port) if str(port) != "80" else "")
+    full_path = scope.get("path", "")
+    http_url = scope.get("scheme", "http") + "://" + server_host + full_path
+    return server_host, port, http_url
+
+
+def collect_request_attributes_asgi(scope: dict[str, Any]):
+    attributes: dict[str] = {}
+    server_host, port, http_url = parser_host_port_url_from_asgi(scope)
+    query_string = scope.get("query_string")
+    if query_string and http_url:
+        if isinstance(query_string, bytes):
+            query_string = query_string.decode("utf8")
+        http_url += "?" + unquote(query_string)
+    attributes[HTTP_REQUEST_METHOD] = scope.get("method", "")
+    attributes[HTTP_FLAVOR] = scope.get("http_version", "")
+    attributes[HTTP_SCHEME] = scope.get("scheme", "")
+    attributes[HTTP_HOST] = server_host
+    attributes[SERVER_PORT] = port
+    attributes[URL_FULL] = remove_url_credentials(http_url)
+    attributes[URL_PATH] = scope.get("path", "")
+    header = scope.get("headers")
+    if header:
+        for key, value in header:
+            if key == b"user-agent":
+                attributes[HTTP_USER_AGENT] = value.decode("utf8")
+
+    client = scope.get("client")
+    if client:
+        attributes[CLIENT_ADDRESS] = client[0]
+        attributes[CLIENT_PORT] = client[1]
+
+    return attributes
+
+
+def _parse_url_query(url: str):
+    parsed_url = urlparse(url)
+    path = parsed_url.path
+    query_params = parsed_url.query
+    return path, query_params

aworld/trace/instrumentation/llm_metrics.py

@@ -0,0 +1,118 @@
+from aworld.metrics.context_manager import MetricContext
+from aworld.metrics.template import MetricTemplate
+from aworld.metrics.metric import MetricType
+
+tokens_usage_histogram = MetricTemplate(
+    type=MetricType.HISTOGRAM,
+    name="llm_token_usage",
+    unit="token",
+    description="Measures number of input and output tokens used"
+)
+
+chat_choice_counter = MetricTemplate(
+    type=MetricType.COUNTER,
+    name="llm_generation_choice_counter",
+    unit="choice",
+    description="Number of choices returned by chat completions call"
+)
+
+duration_histogram = MetricTemplate(
+    type=MetricType.HISTOGRAM,
+    name="llm_chat_duration",
+    unit="s",
+    description="AI chat duration",
+)
+
+chat_exception_counter = MetricTemplate(
+    type=MetricType.COUNTER,
+    name="llm_chat_exception_counter",
+    unit="time",
+    description="Number of exceptions occurred during chat completions",
+)
+
+streaming_time_to_first_token_histogram = MetricTemplate(
+    type=MetricType.HISTOGRAM,
+    name="llm_streaming_time_to_first_token",
+    unit="s",
+    description="Time to first token in streaming chat completions",
+)
+streaming_time_to_generate_histogram = MetricTemplate(
+    type=MetricType.HISTOGRAM,
+    name="streaming_time_to_generate",
+    unit="s",
+    description="Time between first token and completion in streaming chat completions",
+)
+
+
+def record_exception_metric(exception, duration):
+    '''
+    record chat exception to metrics
+    '''
+    if MetricContext.metric_initialized():
+        labels = {
+            "error.type": exception.__class__.__name__,
+        }
+        if duration_histogram:
+            MetricContext.histogram_record(
+                duration_histogram, duration, labels=labels)
+        if chat_exception_counter:
+            MetricContext.count(
+                chat_exception_counter, 1, labels=labels)
+
+
+def record_streaming_time_to_first_token(duration, labels):
+    '''
+    Record duration of start time to first token in stream. 
+    '''
+    if MetricContext.metric_initialized():
+        MetricContext.histogram_record(
+            streaming_time_to_first_token_histogram, duration, labels=labels)
+
+
+def record_streaming_time_to_generate(first_token_to_generate_duration, labels):
+    '''
+    Record duration the first token to response to generation
+    '''
+    if MetricContext.metric_initialized():
+        MetricContext.histogram_record(
+            streaming_time_to_generate_histogram, first_token_to_generate_duration, labels=labels)
+
+
+def record_chat_response_metric(attributes,
+                                prompt_tokens,
+                                completion_tokens,
+                                duration,
+                                choices=None
+                                ):
+    '''
+    Record chat response to metrics
+    '''
+    if MetricContext.metric_initialized():
+        if prompt_tokens and tokens_usage_histogram:
+            labels = {
+                **attributes,
+                "llm.prompt_usage_type": "prompt_tokens"
+            }
+            MetricContext.histogram_record(
+                tokens_usage_histogram, prompt_tokens, labels=labels)
+        if completion_tokens and tokens_usage_histogram:
+            labels = {
+                **attributes,
+                "llm.prompt_usage_type": "completion_tokens"
+            }
+            MetricContext.histogram_record(
+                tokens_usage_histogram, completion_tokens, labels=labels)
+        if duration and duration_histogram:
+            MetricContext.histogram_record(
+                duration_histogram, duration, labels=attributes)
+        if choices and chat_choice_counter:
+            MetricContext.count(chat_choice_counter,
+                                len(choices), labels=attributes)
+            for choice in choices:
+                if choice.get("finish_reason"):
+                    finish_reason_attr = {
+                        **attributes,
+                        "llm.finish_reason": choice.get("finish_reason")
+                    }
+                    MetricContext.count(
+                        chat_choice_counter, 1, labels=finish_reason_attr)

aworld/trace/instrumentation/requests.py

@@ -0,0 +1,213 @@
+from aworld.logs.util import logger
+from aworld.trace.instrumentation.http_util import (
+    collect_attributes_from_request,
+    url_disabled,
+    get_excluded_urls,
+    parse_excluded_urls,
+    HTTP_RESPONSE_STATUS_CODE,
+    HTTP_FLAVOR
+)
+from aworld.metrics.context_manager import MetricContext
+from aworld.metrics.template import MetricTemplate
+from aworld.metrics.metric import MetricType
+from aworld.trace.instrumentation import Instrumentor
+from aworld.trace.propagator.carrier import DictCarrier
+import functools
+from timeit import default_timer
+from requests import sessions
+from requests.models import PreparedRequest, Response
+from requests.structures import CaseInsensitiveDict
+from typing import Collection, Any, Callable
+from aworld.trace.base import TraceProvider, TraceContext, Tracer, SpanType, get_tracer_provider
+from aworld.trace.propagator import get_global_trace_propagator
+
+
+def _wrapped_send(
+    tracer: Tracer = None,
+    excluded_urls=None,
+    request_hook: Callable = None,
+    response_hook: Callable = None,
+    duration_histogram: MetricTemplate = None
+):
+
+    oringinal_send = sessions.Session.send
+
+    @functools.wraps(oringinal_send)
+    def instrumented_send(
+        self: sessions.Session, request: PreparedRequest, **kwargs: Any
+    ):
+        if excluded_urls and url_disabled(request.url, excluded_urls):
+            return oringinal_send(self, request, **kwargs)
+
+        def get_or_create_headers():
+            request.headers = (
+                request.headers
+                if request.headers is not None
+                else CaseInsensitiveDict()
+            )
+            return request.headers
+
+        method = request.method
+        if method is None:
+            method = "HTTP"
+        span_name = method.upper()
+
+        span_attributes = collect_attributes_from_request(request)
+        with tracer.start_as_current_span(
+            span_name, span_type=SpanType.CLIENT, attributes=span_attributes
+        ) as span:
+            exception = None
+            if callable(request_hook):
+                request_hook(span, request)
+
+            headers = get_or_create_headers()
+
+            trace_context = TraceContext(
+                trace_id=span.get_trace_id(),
+                span_id=span.get_span_id(),
+            )
+            propagator = get_global_trace_propagator()
+            if propagator:
+                propagator.inject(trace_context, DictCarrier(headers))
+
+            start_time = default_timer()
+            try:
+                logger.info("Sending headers: %s", request.headers)
+                result = oringinal_send(
+                    self, request, **kwargs
+                )  # *** PROCEED
+            except Exception as exc:  # pylint: disable=W0703
+                exception = exc
+                result = getattr(exc, "response", None)
+            finally:
+                elapsed_time = max(default_timer() - start_time, 0)
+
+            if isinstance(result, Response):
+                span_attributes = {}
+                span_attributes[HTTP_RESPONSE_STATUS_CODE] = result.status_code
+
+                if result.raw is not None:
+                    version = getattr(result.raw, "version", None)
+                    if version:
+                        # Only HTTP/1 is supported by requests
+                        version_text = "1.1" if version == 11 else "1.0"
+                        span_attributes[HTTP_FLAVOR] = version_text
+                span.set_attributes(span_attributes)
+
+                if callable(response_hook):
+                    response_hook(span, request, result)
+
+            if exception is not None:
+                span.record_exception(exception)
+
+            if duration_histogram is not None and MetricContext.metric_initialized():
+                MetricContext.histogram_record(
+                    duration_histogram,
+                    elapsed_time,
+                    span_attributes
+                )
+
+            if exception is not None:
+                raise exception.with_traceback(exception.__traceback__)
+
+        return result
+
+    return instrumented_send
+
+
+class _InstrumentedSession(sessions.Session):
+    """
+    An instrumented requests.Session class.
+    """
+    _excluded_urls = None
+    _tracer_provider: TraceProvider = None
+    _request_hook = None
+    _response_hook = None
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        tracer = self._tracer_provider.get_tracer(
+            "aworld.trace.instrumentation.requests")
+        excluded_urls = kwargs.get("excluded_urls")
+
+        duration_histogram = MetricTemplate(
+            type=MetricType.HISTOGRAM,
+            name="client_request_duration_histogram",
+            unit="s",
+            description="Duration of  HTTP client requests."
+        )
+        self.send = functools.partial(_wrapped_send(
+            tracer=tracer,
+            excluded_urls=excluded_urls,
+            request_hook=self._request_hook,
+            response_hook=self._response_hook,
+            duration_histogram=duration_histogram
+        ), self)
+
+
+class RequestsInstrumentor(Instrumentor):
+    """
+    An instrumentor for the requests module.
+    """
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["requests"]
+
+    def _instrument(self, **kwargs):
+        """
+        Instruments the requests module.
+        """
+        logger.info("requests _instrument entered.")
+        self._original_session = sessions.Session
+        request_hook = kwargs.get("request_hook")
+        response_hook = kwargs.get("response_hook")
+        if callable(request_hook):
+            _InstrumentedSession._request_hook = request_hook
+        if callable(response_hook):
+            _InstrumentedSession._response_hook = response_hook
+        tracer_provider = kwargs.get("tracer_provider")
+        _InstrumentedSession._tracer_provider = tracer_provider
+        excluded_urls = kwargs.get("excluded_urls")
+        _InstrumentedSession._excluded_urls = (
+            get_excluded_urls("FLASK")
+            if excluded_urls is None
+            else parse_excluded_urls(excluded_urls)
+        )
+        sessions.Session = _InstrumentedSession
+        logger.info("requests _instrument exited.")
+
+    def _uninstrument(self, **kwargs):
+        """
+        Uninstruments the requests module.
+        """
+        sessions.Session = self._original_session
+
+
+def instrument_requests(excluded_urls: str = None,
+                        request_hook: Callable = None,
+                        response_hook: Callable = None,
+                        tracer_provider: TraceProvider = None,
+                        **kwargs: Any,
+                        ):
+    """
+    Instruments the requests module.
+    Args:
+        excluded_urls: A comma separated list of URLs to exclude from tracing.
+        request_hook: A function that will be called before a request is sent.  
+            The function will be called with the span and the request.
+        response_hook: A function that will be called after a response is received.
+            The function will be called with the span and the response.
+        tracer_provider: The tracer provider to use. If not provided, the global
+            tracer provider will be used.
+        kwargs: Additional keyword arguments.
+    """
+    all_kwargs = {
+        "excluded_urls": excluded_urls,
+        "request_hook": request_hook,
+        "response_hook": response_hook,
+        "tracer_provider": tracer_provider or get_tracer_provider(),
+        **kwargs
+    }
+    RequestsInstrumentor().instrument(**all_kwargs)
+    logger.info("Requests instrumented.")

aworld/trace/instrumentation/semconv.py

@@ -0,0 +1,29 @@
+# GenAI semconv attribute names
+GEN_AI_SYSTEM = "gen_ai.system"
+GEN_AI_REQUEST_MODEL = "gen_ai.request.model"
+GEN_AI_REQUEST_FREQUENCY_PENALTY = "gen_ai.request.frequency_penalty"
+GEN_AI_REQUEST_MAX_TOKENS = "gen_ai.request.max_tokens"
+GEN_AI_REQUEST_PRESENCE_PENALTY = "gen_ai.request.presence_penalty"
+GEN_AI_REQUEST_STOP_SEQUENCES = "gen_ai.request.stop_sequences"
+GEN_AI_REQUEST_TEMPERATURE = "gen_ai.request.temperature"
+GEN_AI_REQUEST_TOP_K = "gen_ai.request.top_k"
+GEN_AI_REQUEST_TOP_P = "gen_ai.request.top_p"
+GEN_AI_REQUEST_STREAMING = "gen_ai.request.streaming"
+GEN_AI_REQUEST_USER = "gen_ai.request.user"
+GEN_AI_REQUEST_EXTRA_HEADERS = "gen_ai.request.extra_headers"
+GEN_AI_PROMPT = "gen_ai.prompt"
+GEN_AI_PROMPT_TOOLS = "gen_ai.prompt.tools"
+GEN_AI_COMPLETION = "gen_ai.completion"
+GEN_AI_COMPLETION_TOOL_CALLS = "gen_ai.completion.tool_calls"
+GEN_AI_COMPLETION_CONTENT = "gen_ai.completion.content"
+GEN_AI_DURATION = "gen_ai.duration"
+GEN_AI_FIRST_TOKEN_DURATION = "gen_ai.first_token_duration"
+GEN_AI_RESPONSE_FINISH_REASONS = "gen_ai.response.finish_reasons"
+GEN_AI_RESPONSE_ID = "gen_ai.response.id"
+GEN_AI_RESPONSE_MODEL = "gen_ai.response.model"
+GEN_AI_USAGE_INPUT_TOKENS = "gen_ai.usage.input_tokens"
+GEN_AI_USAGE_OUTPUT_TOKENS = "gen_ai.usage.output_tokens"
+GEN_AI_USAGE_TOTAL_TOKENS = "gen_ai.usage.total_tokens"
+GEN_AI_OPERATION_NAME = "gen_ai.operation.name"
+GEN_AI_METHOD_NAME = "gen_ai.method.name"
+GEN_AI_SERVER_ADDRESS = "gen_ai.server.address"

aworld/trace/instrumentation/threading.py

@@ -0,0 +1,149 @@
+import threading
+from typing import Protocol, TypeVar, Any, Callable
+from wrapt import wrap_function_wrapper
+from concurrent import futures
+import aworld.trace as trace
+from aworld.trace.base import TraceContext, Span
+from aworld.trace.propagator import get_global_trace_context
+from aworld.trace.instrumentation import Instrumentor
+from aworld.trace.instrumentation.utils import unwrap
+from aworld.logs.util import logger
+
+
+R = TypeVar("R")
+
+
+class HasTraceContext(Protocol):
+    _trace_context: TraceContext
+
+
+class ThreadingInstrumentor(Instrumentor):
+    '''
+    Trace instrumentor for threading
+    '''
+
+    def instrumentation_dependencies(self) -> str:
+        return ()
+
+    def _instrument(self, **kwargs: Any):
+        self._instrument_thread()
+        self._instrument_timer()
+        self._instrument_thread_pool()
+
+    def _uninstrument(self, **kwargs: Any):
+        self._uninstrument_thread()
+        self._uninstrument_timer()
+        self._uninstrument_thread_pool()
+
+    @staticmethod
+    def _instrument_thread():
+        wrap_function_wrapper(
+            threading.Thread,
+            "start",
+            ThreadingInstrumentor.__wrap_threading_start,
+        )
+        wrap_function_wrapper(
+            threading.Thread,
+            "run",
+            ThreadingInstrumentor.__wrap_threading_run,
+        )
+
+    @staticmethod
+    def _instrument_timer():
+        wrap_function_wrapper(
+            threading.Timer,
+            "start",
+            ThreadingInstrumentor.__wrap_threading_start,
+        )
+        wrap_function_wrapper(
+            threading.Timer,
+            "run",
+            ThreadingInstrumentor.__wrap_threading_run,
+        )
+
+    @staticmethod
+    def _instrument_thread_pool():
+        wrap_function_wrapper(
+            futures.ThreadPoolExecutor,
+            "submit",
+            ThreadingInstrumentor.__wrap_thread_pool_submit,
+        )
+
+    @staticmethod
+    def _uninstrument_thread():
+        unwrap(threading.Thread, "start")
+        unwrap(threading.Thread, "run")
+
+    @staticmethod
+    def _uninstrument_timer():
+        unwrap(threading.Timer, "start")
+        unwrap(threading.Timer, "run")
+
+    @staticmethod
+    def _uninstrument_thread_pool():
+        unwrap(futures.ThreadPoolExecutor, "submit")
+
+    @staticmethod
+    def __wrap_threading_start(
+        call_wrapped: Callable[[], None],
+        instance: HasTraceContext,
+        args: tuple[()],
+        kwargs: dict[str, Any],
+    ) -> None:
+        span: Span = trace.get_current_span()
+        if span:
+            instance._trace_context = TraceContext(
+                trace_id=span.get_trace_id(), span_id=span.get_span_id())
+        return call_wrapped(*args, **kwargs)
+
+    @staticmethod
+    def __wrap_threading_run(
+        call_wrapped: Callable[..., R],
+        instance: HasTraceContext,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> R:
+
+        token = None
+        try:
+            if hasattr(instance, "_trace_context"):
+                if instance._trace_context:
+                    token = get_global_trace_context().set(instance._trace_context)
+            return call_wrapped(*args, **kwargs)
+        finally:
+            if token:
+                get_global_trace_context().reset(token)
+
+    @staticmethod
+    def __wrap_thread_pool_submit(
+        call_wrapped: Callable[..., R],
+        instance: futures.ThreadPoolExecutor,
+        args: tuple[Callable[..., Any], ...],
+        kwargs: dict[str, Any],
+    ) -> R:
+        # obtain the original function and wrapped kwargs
+        original_func = args[0]
+        trace_context = None
+        span: Span = trace.get_current_span()
+        if span and span.get_trace_id() != "":
+            trace_context = TraceContext(
+                trace_id=span.get_trace_id(), span_id=span.get_span_id())
+
+        def wrapped_func(*func_args: Any, **func_kwargs: Any) -> R:
+            token = None
+            try:
+                if trace_context:
+                    token = get_global_trace_context().set(trace_context)
+                return original_func(*func_args, **func_kwargs)
+            finally:
+                if token:
+                    get_global_trace_context().reset(token)
+
+        # replace the original function with the wrapped function
+        new_args: tuple[Callable[..., Any], ...] = (wrapped_func,) + args[1:]
+        return call_wrapped(*new_args, **kwargs)
+
+
+def instrument_theading(**kwargs: Any) -> None:
+    ThreadingInstrumentor().instrument(**kwargs)
+    logger.info("Threading instrumented")

aworld/trace/instrumentation/utils.py

@@ -0,0 +1,31 @@
+from importlib import import_module
+from wrapt import ObjectProxy
+
+
+def unwrap(obj: object, attr: str):
+    """Given a function that was wrapped by wrapt.wrap_function_wrapper, unwrap it
+
+    The object containing the function to unwrap may be passed as dotted module path string.
+
+    Args:
+        obj: Object that holds a reference to the wrapped function or dotted import path as string
+        attr (str): Name of the wrapped function
+    """
+    if isinstance(obj, str):
+        try:
+            module_path, class_name = obj.rsplit(".", 1)
+        except ValueError as exc:
+            raise ImportError(
+                f"Cannot parse '{obj}' as dotted import path"
+            ) from exc
+        module = import_module(module_path)
+        try:
+            obj = getattr(module, class_name)
+        except AttributeError as exc:
+            raise ImportError(
+                f"Cannot import '{class_name}' from '{module}'"
+            ) from exc
+
+    func = getattr(obj, attr, None)
+    if func and isinstance(func, ObjectProxy) and hasattr(func, "__wrapped__"):
+        setattr(obj, attr, func.__wrapped__)