ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter9/gaia-experience/AWorld/aworld/trace/instrumentation/__init__.py

@@ -0,0 +1,87 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+from abc import ABC, abstractmethod
+from typing import Any, Collection
+from importlib_metadata import version, PackageNotFoundError
+from aworld.logs.util import logger
+from aworld.utils.import_package import import_package
+import_package("packaging")  # noqa
+from packaging.requirements import Requirement, InvalidRequirement
+
+
+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.
+        """

chapter9/gaia-experience/AWorld/aworld/trace/instrumentation/agent/__init__.py

@@ -0,0 +1,238 @@
+import wrapt
+import time
+import traceback
+import aworld.trace.constants as trace_constants
+from typing import Collection, Any
+from aworld.trace.instrumentation import Instrumentor
+from aworld.trace.instrumentation import semconv
+from aworld.trace.base import (
+    Tracer,
+    SpanType,
+    get_tracer_provider_silent
+)
+from aworld.logs.util import logger
+from aworld.metrics.context_manager import MetricContext
+from aworld.metrics.template import MetricTemplate
+from aworld.metrics.metric import MetricType
+
+agent_duration_histogram = MetricTemplate(
+    type=MetricType.HISTOGRAM,
+    name="agent_run_duration",
+    unit="s",
+    description="Agent run duration",
+)
+
+agent_run_counter = MetricTemplate(
+    type=MetricType.COUNTER,
+    name="agent_run_counter",
+    unit="time",
+    description="Number of agent run or async run",
+)
+
+agent_usage_histogram = MetricTemplate(
+    type=MetricType.HISTOGRAM,
+    name="agent_token_usage",
+    unit="token",
+    description="Agent token usage"
+)
+
+
+def get_agent_span_attributes(instance, message):
+    return {
+        semconv.AGENT_ID: instance.id(),
+        semconv.AGENT_NAME: instance.name(),
+        semconv.TASK_ID: message.context.task_id if (message.context and message.context.task_id) else "",
+        semconv.SESSION_ID: message.context.session_id if (message.context and message.context.session_id) else message.session_id,
+        semconv.USER_ID: message.context.user if (message.context and message.context.user) else "",
+        trace_constants.ATTRIBUTES_MESSAGE_RUN_TYPE_KEY: trace_constants.RunType.AGNET.value
+    }
+
+
+def _end_span(span):
+    if span:
+        span.end()
+
+
+def _record_metric(duration, attributes, exception=None):
+    if MetricContext.metric_initialized():
+        MetricContext.histogram_record(agent_duration_histogram, duration, labels=attributes)
+        if exception:
+            run_counter_attr = {
+                semconv.AGENT_RUN_SUCCESS: "0",
+                "error.type": exception.__class__.__name__,
+                **attributes
+            }
+        else:
+            run_counter_attr = {
+                semconv.AGENT_RUN_SUCCESS: "1",
+                **attributes
+            }
+        MetricContext.count(agent_run_counter, 1, labels=run_counter_attr)
+
+
+def _record_exception(span, start_time, exception, attributes):
+    try:
+        duration = time.time() - start_time if "start_time" in locals() else 0
+        if span and span.is_recording:
+            span.record_exception(exception=exception)
+        _record_metric(duration, attributes, exception)
+    except Exception as e:
+        logger.warning(f"agent instrument record exception error.{e}")
+
+
+def _record_response(instance,
+                     start_time,
+                     response,
+                     attributes):
+    try:
+        duration = time.time() - start_time if "start_time" in locals() else 0
+        _record_metric(duration, attributes)
+        # if instance and instance.agent_context and instance.agent_context.llm_output and MetricContext.metric_initialized():
+        #     usage = instance.agent_context.llm_output.usage
+        #     for usage_type in ["completion_tokens", "prompt_tokens", "total_tokens"]:
+        #         if usage and usage.get(usage_type):
+        #             labels = {
+        #                 **attributes,
+        #                 semconv.AGENT_USAGE_TYPE: usage_type
+        #             }
+        #             MetricContext.histogram_record(
+        #                 agent_usage_histogram,
+        #                 usage.get(usage_type),
+        #                 labels=labels
+        #             )
+    except Exception as e:
+        logger.warning(f"agent instrument record response error.{e}")
+
+
+def _async_run_class_wrapper(tracer: Tracer):
+    async def _async_run_wrapper(wrapped, instance, args, kwargs):
+        span = None
+        message = args[0] or kwargs.get("message")
+        attributes = get_agent_span_attributes(instance, message)
+        if tracer:
+            span = tracer.start_span(
+                name=trace_constants.SPAN_NAME_PREFIX_AGENT + "async_run",
+                span_type=SpanType.SERVER,
+                attributes=attributes
+            )
+        start_time = time.time()
+        try:
+            response = await wrapped(*args, **kwargs)
+            _record_response(instance, start_time, response, attributes)
+        except Exception as e:
+            _record_exception(span=span,
+                              start_time=start_time,
+                              exception=e,
+                              attributes=attributes
+                              )
+            _end_span(span)
+            raise e
+        _end_span(span)
+        return response
+    return _async_run_wrapper
+
+
+async def _async_run_instance_wrapper(tracer: Tracer):
+
+    @wrapt.decorator
+    async def _awrapper(wrapped, instance, args, kwargs):
+        wrapper_func = _async_run_class_wrapper(tracer=tracer)
+        return await wrapper_func(wrapped, instance, args, kwargs)
+
+    return _awrapper
+
+
+def _call_llm_model_class_wrapper(tracer: Tracer):
+    async def _call_llm_model_wrapper(wrapped, instance, args, kwargs):
+        attributes = {
+            semconv.AGENT_ID: instance.id(),
+            semconv.AGENT_NAME: instance.name()
+        }
+        if hasattr(instance, "context") and instance.context:
+            attributes.update({
+                semconv.TASK_ID: instance.context.task_id if (instance.context and instance.context.task_id) else "",
+                semconv.SESSION_ID: instance.context.session_id if (instance.context and instance.context.session_id) else instance.session_id,
+                semconv.USER_ID: instance.context.user if (instance.context and instance.context.user) else ""
+            })
+        try:
+            response = await wrapped(*args, **kwargs)
+            try:
+                usage = response.usage if hasattr(response, "usage") else None
+                for usage_type in ["completion_tokens", "prompt_tokens", "total_tokens"]:
+                    if usage and usage.get(usage_type):
+                        labels = {
+                            **attributes,
+                            semconv.AGENT_USAGE_TYPE: usage_type
+                        }
+                        MetricContext.histogram_record(
+                            agent_usage_histogram,
+                            usage.get(usage_type),
+                            labels=labels
+                        )
+
+            except Exception as e:
+                logger.warning(f"agent instrument record response error.{e}")
+        except Exception as e:
+            raise e
+        return response
+    return _call_llm_model_wrapper
+
+
+async def _call_llm_model_instance_wrapper(tracer: Tracer):
+
+    @wrapt.decorator
+    async def _awrapper(wrapped, instance, args, kwargs):
+        wrapper_func = _call_llm_model_class_wrapper(tracer=tracer)
+        return await wrapper_func(wrapped, instance, args, kwargs)
+
+    return _awrapper
+
+
+class AgentInstrumentor(Instrumentor):
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ()
+
+    def _instrument(self, **kwargs):
+        agent_trace_enabled = kwargs.get("trace_enabled", False)
+        tracer_provider = get_tracer_provider_silent()
+        tracer = None
+        if tracer_provider and agent_trace_enabled:
+            tracer = tracer_provider.get_tracer(
+                "aworld.trace.instrumentation.agent")
+
+        try:
+            wrapt.wrap_function_wrapper(
+                "aworld.core.agent.base",
+                "BaseAgent.async_run",
+                _async_run_class_wrapper(tracer=tracer)
+            )
+            wrapt.wrap_function_wrapper(
+                "aworld.agents.llm_agent",
+                "Agent.invoke_model",
+                _call_llm_model_class_wrapper(tracer=tracer)
+            )
+        except Exception as err:
+            logger.warning(f"AgentInstrumentor#_instrument failed ,err is {err}")
+
+    def _uninstrument(self, **kwargs: Any):
+        pass
+
+
+def wrap_agent(agent: 'aworld.core.agent.base.BaseAgent'):
+    try:
+        tracer_provider = get_tracer_provider_silent()
+        if not tracer_provider:
+            return agent
+        tracer = tracer_provider.get_tracer(
+            "aworld.trace.instrumentation.agent")
+
+        async_run_wrapper = _async_run_instance_wrapper(tracer)
+        agent.async_run = async_run_wrapper(agent.async_run)
+        if hasattr(agent, "_call_llm_model"):
+            call_llm_model_wrapper = _call_llm_model_instance_wrapper(tracer)
+            agent._call_llm_model = call_llm_model_wrapper(agent._call_llm_model)
+    except Exception:
+        logger.warning(traceback.format_exc())
+
+    return agent

chapter9/gaia-experience/AWorld/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()

chapter9/gaia-experience/AWorld/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

chapter9/gaia-experience/AWorld/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.")

chapter9/gaia-experience/AWorld/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.")

chapter9/gaia-experience/AWorld/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

chapter9/gaia-experience/AWorld/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)

chapter9/gaia-experience/AWorld/aworld/trace/instrumentation/openai/__init__.py

@@ -0,0 +1,333 @@
+import wrapt
+import time
+import openai
+import traceback
+import aworld.trace.instrumentation.semconv as semconv
+from typing import Collection, Any, Union
+from aworld.trace.instrumentation import Instrumentor
+from aworld.trace.base import (
+    Tracer,
+    SpanType,
+    get_tracer_provider_silent
+)
+from aworld.trace.constants import ATTRIBUTES_MESSAGE_RUN_TYPE_KEY, RunType
+from aworld.trace.instrumentation.openai.inout_parse import (
+    run_async,
+    handle_openai_request,
+    is_streaming_response,
+    record_stream_response_chunk,
+    parse_openai_response,
+    record_stream_token_usage,
+    model_as_dict,
+    parse_response_message,
+)
+from aworld.trace.instrumentation.llm_metrics import (
+    record_exception_metric,
+    record_chat_response_metric,
+    record_streaming_time_to_first_token,
+    record_streaming_time_to_generate
+)
+from aworld.logs.util import logger
+
+
+def _chat_wrapper(tracer: Tracer):
+
+    @wrapt.decorator
+    def wrapper(wrapped, instance, args, kwargs):
+        model_name = kwargs.get("model", "")
+        if not model_name:
+            model_name = "OpenAI"
+        span_attributes = {}
+        span_attributes[ATTRIBUTES_MESSAGE_RUN_TYPE_KEY] = RunType.LLM.value
+
+        span = tracer.start_span(
+            name=model_name, span_type=SpanType.CLIENT, attributes=span_attributes)
+
+        run_async(handle_openai_request(span, kwargs, instance))
+        start_time = time.time()
+        try:
+            response = wrapped(*args, **kwargs)
+        except Exception as e:
+            record_exception(span=span,
+                             start_time=start_time,
+                             exception=e
+                             )
+            span.end()
+            raise e
+
+        if is_streaming_response(response):
+            return WrappedStreamResponse(span=span,
+                                         response=response,
+                                         instance=instance,
+                                         start_time=start_time,
+                                         request_kwargs=kwargs
+                                         )
+
+        record_completion(span=span,
+                          start_time=start_time,
+                          response=response,
+                          request_kwargs=kwargs,
+                          instance=instance
+                          )
+        span.end()
+        return response
+
+    return wrapper
+
+
+def _achat_class_wrapper(tracer: Tracer):
+
+    async def awrapper(wrapped, instance, args, kwargs):
+        model_name = kwargs.get("model", "")
+        if not model_name:
+            model_name = "OpenAI"
+        span_attributes = {}
+        span_attributes[ATTRIBUTES_MESSAGE_RUN_TYPE_KEY] = RunType.LLM.value
+
+        span = tracer.start_span(
+            name=model_name, span_type=SpanType.CLIENT, attributes=span_attributes)
+
+        await handle_openai_request(span, kwargs, instance)
+        start_time = time.time()
+        try:
+            response = await wrapped(*args, **kwargs)
+        except Exception as e:
+            record_exception(span=span,
+                             start_time=start_time,
+                             exception=e
+                             )
+            span.end()
+            raise e
+
+        if is_streaming_response(response):
+            return WrappedStreamResponse(span=span,
+                                         response=response,
+                                         instance=instance,
+                                         start_time=start_time,
+                                         request_kwargs=kwargs
+                                         )
+        record_completion(span=span,
+                          start_time=start_time,
+                          response=response,
+                          request_kwargs=kwargs,
+                          instance=instance
+                          )
+        span.end()
+        return response
+
+    return awrapper
+
+
+def _achat_instance_wrapper(tracer: Tracer):
+
+    @wrapt.decorator
+    async def _awrapper(wrapped, instance, args, kwargs):
+        wrapper_func = _achat_class_wrapper(tracer)
+        return await wrapper_func(wrapped, instance, args, kwargs)
+
+    return _awrapper
+
+
+def record_exception(span, start_time, exception):
+    '''
+    record openai chat exception to trace and metrics
+    '''
+    try:
+        duration = time.time() - start_time if "start_time" in locals() else 0
+        if span.is_recording:
+            span.record_exception(exception=exception)
+        record_exception_metric(exception=exception, duration=duration)
+    except Exception as e:
+        logger.warning(f"openai instrument record exception error.{e}")
+
+
+def record_completion(span,
+                      start_time,
+                      response,
+                      request_kwargs,
+                      instance):
+    '''
+    Record chat completion to trace and metrics
+    '''
+    duration = time.time() - start_time if "start_time" in locals() else 0
+    response_dict = model_as_dict(response)
+    attributes = parse_openai_response(
+        response_dict, request_kwargs, instance, False)
+    usage = response_dict.get("usage")
+    choices = response_dict.get("choices")
+    prompt_tokens = usage.get("prompt_tokens")
+    completion_tokens = usage.get("completion_tokens")
+
+    span_attributes = {
+        **attributes,
+        semconv.GEN_AI_USAGE_INPUT_TOKENS: prompt_tokens,
+        semconv.GEN_AI_USAGE_OUTPUT_TOKENS: completion_tokens,
+        semconv.GEN_AI_DURATION: duration
+    }
+    span_attributes.update(parse_response_message(choices))
+    span.set_attributes(span_attributes)
+    record_chat_response_metric(attributes=attributes,
+                                prompt_tokens=prompt_tokens,
+                                completion_tokens=completion_tokens,
+                                duration=duration,
+                                choices=choices
+                                )
+
+
+class WrappedStreamResponse(wrapt.ObjectProxy):
+
+    def __init__(
+        self,
+        span,
+        response,
+        instance=None,
+        start_time=None,
+        request_kwargs=None
+    ):
+        super().__init__(response)
+        self._span = span
+        self._instance = instance
+        self._start_time = start_time
+        self._complete_response = {"choices": [], "model": ""}
+        self._first_token_recorded = False
+        self._request_kwargs = request_kwargs
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.__wrapped__.__exit__(exc_type, exc_val, exc_tb)
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.__wrapped__.__aexit__(exc_type, exc_val, exc_tb)
+
+    def __iter__(self):
+        return self
+
+    def __aiter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            chunk = self.__wrapped__.__next__()
+        except Exception as e:
+            if isinstance(e, StopIteration):
+                self._close_span()
+            raise e
+        else:
+            self._process_stream_chunk(chunk)
+            return chunk
+
+    async def __anext__(self):
+        try:
+            chunk = await self.__wrapped__.__anext__()
+        except Exception as e:
+            if isinstance(e, StopAsyncIteration):
+                self._close_span()
+            raise e
+        else:
+            self._process_stream_chunk(chunk)
+            return chunk
+
+    def _process_stream_chunk(self, chunk):
+        record_stream_response_chunk(chunk, self._complete_response)
+        if not self._first_token_recorded:
+            self._time_of_first_token = time.time()
+            duration = self._time_of_first_token - self._start_time
+            attribute = parse_openai_response(
+                self._complete_response, self._request_kwargs, self._instance, True)
+            record_streaming_time_to_first_token(duration, attribute)
+            self._first_token_recorded = True
+
+    def _close_span(self):
+        duration = None
+        first_token_duration = None
+        first_token_to_generate_duration = None
+        if self._start_time and isinstance(self._start_time, (float, int)):
+            duration = time.time() - self._start_time
+        if self._time_of_first_token and self._start_time and isinstance(self._start_time, (float, int)):
+            first_token_duration = self._time_of_first_token - self._start_time
+            first_token_to_generate_duration = time.time() - self._time_of_first_token
+        prompt_usage, completion_usage = record_stream_token_usage(
+            self._complete_response, self._request_kwargs)
+        attributes = parse_openai_response(
+            self._complete_response, self._request_kwargs, self._instance, True)
+        choices = self._complete_response.get("choices")
+        span_attributes = {
+            **attributes,
+            "llm.prompt_tokens": prompt_usage,
+            "llm.completion_tokens": completion_usage,
+            "llm.duration": duration,
+            "llm.first_token_duration": first_token_duration
+        }
+        span_attributes.update(parse_response_message(choices))
+        self._span.set_attributes(span_attributes)
+        record_chat_response_metric(attributes=attributes,
+                                    prompt_tokens=prompt_usage,
+                                    completion_tokens=completion_usage,
+                                    duration=duration,
+                                    choices=choices
+                                    )
+        record_streaming_time_to_generate(
+            first_token_to_generate_duration, attributes)
+
+        self._span.end()
+
+
+class OpenAIInstrumentor(Instrumentor):
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ("openai >= 1.0.0",)
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = tracer_provider.get_tracer(
+            "aworld.trace.instrumentation.openai")
+
+        wrapt.wrap_function_wrapper(
+            "openai.resources.chat.completions",
+            "Completions.create",
+            _chat_wrapper(tracer=tracer)
+        )
+
+        wrapt.wrap_function_wrapper(
+            "openai.resources.chat.completions",
+            "AsyncCompletions.create",
+            _achat_class_wrapper(tracer)
+        )
+
+    def _instrument(self, **kwargs: Any):
+        pass
+
+
+def wrap_openai(client: Union[openai.OpenAI, openai.AsyncOpenAI]):
+    """Patch the OpenAI client to make it traceable.
+       Example:
+       client = wrap_openai(openai.OpenAI())
+    """
+    try:
+        tracer_provider = get_tracer_provider_silent()
+        if not tracer_provider:
+            return
+        tracer = tracer_provider.get_tracer(
+            "aworld.trace.instrumentation.openai")
+
+        if isinstance(client, openai.OpenAI):
+            wrapper = _chat_wrapper(tracer)
+            client.chat.completions.create = wrapper(
+                client.chat.completions.create)
+            logger.info(
+                f"[{client.__class__}]client.chat.completions.create be warpped")
+        if isinstance(client, openai.AsyncOpenAI):
+            awrapper = _achat_instance_wrapper(tracer)
+            client.chat.completions.create = awrapper(
+                client.chat.completions.create)
+            logger.info(
+                f"[{client.__class__}]client.chat.completions.create be warpped")
+    except Exception:
+        logger.warning(traceback.format_exc())
+
+    return client

chapter9/gaia-experience/AWorld/aworld/trace/instrumentation/openai/inout_parse.py

@@ -0,0 +1,296 @@
+import asyncio
+import os
+import threading
+import copy
+import json
+import openai
+from importlib.metadata import version
+from aworld.logs.util import logger
+from aworld.trace.base import Span
+from aworld.utils import import_package
+import aworld.trace.instrumentation.semconv as semconv
+
+_PYDANTIC_VERSION = version("pydantic")
+
+
+def should_trace_prompts():
+    '''Determine whether it is necessary to record the message
+    '''
+    return (os.getenv("SHOULD_TRACE_PROMPTS") or "true").lower() == "true"
+
+
+def need_flatten_messages():
+    '''Determine whether it is necessary to flatten the messages
+    '''
+    return (os.getenv("TRACE_FLATTEN_MESSAGES") or "false").lower() == "true"
+
+
+def run_async(method):
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        loop = None
+
+    if loop and loop.is_running():
+        thread = threading.Thread(target=lambda: asyncio.run(method))
+        thread.start()
+        thread.join()
+    else:
+        asyncio.run(method)
+
+
+async def handle_openai_request(span: Span, kwargs, instance):
+    if not span or not span.is_recording():
+        return
+    try:
+        attributes = parser_request_params(kwargs, instance)
+        if should_trace_prompts():
+            messages = kwargs.get("messages")
+            if need_flatten_messages():
+                attributes.update(parse_request_message(messages))
+            else:
+                attributes.update({
+                    semconv.GEN_AI_PROMPT: str(messages),
+                })
+        span.set_attributes(attributes)
+    except ValueError as e:
+        logger.warning(f"trace handle openai request error: {e}")
+
+
+def parser_request_params(kwargs, instance):
+    attributes = {
+        semconv.GEN_AI_SYSTEM: "OpenAI",
+        semconv.GEN_AI_REQUEST_MODEL: kwargs.get("model", ""),
+        semconv.GEN_AI_REQUEST_MAX_TOKENS: kwargs.get("max_tokens", ""),
+        semconv.GEN_AI_REQUEST_TEMPERATURE: kwargs.get("temperature", ""),
+        semconv.GEN_AI_REQUEST_TOP_P: kwargs.get("top_p", ""),
+        semconv.GEN_AI_REQUEST_FREQUENCY_PENALTY: kwargs.get("frequency_penalty", ""),
+        semconv.GEN_AI_REQUEST_PRESENCE_PENALTY: kwargs.get("presence_penalty", ""),
+        semconv.GEN_AI_REQUEST_USER: kwargs.get("user", ""),
+        semconv.GEN_AI_REQUEST_EXTRA_HEADERS: kwargs.get("extra_headers", ""),
+        semconv.GEN_AI_REQUEST_STREAMING: kwargs.get("stream", ""),
+        semconv.GEN_AI_OPERATION_NAME: "chat"
+    }
+
+    client = instance._client
+    if isinstance(client, (openai.AsyncOpenAI, openai.OpenAI)):
+        attributes.update({"llm.base_url": str(client.base_url)})
+
+    filterd_attri = {k: v for k, v in attributes.items()
+                     if (v and v != "")}
+    return filterd_attri
+
+
+def is_streaming_response(response):
+    return isinstance(response, openai.Stream) or isinstance(response, openai.AsyncStream)
+
+
+def parse_openai_response(response, request_kwargs, instance, is_streaming):
+    return {
+        semconv.GEN_AI_RESPONSE_MODEL: response.get("model") or request_kwargs.get("model") or None,
+        semconv.GEN_AI_SERVER_ADDRESS: _get_openai_base_url(instance)
+    }
+
+
+def record_stream_token_usage(complete_response, request_kwargs) -> tuple[int, int]:
+    '''
+        return (prompt_usage, completion_usage)
+    '''
+    prompt_usage = 0
+    completion_usage = 0
+
+    # prompt_usage
+    if request_kwargs and request_kwargs.get("messages"):
+        prompt_content = ""
+        model_name = complete_response.get(
+            "model") or request_kwargs.get("model") or "gpt-4"
+        for msg in request_kwargs.get("messages"):
+            if msg.get("content"):
+                prompt_content += msg.get("content")
+        if model_name:
+            prompt_usage = get_token_count_from_string(
+                prompt_content, model_name)
+
+    # completion_usage
+    if complete_response.get("choices"):
+        completion_content = ""
+        model_name = complete_response.get("model") or "gpt-4"
+
+        for choice in complete_response.get("choices"):
+            if choice.get("message") and choice.get("message").get("content"):
+                completion_content += choice["message"]["content"]
+
+        if model_name:
+            completion_usage = get_token_count_from_string(
+                completion_content, model_name)
+
+    return (prompt_usage, completion_usage)
+
+
+def _get_openai_base_url(instance):
+    if hasattr(instance, "_client"):
+        client = instance._client  # pylint: disable=protected-access
+        if isinstance(client, (openai.AsyncOpenAI, openai.OpenAI)):
+            return str(client.base_url)
+
+    return ""
+
+
+def get_token_count_from_string(string: str, model_name: str):
+    import_package("tiktoken")
+    import tiktoken
+
+    if tiktoken_encodings.get(model_name) is None:
+        try:
+            encoding = tiktoken.encoding_for_model(model_name)
+        except KeyError as ex:
+            logger.warning(
+                f"Failed to get tiktoken encoding for model_name {model_name}, error: {str(ex)}")
+            return None
+
+        tiktoken_encodings[model_name] = encoding
+    else:
+        encoding = tiktoken_encodings.get(model_name)
+
+    token_count = len(encoding.encode(string))
+    return token_count
+
+
+def record_stream_response_chunk(chunk, complete_response):
+    chunk = model_as_dict(chunk)
+    complete_response["model"] = chunk.get("model")
+    complete_response["id"] = chunk.get("id")
+
+    # prompt filter results
+    if chunk.get("prompt_filter_results"):
+        complete_response["prompt_filter_results"] = chunk.get(
+            "prompt_filter_results")
+
+    for choice in chunk.get("choices"):
+        index = choice.get("index")
+        if len(complete_response.get("choices")) <= index:
+            complete_response["choices"].append(
+                {"index": index, "message": {"content": "", "role": ""}})
+        complete_choice = complete_response.get("choices")[index]
+        if choice.get("finish_reason"):
+            complete_choice["finish_reason"] = choice.get("finish_reason")
+        if choice.get("content_filter_results"):
+            complete_choice["content_filter_results"] = choice.get(
+                "content_filter_results")
+
+        delta = choice.get("delta")
+
+        if delta and delta.get("content"):
+            complete_choice["message"]["content"] += delta.get("content")
+
+        if delta and delta.get("role"):
+            complete_choice["message"]["role"] = delta.get("role")
+        if delta and delta.get("tool_calls"):
+            tool_calls = delta.get("tool_calls")
+            if not isinstance(tool_calls, list) or len(tool_calls) == 0:
+                continue
+
+            if not complete_choice["message"].get("tool_calls"):
+                complete_choice["message"]["tool_calls"] = []
+
+            for tool_call in tool_calls:
+                i = int(tool_call["index"])
+                if len(complete_choice["message"]["tool_calls"]) <= i:
+                    complete_choice["message"]["tool_calls"].append(
+                        {"id": "", "function": {"name": "", "arguments": ""}}
+                    )
+
+                span_tool_call = complete_choice["message"]["tool_calls"][i]
+                span_function = span_tool_call["function"]
+                tool_call_function = tool_call.get("function")
+
+                if tool_call.get("id"):
+                    span_tool_call["id"] = tool_call.get("id")
+                if tool_call_function and tool_call_function.get("name"):
+                    span_function["name"] = tool_call_function.get("name")
+                if tool_call_function and tool_call_function.get("arguments"):
+                    span_function["arguments"] += tool_call_function.get(
+                        "arguments")
+
+
+def parse_request_message(messages):
+    '''
+    flatten request message to attributes
+    '''
+    attributes = {}
+    for i, msg in enumerate(messages):
+        prefix = f"{semconv.GEN_AI_PROMPT}.{i}"
+        attributes.update({f"{prefix}.role": msg.get("role")})
+        if msg.get("content"):
+            content = copy.deepcopy(msg.get("content"))
+            content = json.dumps(content)
+            attributes.update({f"{prefix}.content": content})
+        if msg.get("tool_call_id"):
+            attributes.update({
+                f"{prefix}.tool_call_id": msg.get("tool_call_id")})
+        tool_calls = msg.get("tool_calls")
+        if tool_calls:
+            for i, tool_call in enumerate(tool_calls):
+                tool_call = model_as_dict(tool_call)
+                function = tool_call.get("function")
+                attributes.update({
+                    f"{prefix}.tool_calls.{i}.id": tool_call.get("id")})
+                attributes.update({
+                    f"{prefix}.tool_calls.{i}.name": function.get("name")})
+                attributes.update({
+                    f"{prefix}.tool_calls.{i}.arguments": function.get("arguments")})
+    return attributes
+
+
+def parse_response_message(choices) -> dict:
+    attributes = {}
+    if not should_trace_prompts():
+        return attributes
+    for choice in choices:
+        index = choice.get("index")
+        prefix = f"{semconv.GEN_AI_COMPLETION}.{index}"
+        attributes.update(
+            {f"{prefix}.finish_reason": choice.get("finish_reason")})
+
+        message = choice.get("message")
+        if not message:
+            continue
+
+        attributes.update({f"{prefix}.role": message.get("role")})
+
+        if message.get("refusal"):
+            attributes.update({f"{prefix}.refusal": message.get("refusal")})
+        else:
+            attributes.update({f"{prefix}.content": message.get("content")})
+
+        function_call = message.get("function_call")
+        if function_call:
+            attributes.update(
+                {f"{prefix}.tool_calls.0.name": function_call.get("name")})
+            attributes.update(
+                {f"{prefix}.tool_calls.0.arguments": function_call.get("arguments")})
+
+        tool_calls = message.get("tool_calls")
+        if tool_calls:
+            for i, tool_call in enumerate(tool_calls):
+                function = tool_call.get("function")
+                attributes.update(
+                    {f"{prefix}.tool_calls.{i}.id": tool_call.get("id")})
+                attributes.update(
+                    {f"{prefix}.tool_calls.{i}.name": function.get("name")})
+                attributes.update(
+                    {f"{prefix}.tool_calls.{i}.arguments": function.get("arguments")})
+    return attributes
+
+
+def model_as_dict(model):
+    if isinstance(model, dict):
+        return model
+    if _PYDANTIC_VERSION < "2.0.0":
+        return model.dict()
+    if hasattr(model, "model_dump"):
+        return model.model_dump()
+    elif hasattr(model, "parse"):  # Raw API response
+        return model_as_dict(model.parse())
+    else:
+        return model

chapter9/gaia-experience/AWorld/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.")

chapter9/gaia-experience/AWorld/aworld/trace/instrumentation/semconv.py

@@ -0,0 +1,46 @@
+# GenAI semconv attribute names
+import os
+
+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"
+
+ATTRIBUTE_NAME_SPACE = os.getenv("ATTRIBUTE_NAME_SPACE", "aworld.")
+AGENT_ID = ATTRIBUTE_NAME_SPACE + "agent.id"
+AGENT_NAME = ATTRIBUTE_NAME_SPACE + "agent.name"
+AGENT_RUN_SUCCESS = ATTRIBUTE_NAME_SPACE + "agent.run.success"
+AGENT_USAGE_TYPE = ATTRIBUTE_NAME_SPACE + "agent.usage_type"
+TOOL_NAME = ATTRIBUTE_NAME_SPACE + "tool.name"
+TOOL_STEP_SUCCESS = ATTRIBUTE_NAME_SPACE + "tool.step.success"
+TASK = ATTRIBUTE_NAME_SPACE + "task"
+TASK_ID = ATTRIBUTE_NAME_SPACE + "task.id"
+TASK_INPUT = ATTRIBUTE_NAME_SPACE + "task.input"
+TASK_IS_SUB_TASK = ATTRIBUTE_NAME_SPACE + "task.is_sub_task"
+TASK_GROUP_ID = ATTRIBUTE_NAME_SPACE + "task.group_id"
+SESSION_ID = ATTRIBUTE_NAME_SPACE + "session.id"
+USER_ID = ATTRIBUTE_NAME_SPACE + "user.id"

chapter9/gaia-experience/AWorld/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")

chapter9/gaia-experience/AWorld/aworld/trace/instrumentation/tool/__init__.py

@@ -0,0 +1,249 @@
+from asyncio import iscoroutine
+import wrapt
+import time
+import traceback
+import aworld.trace.constants as trace_constants
+from typing import Collection, Any, Union, Sequence
+from aworld.trace.instrumentation import Instrumentor
+from aworld.trace.instrumentation import semconv
+from aworld.trace.base import (
+    Tracer,
+    SpanType,
+    get_tracer_provider_silent
+)
+from aworld.logs.util import logger
+from aworld.metrics.context_manager import MetricContext
+from aworld.metrics.template import MetricTemplate
+from aworld.metrics.metric import MetricType
+
+tool_duration_histogram = MetricTemplate(
+    type=MetricType.HISTOGRAM,
+    name="tool_step_duration",
+    unit="s",
+    description="tool step run duration",
+)
+
+tool_step_counter = MetricTemplate(
+    type=MetricType.COUNTER,
+    name="tool_step_counter",
+    unit="time",
+    description="Number of tool step run",
+)
+
+
+def get_tool_name(tool_name: str,
+                  action: Union['ActionModel', Sequence['ActionModel']]) -> tuple[str, trace_constants.RunType]:
+    if tool_name == "mcp" and action:
+        try:
+            if isinstance(action, (list, tuple)):
+                action = action[0]
+            mcp_name = action.action_name.split("__")[0]
+            return (mcp_name, trace_constants.RunType.MCP)
+        except ValueError:
+            logger.warning(traceback.format_exc())
+            return (tool_name, trace_constants.RunType.MCP)
+    return (tool_name, trace_constants.RunType.TOOL)
+
+
+def get_tool_span_attributes(instance, message: 'Message'):
+    run_type = trace_constants.RunType.TOOL
+    action = message.payload
+    agent_id = None
+    tool_name = instance.name()
+    if isinstance(action, (list, tuple)):
+        action = action[0]
+        if action:
+            agent_id = action.agent_name
+            tool_name, run_type = get_tool_name(action.tool_name, action)
+    return {
+        semconv.TOOL_NAME: tool_name,
+        semconv.AGENT_ID: agent_id,
+        semconv.AGENT_NAME: _get_agent_name_from_id(agent_id),
+        semconv.TASK_ID: message.context.task_id if (message.context and message.context.task_id) else "",
+        semconv.SESSION_ID: message.context.session_id if (message.context and message.context.session_id) else "",
+        semconv.USER_ID: message.context.user if (message.context and message.context.user) else "",
+        trace_constants.ATTRIBUTES_MESSAGE_RUN_TYPE_KEY: run_type.value
+    }
+
+
+def _end_span(span):
+    if span:
+        span.end()
+
+
+def _get_agent_name_from_id(agent_id):
+    if agent_id and '---' in agent_id:
+        return agent_id.split('---', 1)[0]
+    return agent_id
+
+
+def _record_metric(duration, attributes, exception=None):
+    if MetricContext.metric_initialized():
+        MetricContext.histogram_record(tool_duration_histogram, duration, labels=attributes)
+        if exception:
+            run_counter_attr = {
+                semconv.TOOL_STEP_SUCCESS: "0",
+                "error.type": exception.__class__.__name__,
+                **attributes
+            }
+        else:
+            run_counter_attr = {
+                semconv.TOOL_STEP_SUCCESS: "1",
+                **attributes
+            }
+        MetricContext.count(tool_step_counter, 1, labels=run_counter_attr)
+
+
+def _record_exception(span, start_time, exception, attributes):
+    try:
+        duration = time.time() - start_time if "start_time" in locals() else 0
+        if span.is_recording:
+            span.record_exception(exception=exception)
+        _record_metric(duration, attributes, exception)
+    except Exception as e:
+        logger.warning(f"tool instrument record exception error.{e}")
+
+
+def _record_response(instance,
+                     start_time,
+                     response,
+                     attributes):
+    try:
+        duration = time.time() - start_time if "start_time" in locals() else 0
+        _record_metric(duration, attributes)
+    except Exception as e:
+        logger.warning(f"tool instrument record response error.{e}")
+
+
+def _async_step_class_wrapper(tracer: Tracer):
+    async def _async_step_wrapper(wrapped, instance, args, kwargs):
+        span = None
+        message = args[0] or kwargs.get("message")
+        attributes = get_tool_span_attributes(instance, message)
+        if tracer:
+            span = tracer.start_span(
+                name=trace_constants.SPAN_NAME_PREFIX_TOOL + "step",
+                span_type=SpanType.SERVER,
+                attributes=attributes
+            )
+        start_time = time.time()
+        try:
+            response = await wrapped(*args, **kwargs)
+            _record_response(instance, start_time, response, attributes)
+        except Exception as e:
+            _record_exception(span=span,
+                              start_time=start_time,
+                              exception=e,
+                              attributes=attributes
+                              )
+            _end_span(span)
+            raise e
+        _end_span(span)
+        return response
+    return _async_step_wrapper
+
+
+def _step_class_wrapper(tracer: Tracer):
+    def _step_wrapper(wrapped, instance, args, kwargs):
+        span = None
+        message = args[0] or kwargs.get("message")
+        attributes = get_tool_span_attributes(instance, message)
+        if tracer:
+            span = tracer.start_span(
+                name=trace_constants.SPAN_NAME_PREFIX_TOOL + "step",
+                span_type=SpanType.SERVER,
+                attributes=attributes
+            )
+        start_time = time.time()
+        try:
+            response = wrapped(*args, **kwargs)
+            _record_response(instance, start_time, response, attributes)
+        except Exception as e:
+            _record_exception(span=span,
+                              start_time=start_time,
+                              exception=e,
+                              attributes=attributes
+                              )
+            _end_span(span)
+            raise e
+        _end_span(span)
+        return response
+    return _step_wrapper
+
+async def _async_step_instance_wrapper(tracer: Tracer):
+
+    @wrapt.decorator
+    async def _awrapper(wrapped, instance, args, kwargs):
+        wrapper_func = _async_step_class_wrapper(tracer=tracer)
+        return await wrapper_func(wrapped, instance, args, kwargs)
+
+    return _awrapper
+
+def _step_instance_wrapper(tracer: Tracer):
+
+    @wrapt.decorator
+    def _wrapper(wrapped, instance, args, kwargs):
+        wrapper_func = _step_class_wrapper(tracer)
+        return wrapper_func(wrapped, instance, args, kwargs)
+
+    return _wrapper
+
+class ToolInstrumentor(Instrumentor):
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ()
+
+    def _instrument(self, **kwargs):
+        agent_trace_enabled = kwargs.get("trace_enabled", False)
+        tracer_provider = get_tracer_provider_silent()
+        tracer = None
+        if tracer_provider and agent_trace_enabled:
+            tracer = tracer_provider.get_tracer(
+                "aworld.trace.instrumentation.tool")
+
+        wrapt.wrap_function_wrapper(
+            "aworld.core.tool.base",
+            "AsyncBaseTool.step",
+            _async_step_class_wrapper(tracer=tracer)
+        )
+
+        wrapt.wrap_function_wrapper(
+            "aworld.core.tool.base",
+            "AsyncTool.step",
+            _async_step_class_wrapper(tracer=tracer)
+        )
+
+        wrapt.wrap_function_wrapper(
+            "aworld.core.tool.base",
+            "BaseTool.step",
+            _step_class_wrapper(tracer=tracer)
+        )
+
+        wrapt.wrap_function_wrapper(
+            "aworld.core.tool.base",
+            "Tool.step",
+            _step_class_wrapper(tracer=tracer)
+        )
+
+    def _uninstrument(self, **kwargs: Any):
+        pass
+
+
+def wrap_tool(tool):
+    try:
+        tracer_provider = get_tracer_provider_silent()
+        if not tracer_provider:
+            return tool
+        tracer = tracer_provider.get_tracer(
+            "aworld.trace.instrumentation.tool")
+
+        async_wrapper = _async_step_instance_wrapper(tracer)
+        wrapper = _step_instance_wrapper(tracer)
+        if iscoroutine(tool.step):
+            tool.step = async_wrapper(tool.step)
+        else:
+            tool.step = wrapper(tool.step)
+    except Exception:
+        logger.warning(traceback.format_exc())
+
+    return tool

chapter9/gaia-experience/AWorld/aworld/trace/instrumentation/uni_llmmodel/__init__.py

@@ -0,0 +1,379 @@
+import wrapt
+import time
+import traceback
+import aworld.trace.instrumentation.semconv as semconv
+from typing import Collection, Any
+from aworld.trace.instrumentation import Instrumentor
+from aworld.trace.base import (
+    Tracer,
+    SpanType,
+    get_tracer_provider_silent
+)
+from aworld.trace.constants import ATTRIBUTES_MESSAGE_RUN_TYPE_KEY, RunType, SPAN_NAME_PREFIX_LLM
+from aworld.trace.instrumentation.llm_metrics import (
+    record_exception_metric,
+    record_chat_response_metric,
+    record_streaming_time_to_first_token,
+    record_streaming_time_to_generate
+)
+from aworld.trace.instrumentation.uni_llmmodel.model_response_parse import (
+    accumulate_stream_response,
+    get_common_attributes_from_response,
+    record_stream_token_usage,
+    parse_response_message,
+    response_to_dic,
+    handle_request
+)
+from aworld.trace.instrumentation.openai.inout_parse import run_async
+
+from aworld.models.model_response import ModelResponse
+from aworld.logs.util import logger
+
+
+def _completion_class_wrapper(tracer: Tracer):
+
+    def wrapper(wrapped, instance, args, kwargs):
+        model_name = instance.provider.model_name
+        if not model_name:
+            model_name = "LLMModel"
+        span_attributes = {}
+        span_attributes[ATTRIBUTES_MESSAGE_RUN_TYPE_KEY] = RunType.LLM.value
+
+        span = tracer.start_span(
+            name=SPAN_NAME_PREFIX_LLM + model_name, span_type=SpanType.CLIENT, attributes=span_attributes)
+
+        run_async(handle_request(span, kwargs, instance))
+        start_time = time.time()
+        try:
+            response = wrapped(*args, **kwargs)
+        except Exception as e:
+            record_exception(span=span,
+                             start_time=start_time,
+                             exception=e
+                             )
+            span.end()
+            raise e
+
+        record_completion(span=span,
+                          start_time=start_time,
+                          response=response,
+                          request_kwargs=kwargs,
+                          instance=instance,
+                          is_async=False
+                          )
+        span.end()
+        return response
+
+    return wrapper
+
+
+def _completion_instance_wrapper(tracer: Tracer):
+
+    @wrapt.decorator
+    def _wrapper(wrapped, instance, args, kwargs):
+        wrapper_func = _completion_class_wrapper(tracer)
+        return wrapper_func(wrapped, instance, args, kwargs)
+
+    return _wrapper
+
+
+def _stream_completion_class_wrapper(tracer: Tracer):
+    def wrapper(wrapped, instance, args, kwargs):
+        model_name = instance.provider.model_name
+        if not model_name:
+            model_name = "LLMModel"
+        span_attributes = {}
+        span_attributes[ATTRIBUTES_MESSAGE_RUN_TYPE_KEY] = RunType.LLM.value
+
+        span = tracer.start_span(
+            name=SPAN_NAME_PREFIX_LLM + model_name, span_type=SpanType.CLIENT, attributes=span_attributes)
+
+        run_async(handle_request(span, kwargs, instance))
+        start_time = time.time()
+        try:
+            response = wrapped(*args, **kwargs)
+        except Exception as e:
+            record_exception(span=span,
+                             start_time=start_time,
+                             exception=e
+                             )
+            span.end()
+            raise e
+        return WrappedGeneratorResponse(span=span,
+                                        response=response,
+                                        instance=instance,
+                                        start_time=start_time,
+                                        request_kwargs=kwargs
+                                        )
+    return wrapper
+
+
+def _stream_completion_instance_wrapper(tracer: Tracer):
+
+    @wrapt.decorator
+    def _stream_wrapper(wrapped, instance, args, kwargs):
+        wrapper_func = _stream_completion_class_wrapper(tracer)
+        return wrapper_func(wrapped, instance, args, kwargs)
+
+    return _stream_wrapper
+
+
+def _acompletion_class_wrapper(tracer: Tracer):
+
+    async def awrapper(wrapped, instance, args, kwargs):
+        model_name = instance.provider.model_name
+        if not model_name:
+            model_name = "LLMModel"
+        span_attributes = {}
+        span_attributes[ATTRIBUTES_MESSAGE_RUN_TYPE_KEY] = RunType.LLM.value
+
+        span = tracer.start_span(
+            name=SPAN_NAME_PREFIX_LLM + model_name, span_type=SpanType.CLIENT, attributes=span_attributes)
+
+        await handle_request(span, kwargs, instance)
+        start_time = time.time()
+        try:
+            response = await wrapped(*args, **kwargs)
+        except Exception as e:
+            record_exception(span=span,
+                             start_time=start_time,
+                             exception=e
+                             )
+            span.end()
+            raise e
+
+        record_completion(span=span,
+                          start_time=start_time,
+                          response=response,
+                          request_kwargs=kwargs,
+                          instance=instance,
+                          is_async=True
+                          )
+        span.end()
+        return response
+
+    return awrapper
+
+
+async def _acompletion_instance_wrapper(tracer: Tracer):
+
+    @wrapt.decorator
+    async def _awrapper(wrapped, instance, args, kwargs):
+        wrapper_func = _acompletion_class_wrapper(tracer)
+        return await wrapper_func(wrapped, instance, args, kwargs)
+
+    return _awrapper
+
+
+def record_exception(span, start_time, exception):
+    '''
+    record openai chat exception to trace and metrics
+    '''
+    try:
+        duration = time.time() - start_time if "start_time" in locals() else 0
+        if span.is_recording:
+            span.record_exception(exception=exception)
+        record_exception_metric(exception=exception, duration=duration)
+    except Exception as e:
+        logger.warning(f"openai instrument record exception error.{e}")
+
+
+def record_completion(span,
+                      start_time,
+                      response,
+                      request_kwargs,
+                      instance,
+                      is_async):
+    '''
+    Record chat completion to trace and metrics
+    '''
+    duration = time.time() - start_time if "start_time" in locals() else 0
+    response_dict = response_to_dic(response)
+    attributes = get_common_attributes_from_response(instance, is_async, False)
+    usage = response_dict.get("usage")
+    content = response_dict.get("content", "")
+    tool_calls = response_dict.get("tool_calls")
+    prompt_tokens = -1
+    completion_tokens = -1
+    total_tokens = -1
+    if usage:
+        prompt_tokens = usage.get("prompt_tokens")
+        completion_tokens = usage.get("completion_tokens")
+        total_tokens = usage.get("total_tokens")
+
+    span_attributes = {
+        **attributes,
+        semconv.GEN_AI_USAGE_INPUT_TOKENS: prompt_tokens,
+        semconv.GEN_AI_USAGE_OUTPUT_TOKENS: completion_tokens,
+        semconv.GEN_AI_USAGE_TOTAL_TOKENS: total_tokens,
+        semconv.GEN_AI_DURATION: duration,
+        semconv.GEN_AI_COMPLETION_CONTENT: content
+    }
+    span_attributes.update(parse_response_message(tool_calls))
+    span.set_attributes(span_attributes)
+    record_chat_response_metric(attributes=attributes,
+                                prompt_tokens=prompt_tokens,
+                                completion_tokens=completion_tokens,
+                                duration=duration
+                                )
+
+
+class WrappedGeneratorResponse():
+
+    def __init__(
+        self,
+        span,
+        response,
+        instance=None,
+        start_time=None,
+        request_kwargs=None
+    ):
+        self._span = span
+        self._response = response
+        self._instance = instance
+        self._start_time = start_time
+        self._complete_response = {
+            "id": "", "model": "", "content": "", "tool_calls": [], "usage": {}}
+        self._first_token_recorded = False
+        self._time_of_first_token = None
+        self._request_kwargs = request_kwargs
+
+    def __iter__(self):
+        return self
+
+    def __aiter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            chunk = self._response.__next__()
+        except Exception as e:
+            if isinstance(e, StopIteration):
+                self._close_span(False)
+            raise e
+        else:
+            self._process_stream_chunk(chunk, False)
+            return chunk
+
+    async def __anext__(self):
+        try:
+            chunk = await self._response.__anext__()
+        except Exception as e:
+            if isinstance(e, StopAsyncIteration):
+                self._close_span(True)
+            raise e
+        else:
+            self._process_stream_chunk(chunk, True)
+            return chunk
+
+    def _process_stream_chunk(self, chunk: ModelResponse, is_async):
+        accumulate_stream_response(chunk, self._complete_response)
+
+        if not self._first_token_recorded:
+            self._time_of_first_token = time.time()
+            duration = self._time_of_first_token - self._start_time
+            attribute = get_common_attributes_from_response(
+                self._instance, is_async, True)
+            record_streaming_time_to_first_token(duration, attribute)
+            self._first_token_recorded = True
+
+    def _close_span(self, is_async):
+        duration = None
+        first_token_duration = None
+        first_token_to_generate_duration = None
+        if self._start_time and isinstance(self._start_time, (float, int)):
+            duration = time.time() - self._start_time
+        if self._time_of_first_token and self._start_time and isinstance(self._start_time, (float, int)):
+            first_token_duration = self._time_of_first_token - self._start_time
+            first_token_to_generate_duration = time.time() - self._time_of_first_token
+
+        prompt_usage, completion_usage = record_stream_token_usage(
+            self._complete_response, self._request_kwargs)
+
+        attributes = get_common_attributes_from_response(
+            self._instance, is_async, True)
+
+        span_attributes = {
+            **attributes,
+            semconv.GEN_AI_USAGE_INPUT_TOKENS: prompt_usage,
+            semconv.GEN_AI_USAGE_OUTPUT_TOKENS: completion_usage,
+            semconv.GEN_AI_USAGE_TOTAL_TOKENS: prompt_usage + completion_usage,
+            semconv.GEN_AI_DURATION: duration,
+            semconv.GEN_AI_FIRST_TOKEN_DURATION: first_token_duration,
+            semconv.GEN_AI_COMPLETION_CONTENT: self._complete_response.get(
+                "content", "")
+        }
+        span_attributes.update(parse_response_message(
+            self._complete_response.get("tool_calls", [])))
+
+        self._span.set_attributes(span_attributes)
+        record_chat_response_metric(attributes=attributes,
+                                    prompt_tokens=prompt_usage,
+                                    completion_tokens=completion_usage,
+                                    duration=duration
+                                    )
+        record_streaming_time_to_generate(
+            first_token_to_generate_duration, attributes)
+
+        self._span.end()
+
+
+class LLMModelInstrumentor(Instrumentor):
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ()
+
+    def _instrument(self, **kwargs):
+        tracer_provider = get_tracer_provider_silent()
+        if not tracer_provider:
+            return
+        tracer = tracer_provider.get_tracer(
+            "aworld.trace.instrumentation.llmmodel")
+
+        wrapt.wrap_function_wrapper(
+            "aworld.models.llm",
+            "LLMModel.completion",
+            _completion_class_wrapper(tracer=tracer)
+        )
+
+        wrapt.wrap_function_wrapper(
+            "aworld.models.llm",
+            "LLMModel.stream_completion",
+            _stream_completion_class_wrapper(tracer=tracer)
+        )
+        wrapt.wrap_function_wrapper(
+            "aworld.models.llm",
+            "LLMModel.acompletion",
+            _acompletion_class_wrapper(tracer)
+        )
+
+        wrapt.wrap_function_wrapper(
+            "aworld.models.llm",
+            "LLMModel.astream_completion",
+            _stream_completion_class_wrapper(tracer)
+        )
+        logger.info(f"LLMModelInstrumentor wrap aworld.models.llm")
+
+    def _uninstrument(self, **kwargs: Any):
+        pass
+
+
+def wrap_llmmodel(client: 'aworld.models.llm.LLMModel'):
+    try:
+        tracer_provider = get_tracer_provider_silent()
+        if not tracer_provider:
+            return client
+        tracer = tracer_provider.get_tracer(
+            "aworld.trace.instrumentation.llmmodel")
+
+        wrapper = _completion_instance_wrapper(tracer)
+        awrapper = _acompletion_instance_wrapper(tracer)
+        stream_wrapper = _stream_completion_instance_wrapper(tracer)
+        client.completion = wrapper(client.completion)
+        client.stream_completion = stream_wrapper(client.stream_completion)
+        client.acompletion = awrapper(client.acompletion)
+        client.astream_completion = stream_wrapper(client.astream_completion)
+    except Exception:
+        logger.warning(traceback.format_exc())
+
+    return client

chapter9/gaia-experience/AWorld/aworld/trace/instrumentation/uni_llmmodel/model_response_parse.py

@@ -0,0 +1,183 @@
+import copy
+import json
+import aworld.trace.instrumentation.semconv as semconv
+from aworld.models.model_response import ModelResponse, ToolCall
+from aworld.trace.base import Span
+from aworld.trace.instrumentation.openai.inout_parse import should_trace_prompts, need_flatten_messages
+from aworld.logs.util import logger
+from aworld.utils.serialized_util import to_serializable
+
+
+def parser_request_params(kwargs, instance: 'aworld.models.llm.LLMModel'):
+    attributes = {
+        semconv.GEN_AI_SYSTEM: instance.provider_name,
+        semconv.GEN_AI_REQUEST_MODEL: instance.provider.model_name,
+        semconv.GEN_AI_REQUEST_MAX_TOKENS: kwargs.get("max_tokens", ""),
+        semconv.GEN_AI_REQUEST_TEMPERATURE: kwargs.get("temperature", ""),
+        semconv.GEN_AI_REQUEST_STOP_SEQUENCES: str(kwargs.get("stop", [])),
+        semconv.GEN_AI_REQUEST_FREQUENCY_PENALTY: kwargs.get("frequency_penalty", ""),
+        semconv.GEN_AI_REQUEST_PRESENCE_PENALTY: kwargs.get("presence_penalty", ""),
+        semconv.GEN_AI_REQUEST_USER: kwargs.get("user", ""),
+        semconv.GEN_AI_REQUEST_EXTRA_HEADERS: kwargs.get("extra_headers", ""),
+        semconv.GEN_AI_REQUEST_STREAMING: kwargs.get("stream", ""),
+        semconv.GEN_AI_REQUEST_TOP_P: kwargs.get("top_p", ""),
+        semconv.GEN_AI_OPERATION_NAME: "chat"
+    }
+    return attributes
+
+
+async def handle_request(span: Span, kwargs, instance):
+    if not span or not span.is_recording():
+        return
+    try:
+        attributes = parser_request_params(kwargs, instance)
+        if should_trace_prompts():
+            messages = kwargs.get("messages")
+            if need_flatten_messages():
+                attributes.update(parse_request_message(messages))
+            else:
+                attributes.update({
+                    semconv.GEN_AI_PROMPT: covert_to_jsonstr(messages)
+                })
+        tools = kwargs.get("tools")
+        if tools:
+            if need_flatten_messages():
+                attributes.update(parse_prompt_tools(tools))
+            else:
+                attributes.update({
+                    semconv.GEN_AI_PROMPT_TOOLS: covert_to_jsonstr(tools)
+                })
+
+        filterd_attri = {k: v for k, v in attributes.items()
+                         if (v and v != "")}
+
+        span.set_attributes(filterd_attri)
+    except Exception as e:
+        logger.warning(f"trace handle openai request error: {e}")
+
+
+def get_common_attributes_from_response(instance: 'LLMModel', is_async, is_streaming):
+    operation = "acompletion" if is_async else "completion"
+    if is_streaming:
+        operation = "astream_completion" if is_async else "stream_completion"
+    return {
+        semconv.GEN_AI_SYSTEM: instance.provider_name,
+        semconv.GEN_AI_RESPONSE_MODEL: instance.provider.model_name,
+        semconv.GEN_AI_METHOD_NAME: operation,
+        semconv.GEN_AI_SERVER_ADDRESS: instance.provider.base_url
+    }
+
+
+def accumulate_stream_response(chunk: ModelResponse, complete_response: dict):
+    from aworld.utils.common import nest_dict_counter
+    # logger.info(f"accumulate_stream_response chunk= {chunk}")
+
+    complete_response["model"] = chunk.model
+    complete_response["id"] = chunk.id
+    if chunk.content:
+        complete_response["content"] += chunk.content
+    if chunk.tool_calls:
+        complete_response["tool_calls"].extend(chunk.tool_calls)
+    if chunk.error:
+        complete_response["error"] = chunk.error
+    complete_response["usage"] = nest_dict_counter(
+        complete_response["usage"], chunk.usage)
+
+
+def record_stream_token_usage(complete_response, request_kwargs) -> tuple[int, int]:
+    '''
+        return (prompt_usage, completion_usage)
+    '''
+    # logger.info(
+    #     f"record_stream_token_usage complete_response= {complete_response}")
+    usage = complete_response.get("usage", {})
+    if usage:
+        prompt_tokens = usage.get("prompt_tokens", 0)
+        completion_tokens = usage.get("completion_tokens", 0)
+        return (prompt_tokens, completion_tokens)
+    return (0, 0)
+
+
+def parse_request_message(messages):
+    '''
+    flatten request message to attributes
+    '''
+    attributes = {}
+    for i, msg in enumerate(messages):
+        prefix = f"{semconv.GEN_AI_PROMPT}.{i}"
+        attributes.update({f"{prefix}.role": msg.get("role")})
+        if msg.get("content"):
+            content = copy.deepcopy(msg.get("content"))
+            content = json.dumps(content, ensure_ascii=False)
+            attributes.update({f"{prefix}.content": content})
+        if msg.get("tool_call_id"):
+            attributes.update({
+                f"{prefix}.tool_call_id": msg.get("tool_call_id")})
+        tool_calls = msg.get("tool_calls")
+        # logger.info(f"input tool_calls={tool_calls}")
+        if tool_calls:
+            for i, tool_call in enumerate(tool_calls):
+                if isinstance(tool_call, dict):
+                    function = tool_call.get('function')
+                    attributes.update({
+                        f"{prefix}.tool_calls.{i}.id": tool_call.get("id")})
+                    attributes.update({
+                        f"{prefix}.tool_calls.{i}.name": function.get("name")})
+                    attributes.update({
+                        f"{prefix}.tool_calls.{i}.arguments": function.get("arguments")})
+                elif isinstance(tool_call, ToolCall):
+                    function = tool_call.function
+                    attributes.update({
+                        f"{prefix}.tool_calls.{i}.id": tool_call.id})
+                    attributes.update({
+                        f"{prefix}.tool_calls.{i}.name": function.name})
+                    attributes.update({
+                        f"{prefix}.tool_calls.{i}.arguments": function.arguments})
+    return attributes
+
+
+def parse_prompt_tools(tools):
+    attributes = {}
+    for i, tool in enumerate(tools):
+        prefix = f"{semconv.GEN_AI_PROMPT_TOOLS}.{i}"
+        if isinstance(tool, dict):
+            tool_type = tool.get("type")
+            attributes.update({
+                f"{prefix}.type": tool_type})
+            if tool.get(tool_type):
+                attributes.update({
+                    f"{prefix}.name": tool.get(tool_type).get("name")})
+    return attributes
+
+
+def parse_response_message(tool_calls) -> dict:
+    attributes = {}
+    prefix = semconv.GEN_AI_COMPLETION_TOOL_CALLS
+    if tool_calls:
+        if need_flatten_messages():
+            for i, tool_call in enumerate(tool_calls):
+                function = tool_call.get("function")
+                attributes.update(
+                    {f"{prefix}.{i}.id": tool_call.get("id")})
+                attributes.update(
+                    {f"{prefix}.{i}.name": function.get("name")})
+                attributes.update(
+                    {f"{prefix}.{i}.arguments": function.get("arguments")})
+        else:
+            attributes.update({
+                prefix: covert_to_jsonstr(tool_calls)
+            })
+    return attributes
+
+
+def response_to_dic(response: ModelResponse) -> dict:
+    logger.info(f"completion response= {response}")
+    return response.to_dict()
+
+
+def covert_to_jsonstr(obj):
+    try:
+        return json.dumps(to_serializable(obj), ensure_ascii=False)
+    except:
+        logger.warning(f"covert_to_jsonstr error: {obj.__class__.__name__}")
+        return str(obj)

chapter9/gaia-experience/AWorld/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__)