chore(weave): Add attributes to Call object and allow explicit root (#…

…1739)

* init

* revised implementation

* fixed

* refactored according to advice

* refactored according to advice

* final refactor

* fixed comment strings

* removed underscore

* moved parent to 3rd param

weave/integrations/llamaindex/llamaindex.py

@@ -95,8 +95,8 @@ def on_event_start(
                     # Create a call object.
                     call = gc.create_call(
                         "llama_index." + event_type.name.lower(),
-                        self._call_map.get(parent_id),
                         process_payload(payload),
+                        self._call_map.get(parent_id),
                     )
 
                     # Add the call to the call map.

weave/monitoring/openai/openai.py

@@ -259,7 +259,7 @@ def log_run(
     client = graph_client_context.require_graph_client()
     parent_run = run_context.get_current_run()
     # TODO: client should not need refs passed in.
-    run = client.create_call(call_name, parent_run, inputs)
+    run = client.create_call(call_name, inputs, parent_run)
 
     def finish_run(output: Any) -> None:
         # TODO: client should not need refs passed in.

weave/tests/test_anonymous_ops.py

@@ -3,7 +3,7 @@
 
 
 def test_named_op(client: weave_client.WeaveClient) -> str:
-    call = client.create_call("anonymous_op", None, {"a": 1})
+    call = client.create_call("anonymous_op", {"a": 1})
     client.finish_call(call, {"c": 3}, None)
 
     call_res = client.server.calls_query(
@@ -25,7 +25,7 @@ def test_named_op(client: weave_client.WeaveClient) -> str:
 
 
 def test_anonymous_op(client: weave_client.WeaveClient) -> str:
-    call = client.create_call("anonymous_op", None, {"a": 1})
+    call = client.create_call("anonymous_op", {"a": 1})
     client.finish_call(call, {"c": 3}, None)
 
     call_res = client.server.calls_query(
@@ -49,7 +49,6 @@ def test_anonymous_op(client: weave_client.WeaveClient) -> str:
 def test_anonymous_op_with_config(client: weave_client.WeaveClient) -> str:
     call = client.create_call(
         weave_client._build_anonymous_op("anonymous_op", {"library_version": "0.42.0"}),
-        None,
         {"a": 1},
     )
     client.finish_call(call, {"c": 3}, None)

weave/tests/test_client_trace.py

@@ -74,6 +74,7 @@ def my_op(a: int) -> int:
         exception=None,
         output=6,
         summary={},
+        attributes={},
     )
 
 
@@ -1721,22 +1722,22 @@ def test_call_stack_order_implicit_depth_first(client):
 
     # This version of the call sequence matches the happy path
     # without any out-of-order calls
-    call_1 = client.create_call("op", None, {})
-    call_2 = client.create_call("op", None, {})
-    call_3 = client.create_call("op", None, {})
+    call_1 = client.create_call("op", {})
+    call_2 = client.create_call("op", {})
+    call_3 = client.create_call("op", {})
     client.finish_call(call_3)
-    call_4 = client.create_call("op", None, {})
+    call_4 = client.create_call("op", {})
     client.finish_call(call_4)
     client.finish_call(call_2)
-    call_5 = client.create_call("op", None, {})
-    call_6 = client.create_call("op", None, {})
+    call_5 = client.create_call("op", {})
+    call_6 = client.create_call("op", {})
     client.finish_call(call_6)
-    call_7 = client.create_call("op", None, {})
+    call_7 = client.create_call("op", {})
     client.finish_call(call_7)
     client.finish_call(call_5)
     client.finish_call(call_1)
 
-    terminal_root_call = client.create_call("op", None, {})
+    terminal_root_call = client.create_call("op", {})
     client.finish_call(terminal_root_call)
 
     inner_res = client.server.calls_query(
@@ -1767,22 +1768,22 @@ def test_call_stack_order_explicit_depth_first(client):
     # This version of the call sequence matches the happy path
     # without any out-of-order calls, but with explicit parentage
     # specified.
-    call_1 = client.create_call("op", None, {})
-    call_2 = client.create_call("op", call_1, {})
-    call_3 = client.create_call("op", call_2, {})
+    call_1 = client.create_call("op", {})
+    call_2 = client.create_call("op", {}, call_1)
+    call_3 = client.create_call("op", {}, call_2)
     client.finish_call(call_3)
-    call_4 = client.create_call("op", call_2, {})
+    call_4 = client.create_call("op", {}, call_2)
     client.finish_call(call_4)
     client.finish_call(call_2)
-    call_5 = client.create_call("op", call_1, {})
-    call_6 = client.create_call("op", call_5, {})
+    call_5 = client.create_call("op", {}, call_1)
+    call_6 = client.create_call("op", {}, call_5)
     client.finish_call(call_6)
-    call_7 = client.create_call("op", call_5, {})
+    call_7 = client.create_call("op", {}, call_5)
     client.finish_call(call_7)
     client.finish_call(call_5)
     client.finish_call(call_1)
 
-    terminal_root_call = client.create_call("op", None, {})
+    terminal_root_call = client.create_call("op", {})
     client.finish_call(terminal_root_call)
 
     inner_res = client.server.calls_query(
@@ -1812,26 +1813,26 @@ def test_call_stack_order_langchain_batch(client):
     #
     # This sequence is pretty much exactly what langchain does when handling
     # a batch of calls. Specifically (prompt | llm).batch([1,2])
-    call_1 = client.create_call("op", None, {})  # <- Implicit Parent, no stack = root
-    call_2 = client.create_call("op", call_1, {})  # <- RunnableSequence1
-    call_5 = client.create_call("op", call_1, {})  # <- RunnableSequence2
-    call_3 = client.create_call("op", call_2, {})  # <- Prompt1
+    call_1 = client.create_call("op", {})  # <- Implicit Parent, no stack = root
+    call_2 = client.create_call("op", {}, call_1)  # <- RunnableSequence1
+    call_5 = client.create_call("op", {}, call_1)  # <- RunnableSequence2
+    call_3 = client.create_call("op", {}, call_2)  # <- Prompt1
     client.finish_call(call_3)
-    call_4 = client.create_call("op", call_2, {})  # <- LLM1
-    call_4gpt = client.create_call("op", None, {})  # <- Openai
+    call_4 = client.create_call("op", {}, call_2)  # <- LLM1
+    call_4gpt = client.create_call("op", {})  # <- Openai
     client.finish_call(call_4gpt)
     client.finish_call(call_4)
-    call_6 = client.create_call("op", call_5, {})  # <- Prompt2
+    call_6 = client.create_call("op", {}, call_5)  # <- Prompt2
     client.finish_call(call_6)
-    call_7 = client.create_call("op", call_5, {})  # <- LLM2
-    call_7gpt = client.create_call("op", None, {})  # <- Openai
+    call_7 = client.create_call("op", {}, call_5)  # <- LLM2
+    call_7gpt = client.create_call("op", {})  # <- Openai
     client.finish_call(call_7gpt)
     client.finish_call(call_7)
     client.finish_call(call_2)
     client.finish_call(call_5)
     client.finish_call(call_1)
 
-    terminal_root_call = client.create_call("op", None, {})
+    terminal_root_call = client.create_call("op", {})
     client.finish_call(terminal_root_call)
 
     inner_res = client.server.calls_query(
@@ -1864,16 +1865,16 @@ def test_call_stack_order_out_of_order_pop(client):
     #
     # This ordering is a specifically challenging case where we return to
     # a parent that that was not the top of stack
-    call_1 = client.create_call("op", None, {})
-    call_2 = client.create_call("op", None, {})
-    call_3 = client.create_call("op", None, {})
+    call_1 = client.create_call("op", {})
+    call_2 = client.create_call("op", {})
+    call_3 = client.create_call("op", {})
     # Purposely swap 4 & 5
-    call_5 = client.create_call("op", call_1, {})  # <- Explicit Parent (call_1)
-    call_4 = client.create_call("op", call_2, {})  # <- Explicit Parent (call_2)
-    call_6 = client.create_call("op", call_5, {})  # <- Explicit Parent (call_5)
+    call_5 = client.create_call("op", {}, call_1)  # <- Explicit Parent (call_1)
+    call_4 = client.create_call("op", {}, call_2)  # <- Explicit Parent (call_2)
+    call_6 = client.create_call("op", {}, call_5)  # <- Explicit Parent (call_5)
     client.finish_call(call_6)  # <- Finish call_6
     # (should change stack to call_6.parent which is call_5)
-    call_7 = client.create_call("op", None, {})  # <- Implicit Parent (call_5)
+    call_7 = client.create_call("op", {})  # <- Implicit Parent (call_5)
     # (current stack implementation will think this is 4)
 
     # Finish them in completely reverse order, because why not?
@@ -1884,7 +1885,7 @@ def test_call_stack_order_out_of_order_pop(client):
     client.finish_call(call_5)
     client.finish_call(call_7)
 
-    terminal_root_call = client.create_call("op", None, {})
+    terminal_root_call = client.create_call("op", {})
     client.finish_call(terminal_root_call)
 
     inner_res = client.server.calls_query(
@@ -1935,13 +1936,13 @@ def test_call_stack_order_height_ordering(client):
     #
     #
     # This ordering calls ops in the order of their height in the tree
-    call_1 = client.create_call("op", None, {})
-    call_2 = client.create_call("op", call_1, {})
-    call_5 = client.create_call("op", call_1, {})
-    call_3 = client.create_call("op", call_2, {})
-    call_6 = client.create_call("op", call_5, {})
-    call_4 = client.create_call("op", call_2, {})
-    call_7 = client.create_call("op", call_5, {})
+    call_1 = client.create_call("op", {})
+    call_2 = client.create_call("op", {}, call_1)
+    call_5 = client.create_call("op", {}, call_1)
+    call_3 = client.create_call("op", {}, call_2)
+    call_6 = client.create_call("op", {}, call_5)
+    call_4 = client.create_call("op", {}, call_2)
+    call_7 = client.create_call("op", {}, call_5)
 
     # Finish them in completely reverse order
     client.finish_call(call_1)
@@ -1951,7 +1952,7 @@ def test_call_stack_order_height_ordering(client):
     client.finish_call(call_5)
     client.finish_call(call_7)
 
-    terminal_root_call = client.create_call("op", None, {})
+    terminal_root_call = client.create_call("op", {})
     client.finish_call(terminal_root_call)
 
     inner_res = client.server.calls_query(
@@ -1980,21 +1981,21 @@ def test_call_stack_order_mixed(client):
     #
     #
     # This ordering is as mixed up as I could make it
-    call_1 = client.create_call("op", None, {})
-    call_5 = client.create_call("op", call_1, {})
-    call_7 = client.create_call("op", call_5, {})
+    call_1 = client.create_call("op", {})
+    call_5 = client.create_call("op", {}, call_1)
+    call_7 = client.create_call("op", {}, call_5)
     client.finish_call(call_7)
-    call_6 = client.create_call("op", call_5, {})
+    call_6 = client.create_call("op", {}, call_5)
     client.finish_call(call_5)
-    call_2 = client.create_call("op", call_1, {})
+    call_2 = client.create_call("op", {}, call_1)
     client.finish_call(call_1)
-    call_4 = client.create_call("op", call_2, {})
-    call_3 = client.create_call("op", call_2, {})
+    call_4 = client.create_call("op", {}, call_2)
+    call_3 = client.create_call("op", {}, call_2)
     client.finish_call(call_2)
     client.finish_call(call_3)
     client.finish_call(call_4)
 
-    terminal_root_call = client.create_call("op", None, {})
+    terminal_root_call = client.create_call("op", {})
     client.finish_call(terminal_root_call)
 
     inner_res = client.server.calls_query(

weave/tests/test_weave_client.py

@@ -197,7 +197,7 @@ class B(A):
 
 
 def test_call_create(client):
-    call = client.create_call("x", None, {"a": 5, "b": 10})
+    call = client.create_call("x", {"a": 5, "b": 10})
     client.finish_call(call, "hello")
     result = client.call(call.id)
     expected = weave_client.Call(
@@ -211,14 +211,15 @@ def test_call_create(client):
         exception=None,
         summary={},
         _children=[],
+        attributes={},
     )
     assert dataclasses.asdict(result._val) == dataclasses.asdict(expected)
 
 
 def test_calls_query(client):
-    call0 = client.create_call("x", None, {"a": 5, "b": 10})
-    call1 = client.create_call("x", None, {"a": 6, "b": 11})
-    call2 = client.create_call("y", None, {"a": 5, "b": 10})
+    call0 = client.create_call("x", {"a": 5, "b": 10})
+    call1 = client.create_call("x", {"a": 6, "b": 11})
+    call2 = client.create_call("y", {"a": 5, "b": 10})
     result = list(client.calls(weave_client._CallsFilter(op_names=[call1.op_name])))
     assert len(result) == 2
     assert result[0] == weave_client.Call(
@@ -228,6 +229,7 @@ def test_calls_query(client):
         parent_id=None,
         inputs={"a": 5, "b": 10},
         id=call0.id,
+        attributes={},
     )
     assert result[1] == weave_client.Call(
         op_name="weave:///shawn/test-project/op/x:tzUhDyzVm5bqQsuqh5RT4axEXSosyLIYZn9zbRyenaw",
@@ -236,6 +238,7 @@ def test_calls_query(client):
         parent_id=call0.id,
         inputs={"a": 6, "b": 11},
         id=call1.id,
+        attributes={},
     )
     client.finish_call(call2, None)
     client.finish_call(call1, None)
@@ -257,10 +260,10 @@ def patched_delete(req: tsi.CallsDeleteReq) -> tsi.CallsDeleteRes:
     # Patch calls_delete with our fake middlewear
     client.server.calls_delete = patched_delete
 
-    call0 = client.create_call("x", None, {"a": 5, "b": 10})
-    call0_child1 = client.create_call("x", call0, {"a": 5, "b": 11})
-    _call0_child2 = client.create_call("x", call0_child1, {"a": 5, "b": 12})
-    call1 = client.create_call("y", None, {"a": 6, "b": 11})
+    call0 = client.create_call("x", {"a": 5, "b": 10})
+    call0_child1 = client.create_call("x", {"a": 5, "b": 11}, call0)
+    _call0_child2 = client.create_call("x", {"a": 5, "b": 12}, call0_child1)
+    call1 = client.create_call("y", {"a": 6, "b": 11})
 
     assert len(list(client.calls())) == 4
 
@@ -349,7 +352,7 @@ def patched_update(req: tsi.CallUpdateReq) -> tsi.CallUpdateRes:
     # Patch call_update with our fake middlewear
     client.server.call_update = patched_update
 
-    call0 = client.create_call("x", None, {"a": 5, "b": 10})
+    call0 = client.create_call("x", {"a": 5, "b": 10})
 
     # Rename using the client method
     client.set_call_display_name(call0, "updated_name")
@@ -427,7 +430,7 @@ def test_dataset_calls(client):
         "my-dataset",
     )
     for row in ref.rows:
-        call = client.create_call("x", None, {"a": row["doc"]})
+        call = client.create_call("x", {"a": row["doc"]})
         client.finish_call(call, None)
 
     calls = list(client.calls({"op_name": "x"}))

weave/trace/op.py

@@ -82,8 +82,8 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
         attributes = call_attributes.get()
         run = client.create_call(
             self,
-            parent_run,
             inputs_with_defaults,
+            parent_run,
             attributes=attributes,
             display_name=self.display_name,
         )