Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 15 commits into from
Jun 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion weave/integrations/llamaindex/llamaindex.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
2 changes: 1 addition & 1 deletion weave/monitoring/openai/openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
5 changes: 2 additions & 3 deletions weave/tests/test_anonymous_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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(
Expand All @@ -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)
Expand Down
101 changes: 51 additions & 50 deletions weave/tests/test_client_trace.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,7 @@ def my_op(a: int) -> int:
exception=None,
output=6,
summary={},
attributes={},
)


Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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?
Expand All @@ -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(
Expand Down Expand Up @@ -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)
Expand All @@ -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(
Expand Down Expand Up @@ -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(
Expand Down
23 changes: 13 additions & 10 deletions weave/tests/test_weave_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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(
Expand All @@ -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",
Expand All @@ -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)
Expand All @@ -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

Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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"}))
Expand Down
2 changes: 1 addition & 1 deletion weave/trace/op.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down
Loading
Loading