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

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()