feat: add fetch_sessions() (langfuse#794)

marcklingen · web-flow · commit 8c7547e5f77f · 2024-07-10T19:20:32.000+02:00
diff --git a/langfuse/client.py b/langfuse/client.py
@@ -27,6 +27,7 @@
     DatasetRunWithItems,
 )
 from langfuse.api.resources.commons.types.observations_view import ObservationsView
+from langfuse.api.resources.commons.types.session import Session
 from langfuse.api.resources.commons.types.trace_with_details import TraceWithDetails
 from langfuse.api.resources.datasets.types.paginated_dataset_runs import (
     PaginatedDatasetRuns,
@@ -117,6 +118,14 @@ class FetchObservationResponse:
     data: Observation
 
 
+@dataclass
+class FetchSessionsResponse:
+    """Response object for fetch_sessions method."""
+
+    data: typing.List[Session]
+    meta: MetaResponse
+
+
 class Langfuse(object):
     """Langfuse Python client.
 
@@ -584,8 +593,8 @@ def fetch_traces(
             name (Optional[str]): Filter by name of traces. Defaults to None.
             user_id (Optional[str]): Filter by user_id. Defaults to None.
             session_id (Optional[str]): Filter by session_id. Defaults to None.
-            from_timestamp (Optional[dt.datetime]): Retrieve only traces with a timestamp on or after this datetime (ISO 8601). Defaults to None.
-            to_timestamp (Optional[dt.datetime]): Retrieve only traces with a timestamp before this datetime (ISO 8601). Defaults to None.
+            from_timestamp (Optional[dt.datetime]): Retrieve only traces with a timestamp on or after this datetime. Defaults to None.
+            to_timestamp (Optional[dt.datetime]): Retrieve only traces with a timestamp before this datetime. Defaults to None.
             order_by (Optional[str]): Format of the string `[field].[asc/desc]`. Fields: id, timestamp, name, userId, release, version, public, bookmarked, sessionId. Example: `timestamp.asc`. Defaults to None.
             tags (Optional[Union[str, Sequence[str]]]): Filter by tags. Defaults to None.
 
@@ -636,8 +645,8 @@ def get_traces(
             name (Optional[str]): Filter by name of traces. Defaults to None.
             user_id (Optional[str]): Filter by user_id. Defaults to None.
             session_id (Optional[str]): Filter by session_id. Defaults to None.
-            from_timestamp (Optional[dt.datetime]): Retrieve only traces with a timestamp on or after this datetime (ISO 8601). Defaults to None.
-            to_timestamp (Optional[dt.datetime]): Retrieve only traces with a timestamp before this datetime (ISO 8601). Defaults to None.
+            from_timestamp (Optional[dt.datetime]): Retrieve only traces with a timestamp on or after this datetime. Defaults to None.
+            to_timestamp (Optional[dt.datetime]): Retrieve only traces with a timestamp before this datetime. Defaults to None.
             order_by (Optional[str]): Format of the string `[field].[asc/desc]`. Fields: id, timestamp, name, userId, release, version, public, bookmarked, sessionId. Example: `timestamp.asc`. Defaults to None.
             tags (Optional[Union[str, Sequence[str]]]): Filter by tags. Defaults to None.
 
@@ -692,8 +701,8 @@ def fetch_observations(
             user_id (Optional[str]): User identifier. Defaults to None.
             trace_id (Optional[str]): Trace identifier. Defaults to None.
             parent_observation_id (Optional[str]): Parent observation identifier. Defaults to None.
-            from_start_time (Optional[dt.datetime]): Retrieve only observations with a start_time on or after this datetime (ISO 8601). Defaults to None.
-            to_start_time (Optional[dt.datetime]): Retrieve only observations with a start_time before this datetime (ISO 8601). Defaults to None.
+            from_start_time (Optional[dt.datetime]): Retrieve only observations with a start_time on or after this datetime. Defaults to None.
+            to_start_time (Optional[dt.datetime]): Retrieve only observations with a start_time before this datetime. Defaults to None.
             type (Optional[str]): Type of the observation. Defaults to None.
 
         Returns:
@@ -744,8 +753,8 @@ def get_observations(
             user_id (Optional[str]): User identifier. Defaults to None.
             trace_id (Optional[str]): Trace identifier. Defaults to None.
             parent_observation_id (Optional[str]): Parent observation identifier. Defaults to None.
-            from_start_time (Optional[dt.datetime]): Retrieve only observations with a start_time on or after this datetime (ISO 8601). Defaults to None.
-            to_start_time (Optional[dt.datetime]): Retrieve only observations with a start_time before this datetime (ISO 8601). Defaults to None.
+            from_start_time (Optional[dt.datetime]): Retrieve only observations with a start_time on or after this datetime. Defaults to None.
+            to_start_time (Optional[dt.datetime]): Retrieve only observations with a start_time before this datetime. Defaults to None.
             type (Optional[str]): Type of the observation. Defaults to None.
 
         Returns:
@@ -797,8 +806,8 @@ def get_generations(
             name (Optional[str]): Name of the generations to return. Defaults to None.
             user_id (Optional[str]): User identifier of the generations to return. Defaults to None.
             trace_id (Optional[str]): Trace identifier of the generations to return. Defaults to None.
-            from_start_time (Optional[dt.datetime]): Retrieve only observations with a start_time on or after this datetime (ISO 8601). Defaults to None.
-            to_start_time (Optional[dt.datetime]): Retrieve only observations with a start_time before this datetime (ISO 8601). Defaults to None.
+            from_start_time (Optional[dt.datetime]): Retrieve only observations with a start_time on or after this datetime. Defaults to None.
+            to_start_time (Optional[dt.datetime]): Retrieve only observations with a start_time before this datetime. Defaults to None.
             parent_observation_id (Optional[str]): Parent observation identifier of the generations to return. Defaults to None.
 
         Returns:
@@ -869,6 +878,43 @@ def get_observation(
             self.log.exception(e)
             raise e
 
+    def fetch_sessions(
+        self,
+        *,
+        page: typing.Optional[int] = None,
+        limit: typing.Optional[int] = None,
+        from_timestamp: typing.Optional[dt.datetime] = None,
+        to_timestamp: typing.Optional[dt.datetime] = None,
+    ) -> FetchSessionsResponse:
+        """Get a list of sessions in the current project.
+
+        Args:
+            page (Optional[int]): Page number of the sessions to return. Defaults to None.
+            limit (Optional[int]): Maximum number of sessions to return. Defaults to None.
+            from_timestamp (Optional[dt.datetime]): Retrieve only sessions with a timestamp on or after this datetime. Defaults to None.
+            to_timestamp (Optional[dt.datetime]): Retrieve only sessions with a timestamp before this datetime. Defaults to None.
+
+        Returns:
+            FetchSessionsResponse, list of sessions on `data` and metadata on `meta`.
+
+        Raises:
+            Exception: If an error occurred during the request.
+        """
+        try:
+            self.log.debug(
+                f"Getting sessions... {page}, {limit}, {from_timestamp}, {to_timestamp}"
+            )
+            res = self.client.sessions.list(
+                page=page,
+                limit=limit,
+                from_timestamp=from_timestamp,
+                to_timestamp=to_timestamp,
+            )
+            return FetchSessionsResponse(data=res.data, meta=res.meta)
+        except Exception as e:
+            self.log.exception(e)
+            raise e
+
     @overload
     def get_prompt(
         self,
diff --git a/tests/test_core_sdk.py b/tests/test_core_sdk.py
@@ -6,6 +6,7 @@
 from langfuse.client import (
     FetchObservationResponse,
     FetchObservationsResponse,
+    FetchSessionsResponse,
     FetchTraceResponse,
     FetchTracesResponse,
 )
@@ -1401,3 +1402,34 @@ def test_fetch_observations_empty():
     assert isinstance(response, FetchObservationsResponse)
     assert len(response.data) == 0
     assert response.meta.total_items == 0
+
+
+def test_fetch_sessions():
+    langfuse = Langfuse()
+
+    # unique name
+    name = create_uuid()
+    session1 = create_uuid()
+    session2 = create_uuid()
+    session3 = create_uuid()
+
+    # Create multiple traces
+    langfuse.trace(name=name, session_id=session1)
+    langfuse.trace(name=name, session_id=session2)
+    langfuse.trace(name=name, session_id=session3)
+    langfuse.flush()
+
+    # Fetch traces
+    response = langfuse.fetch_sessions()
+
+    # Assert the structure of the response, cannot check for the exact number of sessions as the table is not cleared between tests
+    assert isinstance(response, FetchSessionsResponse)
+    assert hasattr(response, "data")
+    assert hasattr(response, "meta")
+    assert isinstance(response.data, list)
+    assert response.data[0].id in [session1, session2, session3]
+
+    # fetch only one, cannot check for the exact number of sessions as the table is not cleared between tests
+    response = langfuse.fetch_sessions(limit=1, page=2)
+    assert len(response.data) == 1
+    assert response.data[0].id in [session1, session2, session3]
