| 123456789101112131415161718192021222324252627282930313233343536373839404142 |
- """Middleware for making it easier to do distributed tracing."""
- class TracingMiddleware:
- """Middleware for propagating distributed tracing context using LangSmith.
- This middleware checks for the `'langsmith-trace'` header and propagates the
- tracing context if present. It does not start new traces by default.
- Designed to work with ASGI applications.
- Attributes:
- app: The ASGI application being wrapped.
- """
- def __init__(self, app):
- """Initialize the middleware."""
- from langsmith.run_helpers import tracing_context # type: ignore
- self._with_headers = tracing_context
- self.app = app
- async def __call__(self, scope: dict, receive, send):
- """Handle incoming requests and propagate tracing context if applicable.
- Args:
- scope: A dict containing ASGI connection scope.
- receive (callable): An awaitable callable for receiving ASGI events.
- send (callable): An awaitable callable for sending ASGI events.
- If the request is HTTP and contains the `'langsmith-trace'` header,
- it propagates the tracing context before calling the wrapped application.
- Otherwise, it calls the application directly without modifying the context.
- """
- if scope["type"] == "http" and "headers" in scope:
- headers = dict(scope["headers"])
- if b"langsmith-trace" in headers:
- with self._with_headers(parent=headers):
- await self.app(scope, receive, send)
- return
- await self.app(scope, receive, send)
|