From 19f879cf636d2aaf91c0445ad1733bca3ec11c4a Mon Sep 17 00:00:00 2001
From: Eugene Yurtsev <eyurtsev@gmail.com>
Date: Fri, 22 Mar 2024 14:32:08 -0400
Subject: [PATCH] x

---
 backend/server/main.py       |  4 ++++
 backend/server/middleware.py | 31 +++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+)
 create mode 100644 backend/server/middleware.py

diff --git a/backend/server/main.py b/backend/server/main.py
index 7b2839f..78b675d 100644
--- a/backend/server/main.py
+++ b/backend/server/main.py
@@ -14,6 +14,7 @@
     ExtractResponse,
     extraction_runnable,
 )
+from server.middleware import add_langsmith_tracing
 
 logger = logging.getLogger(__name__)
 
@@ -30,6 +31,9 @@
 )
 
 
+add_langsmith_tracing = app.middleware("http")(add_langsmith_tracing)
+
+
 ROOT = Path(__file__).parent.parent
 
 ORIGINS = os.environ.get("CORS_ORIGINS", "").split(",")
diff --git a/backend/server/middleware.py b/backend/server/middleware.py
new file mode 100644
index 0000000..db19a96
--- /dev/null
+++ b/backend/server/middleware.py
@@ -0,0 +1,31 @@
+"""Add example middleware to enable langsmith tracing on selected endpoints.
+
+This enables langsmith tracing on selected endpoints, showing how to add
+tags and metadata to the traces.
+"""
+from urllib.parse import urlparse
+
+from langchain_core.tracers.context import tracing_v2_enabled
+from langsmith import trace
+from starlette.requests import Request
+
+
+async def add_langsmith_tracing(request: Request, call_next):
+    """Add langsmith middleware"""
+    with tracing_v2_enabled():
+        parsed_url = urlparse(str(request.url))
+        if parsed_url.path not in {"/suggest", "/extract", "/extract/shared"}:
+            # Skip tracing for all other endpoints
+            return await call_next(request)
+
+        method = request.method
+
+        with trace(
+            name=parsed_url.path,
+            tags=[method],
+            metadata={
+                "__useragent": request.headers.get("user-agent"),
+            },
+        ):
+            response = await call_next(request)
+    return response