Coverage for langsmith/middleware.py: 23%

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 the 

8 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: ignore 

19 

20 self._with_headers = tracing_context 

21 self.app = app 

22 

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 return 

42 await self.app(scope, receive, send)