codekingpro/portable-devtools

1"""Middleware for making it easier to do distributed tracing."""2 3 4class TracingMiddleware:5    """Middleware for propagating distributed tracing context using LangSmith.6 7    This middleware checks for the `'langsmith-trace'` header and propagates the8    tracing context if present. It does not start new traces by default.9 10    Designed to work with ASGI applications.11 12    Attributes:13        app: The ASGI application being wrapped.14    """15 16    def __init__(self, app):17        """Initialize the middleware."""18        from langsmith.run_helpers import tracing_context  # type: ignore19 20        self._with_headers = tracing_context21        self.app = app22 23    async def __call__(self, scope: dict, receive, send):24        """Handle incoming requests and propagate tracing context if applicable.25 26        Args:27            scope: A dict containing ASGI connection scope.28            receive (callable): An awaitable callable for receiving ASGI events.29            send (callable): An awaitable callable for sending ASGI events.30 31        If the request is HTTP and contains the `'langsmith-trace'` header,32        it propagates the tracing context before calling the wrapped application.33 34        Otherwise, it calls the application directly without modifying the context.35        """36        if scope["type"] == "http" and "headers" in scope:37            headers = dict(scope["headers"])38            if b"langsmith-trace" in headers:39                with self._with_headers(parent=headers):40                    await self.app(scope, receive, send)41                return42        await self.app(scope, receive, send)43