middleware.py 1.6 KB

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