feat: Instrumentation reads session id from metadata (#446)

Co-authored-by: Xander Song <axiomofjoy@gmail.com>
Arize-ai · May 10, 2024 · 5490f68 · 5490f68
1 parent 108c8bd
commit 5490f68
Show file tree

Hide file tree

Showing 5 changed files with 523 additions and 45 deletions.
diff --git a/python/instrumentation/openinference-instrumentation-langchain/examples/chain_metadata.py b/python/instrumentation/openinference-instrumentation-langchain/examples/chain_metadata.py
@@ -30,6 +30,12 @@
                 "key-2": "val-2",
             },
         },
+        prompt_template="Who won the soccer match in {city} on {date}",
+        prompt_template_version="v1.0",
+        prompt_template_variables={
+            "city": "Johannesburg",
+            "date": "July 11th",
+        },
         tags=["tag-1", "tag-2"],
     ):
         prompt_template = "Tell me a {adjective} joke"

diff --git a/...on/instrumentation/openinference-instrumentation-langchain/examples/openai_chat_stream.py b/...on/instrumentation/openinference-instrumentation-langchain/examples/openai_chat_stream.py
@@ -29,6 +29,12 @@
             },
         },
         tags=["tag-1", "tag-2"],
+        prompt_template="Who won the soccer match in {city} on {date}",
+        prompt_template_version="v1.0",
+        prompt_template_variables={
+            "city": "Johannesburg",
+            "date": "July 11th",
+        },
     ):
         for chunk in ChatOpenAI(model_name="gpt-3.5-turbo").stream("Write a haiku."):
             print(chunk.content, end="", flush=True)
diff --git a/python/instrumentation/openinference-instrumentation-langchain/pyproject.toml b/python/instrumentation/openinference-instrumentation-langchain/pyproject.toml
@@ -28,7 +28,7 @@ dependencies = [
   "opentelemetry-api",
   "opentelemetry-instrumentation",
   "opentelemetry-semantic-conventions",
-  "openinference-instrumentation>=0.1.3",
+  "openinference-instrumentation>=0.1.5",
   "openinference-semantic-conventions",
   "wrapt",
 ]

diff --git a/...nference-instrumentation-langchain/src/openinference/instrumentation/langchain/_tracer.py b/...nference-instrumentation-langchain/src/openinference/instrumentation/langchain/_tracer.py
@@ -192,6 +192,7 @@ def _update_span(span: trace_api.Span, run: Run) -> None:
         else _langchain_run_type_to_span_kind(run.run_type)
     )
     span.set_attribute(OPENINFERENCE_SPAN_KIND, span_kind.value)
+    span.set_attributes(dict(get_attributes_from_context()))
     span.set_attributes(
         dict(
             _flatten(
@@ -213,7 +214,6 @@ def _update_span(span: trace_api.Span, run: Run) -> None:
             )
         )
     )
-    span.set_attributes(dict(get_attributes_from_context()))
 
 
 def _langchain_run_type_to_span_kind(run_type: str) -> OpenInferenceSpanKindValues:
@@ -569,6 +569,12 @@ def _metadata(run: Run) -> Iterator[Tuple[str, str]]:
     if not run.extra or not (metadata := run.extra.get("metadata")):
         return
     assert isinstance(metadata, Mapping), f"expected Mapping, found {type(metadata)}"
+    if session_id := (
+        metadata.get(LANGCHAIN_SESSION_ID)
+        or metadata.get(LANGCHAIN_CONVERSATION_ID)
+        or metadata.get(LANGCHAIN_THREAD_ID)
+    ):
+        yield SESSION_ID, session_id
     yield METADATA, json.dumps(metadata)
 
 
@@ -599,6 +605,10 @@ def _as_utc_nano(dt: datetime) -> int:
     return int(dt.astimezone(timezone.utc).timestamp() * 1_000_000_000)
 
 
+LANGCHAIN_SESSION_ID = "session_id"
+LANGCHAIN_CONVERSATION_ID = "conversation_id"
+LANGCHAIN_THREAD_ID = "thread_id"
+
 DOCUMENT_CONTENT = DocumentAttributes.DOCUMENT_CONTENT
 DOCUMENT_ID = DocumentAttributes.DOCUMENT_ID
 DOCUMENT_METADATA = DocumentAttributes.DOCUMENT_METADATA
@@ -641,3 +651,4 @@ def _as_utc_nano(dt: datetime) -> int:
 TOOL_NAME = SpanAttributes.TOOL_NAME
 TOOL_PARAMETERS = SpanAttributes.TOOL_PARAMETERS
 METADATA = SpanAttributes.METADATA
+SESSION_ID = SpanAttributes.SESSION_ID