From 0480b2dccf87933ebe6d18264ef0804d703df33f Mon Sep 17 00:00:00 2001
From: Jamie Rasmussen <jamie.rasmussen@wandb.com>
Date: Wed, 26 Feb 2025 10:32:39 -0600
Subject: [PATCH] fix(weave): use of deprecated /calls/query endpoint

---
 weave/trace/weave_client.py                   | 25 ++++++++++---------
 .../remote_http_trace_server.py               | 11 +++++---
 2 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/weave/trace/weave_client.py b/weave/trace/weave_client.py
index e56d2ebd0f8b..b74fd1cfb63a 100644
--- a/weave/trace/weave_client.py
+++ b/weave/trace/weave_client.py
@@ -308,20 +308,21 @@ def _make_calls_iterator(
     page_size: int = DEFAULT_CALLS_PAGE_SIZE,
 ) -> CallsIter:
     def fetch_func(offset: int, limit: int) -> list[CallSchema]:
-        response = server.calls_query(
-            CallsQueryReq(
-                project_id=project_id,
-                filter=filter,
-                offset=offset,
-                limit=limit,
-                include_costs=include_costs,
-                include_feedback=include_feedback,
-                query=query,
-                sort_by=sort_by,
-                columns=columns,
+        return list(
+            server.calls_query_stream(
+                CallsQueryReq(
+                    project_id=project_id,
+                    filter=filter,
+                    offset=offset,
+                    limit=limit,
+                    include_costs=include_costs,
+                    include_feedback=include_feedback,
+                    query=query,
+                    sort_by=sort_by,
+                    columns=columns,
+                )
             )
         )
-        return response.calls
 
     # TODO: Should be Call, not WeaveObject
     def transform_func(call: CallSchema) -> WeaveObject:
diff --git a/weave/trace_server_bindings/remote_http_trace_server.py b/weave/trace_server_bindings/remote_http_trace_server.py
index ad7b86759182..9180f38a6f2c 100644
--- a/weave/trace_server_bindings/remote_http_trace_server.py
+++ b/weave/trace_server_bindings/remote_http_trace_server.py
@@ -8,6 +8,7 @@
 from pydantic import BaseModel, ValidationError
 
 from weave.trace.env import weave_trace_server_url
+from weave.trace.util import deprecated
 from weave.trace_server import requests
 from weave.trace_server import trace_server_interface as tsi
 from weave.trace_server.async_batch_processor import AsyncBatchProcessor
@@ -287,14 +288,16 @@ def call_read(self, req: Union[tsi.CallReadReq, dict[str, Any]]) -> tsi.CallRead
             "/call/read", req, tsi.CallReadReq, tsi.CallReadRes
         )
 
+    @deprecated(new_name="calls_query_stream")
     def calls_query(
         self, req: Union[tsi.CallsQueryReq, dict[str, Any]]
     ) -> tsi.CallsQueryRes:
-        return self._generic_request(
-            "/calls/query", req, tsi.CallsQueryReq, tsi.CallsQueryRes
-        )
+        # This previously called the deprecated /calls/query endpoint.
+        return tsi.CallsQueryRes(calls=list(self.calls_query_stream(req)))
 
-    def calls_query_stream(self, req: tsi.CallsQueryReq) -> Iterator[tsi.CallSchema]:
+    def calls_query_stream(
+        self, req: Union[tsi.CallsQueryReq, dict[str, Any]]
+    ) -> Iterator[tsi.CallSchema]:
         return self._generic_stream_request(
             "/calls/stream_query", req, tsi.CallsQueryReq, tsi.CallSchema
         )