From d7643f9c864e62f49e904dba162ab952134c74e8 Mon Sep 17 00:00:00 2001
From: William FH <13333726+hinthornw@users.noreply.github.com>
Date: Tue, 2 Apr 2024 21:52:30 -0700
Subject: [PATCH] Add select keys (#571)

---
 js/package.json             |  4 +--
 js/src/client.ts            | 45 +++++++++++++++++++++++++++++-
 js/src/index.ts             |  2 +-
 js/src/schemas.ts           |  3 ++
 python/langsmith/client.py  | 55 +++++++++++++++++++++++++++++--------
 python/langsmith/schemas.py |  2 ++
 python/pyproject.toml       |  2 +-
 7 files changed, 97 insertions(+), 16 deletions(-)

diff --git a/js/package.json b/js/package.json
index 986084022..4d4272981 100644
--- a/js/package.json
+++ b/js/package.json
@@ -1,6 +1,6 @@
 {
   "name": "langsmith",
-  "version": "0.1.15",
+  "version": "0.1.16",
   "description": "Client library to connect to the LangSmith LLM Tracing and Evaluation Platform.",
   "packageManager": "yarn@1.22.19",
   "files": [
@@ -173,4 +173,4 @@
     },
     "./package.json": "./package.json"
   }
-}
+}
\ No newline at end of file
diff --git a/js/src/client.ts b/js/src/client.ts
index 05acaf871..749f02afa 100644
--- a/js/src/client.ts
+++ b/js/src/client.ts
@@ -67,6 +67,10 @@ interface ListRunsParams {
    * The ID of the trace to filter by.
    */
   traceId?: string;
+  /**
+   * isRoot - Whether to only include root runs.
+   *  */
+  isRoot?: boolean;
 
   /**
    * The execution order to filter by.
@@ -146,6 +150,11 @@ interface ListRunsParams {
    * conjunction with the regular `filter` parameter to let you filter runs by attributes of any run within a trace.
    */
   treeFilter?: string;
+  /**
+   * The values to include in the response.
+   *
+   */
+  select?: string[];
 }
 
 interface UploadCSVParams {
@@ -1026,7 +1035,7 @@ export class Client {
    * @param traceId - The ID of the trace to filter by.
    * @param referenceExampleId - The ID of the reference example to filter by.
    * @param startTime - The start time to filter by.
-   * @param executionOrder - The execution order to filter by.
+   * @param isRoot - Indicates whether to only return root runs.
    * @param runType - The run type to filter by.
    * @param error - Indicates whether to filter by error runs.
    * @param id - The ID of the run to filter by.
@@ -1108,6 +1117,7 @@ export class Client {
       referenceExampleId,
       startTime,
       executionOrder,
+      isRoot,
       runType,
       error,
       id,
@@ -1116,6 +1126,7 @@ export class Client {
       traceFilter,
       treeFilter,
       limit,
+      select,
     } = props;
     let projectIds: string[] = [];
     if (projectId) {
@@ -1132,6 +1143,36 @@ export class Client {
       );
       projectIds.push(...projectIds_);
     }
+    const default_select = [
+      "app_path",
+      "child_run_ids",
+      "completion_cost",
+      "completion_tokens",
+      "dotted_order",
+      "end_time",
+      "error",
+      "events",
+      "extra",
+      "feedback_stats",
+      "first_token_time",
+      "id",
+      "inputs",
+      "name",
+      "outputs",
+      "parent_run_id",
+      "parent_run_ids",
+      "prompt_cost",
+      "prompt_tokens",
+      "reference_example_id",
+      "run_type",
+      "session_id",
+      "start_time",
+      "status",
+      "tags",
+      "total_cost",
+      "total_tokens",
+      "trace_id",
+    ];
     const body = {
       session: projectIds.length ? projectIds : null,
       run_type: runType,
@@ -1147,6 +1188,8 @@ export class Client {
       id,
       limit,
       trace: traceId,
+      select: select ? select : default_select,
+      is_root: isRoot,
     };
 
     for await (const runs of this._getCursorPaginatedList<Run>(
diff --git a/js/src/index.ts b/js/src/index.ts
index 3454f1b8d..4a61115ee 100644
--- a/js/src/index.ts
+++ b/js/src/index.ts
@@ -11,4 +11,4 @@ export type {
 export { RunTree, type RunTreeConfig } from "./run_trees.js";
 
 // Update using yarn bump-version
-export const __version__ = "0.1.15";
+export const __version__ = "0.1.16";
diff --git a/js/src/schemas.ts b/js/src/schemas.ts
index cf207c674..e6a43dbf1 100644
--- a/js/src/schemas.ts
+++ b/js/src/schemas.ts
@@ -168,6 +168,9 @@ export interface Run extends BaseRun {
 
   /** IDs of parent runs, if multiple exist. */
   parent_run_ids?: string[];
+
+  /** Whether the run is included in a dataset. */
+  in_dataset?: boolean;
 }
 
 export interface RunCreate extends BaseRun {
diff --git a/python/langsmith/client.py b/python/langsmith/client.py
index 8716561d3..5caccc355 100644
--- a/python/langsmith/client.py
+++ b/python/langsmith/client.py
@@ -1440,11 +1440,12 @@ def list_runs(
         filter: Optional[str] = None,
         trace_filter: Optional[str] = None,
         tree_filter: Optional[str] = None,
-        execution_order: Optional[int] = None,
+        is_root: Optional[bool] = None,
         parent_run_id: Optional[ID_TYPE] = None,
         start_time: Optional[datetime.datetime] = None,
         error: Optional[bool] = None,
-        run_ids: Optional[List[ID_TYPE]] = None,
+        run_ids: Optional[Sequence[ID_TYPE]] = None,
+        select: Optional[Sequence[str]] = None,
         limit: Optional[int] = None,
         **kwargs: Any,
     ) -> Iterator[ls_schemas.Run]:
@@ -1475,10 +1476,8 @@ def list_runs(
             sibling and child runs. This is meant to be used in conjunction with
             the regular `filter` parameter to let you filter runs by attributes
             of any run within a trace.
-        execution_order : int or None, default=None
-            The execution order to filter by. Execution order is the position
-            of the run in the full trace's execution sequence.
-                All root run traces have execution_order 1.
+        is_root : bool or None, default=None
+            Whether to filter by root runs.
         parent_run_id : UUID or None, default=None
             The ID of the parent run to filter by.
         start_time : datetime or None, default=None
@@ -1511,12 +1510,15 @@ def list_runs(
                 run_type="llm",
             )
 
-            # List traces in a project
-            root_runs = client.list_runs(project_name="<your_project>", execution_order=1)
+            # List root traces in a project
+            root_runs = client.list_runs(project_name="<your_project>", is_root=1)
 
             # List runs without errors
             correct_runs = client.list_runs(project_name="<your_project>", error=False)
 
+            # List runs and only return their inputs/outputs (to speed up the query)
+            input_output_runs = client.list_runs(project_name="<your_project>", select=["inputs", "outputs"])
+
             # List runs by run ID
             run_ids = [
                 "a36092d2-4ad5-4fb4-9c0d-0dba9a2ed836",
@@ -1561,7 +1563,37 @@ def list_runs(
             project_ids.extend(
                 [self.read_project(project_name=name).id for name in project_name]
             )
-
+        default_select = [
+            "app_path",
+            "child_run_ids",
+            "completion_cost",
+            "completion_tokens",
+            "dotted_order",
+            "end_time",
+            "error",
+            "events",
+            "extra",
+            "feedback_stats",
+            "first_token_time",
+            "id",
+            "inputs",
+            "name",
+            "outputs",
+            "parent_run_id",
+            "parent_run_ids",
+            "prompt_cost",
+            "prompt_tokens",
+            "reference_example_id",
+            "run_type",
+            "session_id",
+            "start_time",
+            "status",
+            "tags",
+            "total_cost",
+            "total_tokens",
+            "trace_id",
+        ]
+        select = select or default_select
         body_query: Dict[str, Any] = {
             "session": project_ids if project_ids else None,
             "run_type": run_type,
@@ -1572,12 +1604,13 @@ def list_runs(
             "filter": filter,
             "trace_filter": trace_filter,
             "tree_filter": tree_filter,
-            "execution_order": execution_order,
+            "is_root": is_root,
             "parent_run": parent_run_id,
             "start_time": start_time.isoformat() if start_time else None,
             "error": error,
             "id": run_ids,
             "trace": trace_id,
+            "select": select,
             **kwargs,
         }
         body_query = {k: v for k, v in body_query.items() if v is not None}
@@ -2051,7 +2084,7 @@ def get_test_results(
         import pandas as pd  # type: ignore
 
         runs = self.list_runs(
-            project_id=project_id, project_name=project_name, execution_order=1
+            project_id=project_id, project_name=project_name, is_root=True
         )
         results = []
         example_ids = []
diff --git a/python/langsmith/schemas.py b/python/langsmith/schemas.py
index 50947e6d4..5823d278c 100644
--- a/python/langsmith/schemas.py
+++ b/python/langsmith/schemas.py
@@ -304,6 +304,8 @@ class Run(RunBase):
         - 20230914T223155647Z1b64098b-4ab7-43f6-afee-992304f198d8.20230914T223155649Z809ed3a2-0172-4f4d-8a02-a64e9b7a0f8a
         - 20230915T223155647Z1b64098b-4ab7-43f6-afee-992304f198d8.20230914T223155650Zc8d9f4c5-6c5a-4b2d-9b1c-3d9d7a7c5c7c
     """  # noqa: E501
+    in_dataset: Optional[bool] = None
+    """Whether this run is in a dataset."""
     _host_url: Optional[str] = PrivateAttr(default=None)
 
     def __init__(self, _host_url: Optional[str] = None, **kwargs: Any) -> None:
diff --git a/python/pyproject.toml b/python/pyproject.toml
index 17504e4c4..ceaaea5a1 100644
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langsmith"
-version = "0.1.38"
+version = "0.1.39"
 description = "Client library to connect to the LangSmith LLM Tracing and Evaluation Platform."
 authors = ["LangChain <support@langchain.dev>"]
 license = "MIT"