From af590566040655b04d9041e513f34d4f754d1dc8 Mon Sep 17 00:00:00 2001 From: anovazzi1 Date: Mon, 17 Feb 2025 11:50:14 -0300 Subject: [PATCH] feat: Add session ID support to tracing service (#5820) * feat: Add session ID support to tracing service This commit adds support for setting the session ID in the tracing service. Two methods have been modified: - In the `Graph` class, the `set_session_id` method has been added to set the ID of the current session. - In the `ArizePhoenixTracer` class, the `__init__` method has been modified to accept a `session_id` parameter and store it. These changes enable the tracing service to associate traces with specific sessions, improving trace management and analysis. * refactor: Improve tracing session ID handling in Graph class - Rename `set_session_id()` method to `set_tracing_session_id()` - Add null check for tracing service when setting session ID - Enhance method to only set session ID when tracing service is available --------- Co-authored-by: Gabriel Luiz Freitas Almeida --- src/backend/base/langflow/graph/graph/base.py | 11 +++++++++++ .../base/langflow/services/tracing/arize_phoenix.py | 7 +++++-- src/backend/base/langflow/services/tracing/service.py | 6 ++++++ 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/src/backend/base/langflow/graph/graph/base.py b/src/backend/base/langflow/graph/graph/base.py index 557f0d373..6a6a4b422 100644 --- a/src/backend/base/langflow/graph/graph/base.py +++ b/src/backend/base/langflow/graph/graph/base.py @@ -635,6 +635,15 @@ class Graph: raise ValueError(msg) return self._run_id + def set_tracing_session_id(self) -> None: + """Sets the ID of the current session. + + Args: + session_id (str): The session ID. + """ + if self.tracing_service: + self.tracing_service.set_session_id(self._session_id) + def set_run_id(self, run_id: uuid.UUID | None = None) -> None: """Sets the ID of the current run. @@ -647,6 +656,8 @@ class Graph: self._run_id = str(run_id) if self.tracing_service: self.tracing_service.set_run_id(run_id) + if self._session_id and self.tracing_service is not None: + self.tracing_service.set_session_id(self.session_id) def set_run_name(self) -> None: # Given a flow name, flow_id diff --git a/src/backend/base/langflow/services/tracing/arize_phoenix.py b/src/backend/base/langflow/services/tracing/arize_phoenix.py index 84e9273b2..a65a84719 100644 --- a/src/backend/base/langflow/services/tracing/arize_phoenix.py +++ b/src/backend/base/langflow/services/tracing/arize_phoenix.py @@ -39,7 +39,9 @@ class ArizePhoenixTracer(BaseTracer): chat_input_value: str chat_output_value: str - def __init__(self, trace_name: str, trace_type: str, project_name: str, trace_id: UUID): + def __init__( + self, trace_name: str, trace_type: str, project_name: str, trace_id: UUID, session_id: str | None = None + ): """Initializes the ArizePhoenixTracer instance and sets up a root span.""" self.trace_name = trace_name self.trace_type = trace_type @@ -49,6 +51,7 @@ class ArizePhoenixTracer(BaseTracer): self.flow_id = trace_name.split(" - ")[-1] self.chat_input_value = "" self.chat_output_value = "" + self.session_id = session_id try: self._ready = self.setup_arize_phoenix() @@ -63,7 +66,7 @@ class ArizePhoenixTracer(BaseTracer): name=self.flow_id, start_time=self._get_current_timestamp(), ) - self.root_span.set_attribute(SpanAttributes.SESSION_ID, self.flow_id) + self.root_span.set_attribute(SpanAttributes.SESSION_ID, self.session_id or self.flow_id) self.root_span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, self.trace_type) self.root_span.set_attribute("langflow.project.name", self.project_name) self.root_span.set_attribute("langflow.flow.name", self.flow_name) diff --git a/src/backend/base/langflow/services/tracing/service.py b/src/backend/base/langflow/services/tracing/service.py index 4742fb0ed..65f16f9ab 100644 --- a/src/backend/base/langflow/services/tracing/service.py +++ b/src/backend/base/langflow/services/tracing/service.py @@ -65,6 +65,7 @@ class TracingService(Service): self.worker_task: asyncio.Task | None = None self.end_trace_tasks: set[asyncio.Task] = set() self.deactivated = self.settings_service.settings.deactivate_tracing + self.session_id: str | None = None async def log_worker(self) -> None: while self.running or not self.logs_queue.empty(): @@ -163,6 +164,7 @@ class TracingService(Service): trace_type="chain", project_name=self.project_name, trace_id=self.run_id, + session_id=self.session_id, ) def set_run_name(self, name: str) -> None: @@ -286,3 +288,7 @@ class TracingService(Service): if langchain_callback: callbacks.append(langchain_callback) return callbacks + + def set_session_id(self, session_id: str) -> None: + """Set the session ID for tracing.""" + self.session_id = session_id