From b2b8817ba4aea1dbd7c7f7accd81ea8dbd5595d8 Mon Sep 17 00:00:00 2001
From: Dustin Ngo <dustin@arize.com>
Date: Thu, 14 Mar 2024 16:15:10 -0400
Subject: [PATCH] Add project name to client helper functions

---
 src/phoenix/trace/dsl/helpers.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/phoenix/trace/dsl/helpers.py b/src/phoenix/trace/dsl/helpers.py
index 5e7706f090..47d277cde7 100644
--- a/src/phoenix/trace/dsl/helpers.py
+++ b/src/phoenix/trace/dsl/helpers.py
@@ -21,12 +21,13 @@
 
 
 class CanQuerySpans(Protocol):
+    # Implemented by phoenix.session.client.Client
     def query_spans(
-        self, *query: SpanQuery
+        self, *query: SpanQuery, project_name: Optional[str] = None
     ) -> Optional[Union[pd.DataFrame, List[pd.DataFrame]]]: ...
 
 
-def get_retrieved_documents(obj: CanQuerySpans) -> pd.DataFrame:
+def get_retrieved_documents(obj: CanQuerySpans, project_name: Optional[str] = None) -> pd.DataFrame:
     return cast(
         pd.DataFrame,
         obj.query_spans(
@@ -37,12 +38,13 @@ def get_retrieved_documents(obj: CanQuerySpans) -> pd.DataFrame:
                 RETRIEVAL_DOCUMENTS,
                 reference=DOCUMENT_CONTENT,
                 document_score=DOCUMENT_SCORE,
-            )
+            ),
+            project_name=project_name,
         ),
     )
 
 
-def get_qa_with_reference(obj: CanQuerySpans) -> pd.DataFrame:
+def get_qa_with_reference(obj: CanQuerySpans, project_name: Optional[str] = None) -> pd.DataFrame:
     return pd.concat(
         cast(
             List[pd.DataFrame],
@@ -55,6 +57,7 @@ def get_qa_with_reference(obj: CanQuerySpans) -> pd.DataFrame:
                     RETRIEVAL_DOCUMENTS,
                     reference=DOCUMENT_CONTENT,
                 ),
+                project_name=project_name,
             ),
         ),
         axis=1,