diff --git a/tests/contrib/langgraph/test_langgraph.py b/tests/contrib/langgraph/test_langgraph.py index 1894b8b9420..1ea7f8714ab 100644 --- a/tests/contrib/langgraph/test_langgraph.py +++ b/tests/contrib/langgraph/test_langgraph.py @@ -1,134 +1,159 @@ +from collections import Counter + + +def assert_has_spans(spans, expected): + resources = [span.resource for span in spans] + assert len(resources) == len(expected) + assert Counter(resources) == Counter(expected) + + def assert_simple_graph_spans(spans): - assert len(spans) == 3 - assert spans[0].resource == "langgraph.graph.state.CompiledStateGraph.LangGraph" - assert spans[1].resource == "langgraph.utils.runnable.RunnableSeq.a" - assert spans[2].resource == "langgraph.utils.runnable.RunnableSeq.b" + assert_has_spans( + spans, + expected=[ + "langgraph.graph.state.CompiledStateGraph.LangGraph", + "langgraph.utils.runnable.RunnableSeq.a", + "langgraph.utils.runnable.RunnableSeq.b", + ], + ) def assert_conditional_graph_spans(spans, which): - assert len(spans) == 3 - assert spans[0].resource == "langgraph.graph.state.CompiledStateGraph.LangGraph" - assert spans[1].resource == "langgraph.utils.runnable.RunnableSeq.a" - assert spans[2].resource == f"langgraph.utils.runnable.RunnableSeq.{which}" + assert_has_spans( + spans, + expected=[ + "langgraph.graph.state.CompiledStateGraph.LangGraph", + "langgraph.utils.runnable.RunnableSeq.a", + f"langgraph.utils.runnable.RunnableSeq.{which}", + ], + ) def assert_subgraph_spans(spans): - assert len(spans) == 6 - assert spans[0].resource == "langgraph.graph.state.CompiledStateGraph.LangGraph" - assert spans[1].resource == "langgraph.utils.runnable.RunnableSeq.a" - assert spans[2].resource == "langgraph.graph.state.CompiledStateGraph.LangGraph" - assert spans[3].resource == "langgraph.utils.runnable.RunnableSeq.b1" - assert spans[4].resource == "langgraph.utils.runnable.RunnableSeq.b2" - assert spans[5].resource == "langgraph.utils.runnable.RunnableSeq.b3" + assert_has_spans( + spans, + expected=[ + "langgraph.graph.state.CompiledStateGraph.LangGraph", + "langgraph.utils.runnable.RunnableSeq.a", + "langgraph.graph.state.CompiledStateGraph.LangGraph", + "langgraph.utils.runnable.RunnableSeq.b1", + "langgraph.utils.runnable.RunnableSeq.b2", + "langgraph.utils.runnable.RunnableSeq.b3", + ], + ) def assert_fanning_graph_spans(spans): - assert len(spans) == 5 - assert spans[0].resource == "langgraph.graph.state.CompiledStateGraph.LangGraph" - assert spans[1].resource == "langgraph.utils.runnable.RunnableSeq.a" - assert spans[2].resource == "langgraph.utils.runnable.RunnableSeq.b" - assert spans[3].resource == "langgraph.utils.runnable.RunnableSeq.c" - assert spans[4].resource == "langgraph.utils.runnable.RunnableSeq.d" - - -def test_simple_graph(langgraph, simple_graph, mock_tracer): + assert_has_spans( + spans, + expected=[ + "langgraph.graph.state.CompiledStateGraph.LangGraph", + "langgraph.utils.runnable.RunnableSeq.a", + "langgraph.utils.runnable.RunnableSeq.b", + "langgraph.utils.runnable.RunnableSeq.c", + "langgraph.utils.runnable.RunnableSeq.d", + ], + ) + + +def test_simple_graph(simple_graph, mock_tracer): simple_graph.invoke({"a_list": [], "which": "a"}) spans = mock_tracer.pop_traces()[0] assert_simple_graph_spans(spans) -async def test_simple_graph_async(langgraph, simple_graph, mock_tracer): +async def test_simple_graph_async(simple_graph, mock_tracer): await simple_graph.ainvoke({"a_list": [], "which": "a"}) spans = mock_tracer.pop_traces()[0] assert_simple_graph_spans(spans) -def test_simple_graph_stream(langgraph, simple_graph, mock_tracer): +def test_simple_graph_stream(simple_graph, mock_tracer): for _ in simple_graph.stream({"a_list": [], "which": "a"}): pass spans = mock_tracer.pop_traces()[0] assert_simple_graph_spans(spans) -async def test_simple_graph_stream_async(langgraph, simple_graph, mock_tracer): +async def test_simple_graph_stream_async(simple_graph, mock_tracer): async for _ in simple_graph.astream({"a_list": [], "which": "a"}): pass spans = mock_tracer.pop_traces()[0] assert_simple_graph_spans(spans) -def test_conditional_graph(langgraph, conditional_graph, mock_tracer): +def test_conditional_graph(conditional_graph, mock_tracer): conditional_graph.invoke({"a_list": [], "which": "c"}) spans = mock_tracer.pop_traces()[0] assert_conditional_graph_spans(spans, which="c") -async def test_conditional_graph_async(langgraph, conditional_graph, mock_tracer): +async def test_conditional_graph_async(conditional_graph, mock_tracer): await conditional_graph.ainvoke({"a_list": [], "which": "b"}) spans = mock_tracer.pop_traces()[0] assert_conditional_graph_spans(spans, which="b") -def test_conditional_graph_stream(langgraph, conditional_graph, mock_tracer): +def test_conditional_graph_stream(conditional_graph, mock_tracer): for _ in conditional_graph.stream({"a_list": [], "which": "c"}): pass spans = mock_tracer.pop_traces()[0] assert_conditional_graph_spans(spans, which="c") -async def test_conditional_graph_stream_async(langgraph, conditional_graph, mock_tracer): +async def test_conditional_graph_stream_async(conditional_graph, mock_tracer): async for _ in conditional_graph.astream({"a_list": [], "which": "b"}): pass spans = mock_tracer.pop_traces()[0] assert_conditional_graph_spans(spans, which="b") -def test_subgraph(langgraph, complex_graph, mock_tracer): +def test_subgraph(complex_graph, mock_tracer): complex_graph.invoke({"a_list": [], "which": "b"}) spans = mock_tracer.pop_traces()[0] assert_subgraph_spans(spans) -async def test_subgraph_async(langgraph, complex_graph, mock_tracer): +async def test_subgraph_async(complex_graph, mock_tracer): await complex_graph.ainvoke({"a_list": [], "which": "b"}) spans = mock_tracer.pop_traces()[0] assert_subgraph_spans(spans) -def test_subgraph_stream(langgraph, complex_graph, mock_tracer): +def test_subgraph_stream(complex_graph, mock_tracer): for _ in complex_graph.stream({"a_list": [], "which": "b"}): pass spans = mock_tracer.pop_traces()[0] assert_subgraph_spans(spans) -async def test_subgraph_stream_async(langgraph, complex_graph, mock_tracer): +async def test_subgraph_stream_async(complex_graph, mock_tracer): async for _ in complex_graph.astream({"a_list": [], "which": "b"}): pass spans = mock_tracer.pop_traces()[0] assert_subgraph_spans(spans) -def test_fanning_graph(langgraph, fanning_graph, mock_tracer): +def test_fanning_graph(fanning_graph, mock_tracer): fanning_graph.invoke({"a_list": [], "which": "b"}) spans = mock_tracer.pop_traces()[0] assert_fanning_graph_spans(spans) -async def test_fanning_graph_async(langgraph, fanning_graph, mock_tracer): +async def test_fanning_graph_async(fanning_graph, mock_tracer): await fanning_graph.ainvoke({"a_list": [], "which": "b"}) spans = mock_tracer.pop_traces()[0] assert_fanning_graph_spans(spans) -def test_fanning_graph_stream(langgraph, fanning_graph, mock_tracer): +def test_fanning_graph_stream(fanning_graph, mock_tracer): for _ in fanning_graph.stream({"a_list": [], "which": "b"}): pass spans = mock_tracer.pop_traces()[0] assert_fanning_graph_spans(spans) -async def test_fanning_graph_stream_async(langgraph, fanning_graph, mock_tracer): +async def test_fanning_graph_stream_async(fanning_graph, mock_tracer): async for _ in fanning_graph.astream({"a_list": [], "which": "b"}): pass spans = mock_tracer.pop_traces()[0] diff --git a/tests/contrib/langgraph/test_langgraph_llmobs.py b/tests/contrib/langgraph/test_langgraph_llmobs.py index 968a98bdb26..46728dcdab0 100644 --- a/tests/contrib/langgraph/test_langgraph_llmobs.py +++ b/tests/contrib/langgraph/test_langgraph_llmobs.py @@ -6,8 +6,8 @@ def _assert_span_link(from_span_event, to_span_event, from_io, to_io): Assert that a span link exists between two span events, specifically the correct span ID and from/to specification. """ found = False - expected_to_span_id = "undefined" if not to_span_event else to_span_event["span_id"] - for span_link in from_span_event["span_links"]: + expected_to_span_id = "undefined" if not from_span_event else from_span_event["span_id"] + for span_link in to_span_event["span_links"]: if span_link["span_id"] == expected_to_span_id: assert span_link["attributes"] == {"from": from_io, "to": to_io} found = True @@ -15,20 +15,24 @@ def _assert_span_link(from_span_event, to_span_event, from_io, to_io): assert found +def _find_span_by_name(llmobs_events, name): + for event in llmobs_events: + if event["name"] == name: + return event + assert False, f"Span '{name}' not found in llmobs span events" + + class TestLangGraphLLMObs: def test_simple_graph(self, llmobs_events, simple_graph): simple_graph.invoke({"a_list": [], "which": "a"}, stream_mode=["values"]) - graph_span = llmobs_events[2] - a_span = llmobs_events[0] - b_span = llmobs_events[1] + graph_span = _find_span_by_name(llmobs_events, "LangGraph") + a_span = _find_span_by_name(llmobs_events, "a") + b_span = _find_span_by_name(llmobs_events, "b") - assert graph_span["name"] == "LangGraph" - _assert_span_link(graph_span, None, "input", "input") - _assert_span_link(graph_span, b_span, "output", "output") - assert a_span["name"] == "a" - _assert_span_link(a_span, graph_span, "input", "input") - assert b_span["name"] == "b" - _assert_span_link(b_span, a_span, "output", "input") + _assert_span_link(None, graph_span, "input", "input") + _assert_span_link(b_span, graph_span, "output", "output") + _assert_span_link(graph_span, a_span, "input", "input") + _assert_span_link(a_span, b_span, "output", "input") # Check that the graph span has the appropriate output # stream_mode=["values"] should result in the last yield being a tuple @@ -36,17 +40,14 @@ def test_simple_graph(self, llmobs_events, simple_graph): async def test_simple_graph_async(self, llmobs_events, simple_graph): await simple_graph.ainvoke({"a_list": [], "which": "a"}) - graph_span = llmobs_events[2] - a_span = llmobs_events[0] - b_span = llmobs_events[1] + graph_span = _find_span_by_name(llmobs_events, "LangGraph") + a_span = _find_span_by_name(llmobs_events, "a") + b_span = _find_span_by_name(llmobs_events, "b") - assert graph_span["name"] == "LangGraph" - _assert_span_link(graph_span, None, "input", "input") - _assert_span_link(graph_span, b_span, "output", "output") - assert a_span["name"] == "a" - _assert_span_link(a_span, graph_span, "input", "input") - assert b_span["name"] == "b" - _assert_span_link(b_span, a_span, "output", "input") + _assert_span_link(None, graph_span, "input", "input") + _assert_span_link(b_span, graph_span, "output", "output") + _assert_span_link(graph_span, a_span, "input", "input") + _assert_span_link(a_span, b_span, "output", "input") # Check that the graph span has the appropriate output # default stream_mode of "values" should result in the last yield being an object @@ -54,118 +55,91 @@ async def test_simple_graph_async(self, llmobs_events, simple_graph): def test_conditional_graph(self, llmobs_events, conditional_graph): conditional_graph.invoke({"a_list": [], "which": "c"}) - graph_span = llmobs_events[2] - a_span = llmobs_events[0] - c_span = llmobs_events[1] + graph_span = _find_span_by_name(llmobs_events, "LangGraph") + a_span = _find_span_by_name(llmobs_events, "a") + c_span = _find_span_by_name(llmobs_events, "c") - assert graph_span["name"] == "LangGraph" - _assert_span_link(graph_span, None, "input", "input") - _assert_span_link(graph_span, c_span, "output", "output") - assert a_span["name"] == "a" - _assert_span_link(a_span, graph_span, "input", "input") - assert c_span["name"] == "c" - _assert_span_link(c_span, a_span, "output", "input") + _assert_span_link(None, graph_span, "input", "input") + _assert_span_link(c_span, graph_span, "output", "output") + _assert_span_link(graph_span, a_span, "input", "input") + _assert_span_link(a_span, c_span, "output", "input") async def test_conditional_graph_async(self, llmobs_events, conditional_graph): await conditional_graph.ainvoke({"a_list": [], "which": "b"}) - graph_span = llmobs_events[2] - a_span = llmobs_events[0] - c_span = llmobs_events[1] + graph_span = _find_span_by_name(llmobs_events, "LangGraph") + a_span = _find_span_by_name(llmobs_events, "a") + b_span = _find_span_by_name(llmobs_events, "b") - assert graph_span["name"] == "LangGraph" - _assert_span_link(graph_span, None, "input", "input") - _assert_span_link(graph_span, c_span, "output", "output") - assert a_span["name"] == "a" - _assert_span_link(a_span, graph_span, "input", "input") - assert c_span["name"] == "b" - _assert_span_link(c_span, a_span, "output", "input") + _assert_span_link(None, graph_span, "input", "input") + _assert_span_link(b_span, graph_span, "output", "output") + _assert_span_link(graph_span, a_span, "input", "input") + _assert_span_link(a_span, b_span, "output", "input") def test_subgraph(self, llmobs_events, complex_graph): complex_graph.invoke({"a_list": [], "which": "b"}) - graph_span = llmobs_events[5] - a_span = llmobs_events[0] - b_span = llmobs_events[4] - b1_span = llmobs_events[1] - b2_span = llmobs_events[2] - b3_span = llmobs_events[3] - - assert graph_span["name"] == "LangGraph" - _assert_span_link(graph_span, None, "input", "input") - _assert_span_link(graph_span, b_span, "output", "output") - assert a_span["name"] == "a" - _assert_span_link(a_span, graph_span, "input", "input") - assert b_span["name"] == "b" - _assert_span_link(b_span, a_span, "output", "input") - _assert_span_link(b_span, b3_span, "output", "output") - assert b1_span["name"] == "b1" - _assert_span_link(b1_span, b_span, "input", "input") - assert b2_span["name"] == "b2" - _assert_span_link(b2_span, b1_span, "output", "input") - assert b3_span["name"] == "b3" - _assert_span_link(b3_span, b2_span, "output", "input") + graph_span = _find_span_by_name(llmobs_events, "LangGraph") + a_span = _find_span_by_name(llmobs_events, "a") + b_span = _find_span_by_name(llmobs_events, "b") + b1_span = _find_span_by_name(llmobs_events, "b1") + b2_span = _find_span_by_name(llmobs_events, "b2") + b3_span = _find_span_by_name(llmobs_events, "b3") + + _assert_span_link(None, graph_span, "input", "input") + _assert_span_link(b3_span, b_span, "output", "output") + _assert_span_link(graph_span, a_span, "input", "input") + _assert_span_link(a_span, b_span, "output", "input") + _assert_span_link(b3_span, b_span, "output", "output") + _assert_span_link(b_span, b1_span, "input", "input") + _assert_span_link(b1_span, b2_span, "output", "input") + _assert_span_link(b2_span, b3_span, "output", "input") async def test_subgraph_async(self, llmobs_events, complex_graph): await complex_graph.ainvoke({"a_list": [], "which": "b"}) - graph_span = llmobs_events[5] - a_span = llmobs_events[0] - b_span = llmobs_events[4] - b1_span = llmobs_events[1] - b2_span = llmobs_events[2] - b3_span = llmobs_events[3] - - assert graph_span["name"] == "LangGraph" - _assert_span_link(graph_span, None, "input", "input") - _assert_span_link(graph_span, b_span, "output", "output") - assert a_span["name"] == "a" - _assert_span_link(a_span, graph_span, "input", "input") - assert b_span["name"] == "b" - _assert_span_link(b_span, a_span, "output", "input") - _assert_span_link(b_span, b3_span, "output", "output") - assert b1_span["name"] == "b1" - _assert_span_link(b1_span, b_span, "input", "input") - assert b2_span["name"] == "b2" - _assert_span_link(b2_span, b1_span, "output", "input") - assert b3_span["name"] == "b3" - _assert_span_link(b3_span, b2_span, "output", "input") + graph_span = _find_span_by_name(llmobs_events, "LangGraph") + a_span = _find_span_by_name(llmobs_events, "a") + b_span = _find_span_by_name(llmobs_events, "b") + b1_span = _find_span_by_name(llmobs_events, "b1") + b2_span = _find_span_by_name(llmobs_events, "b2") + b3_span = _find_span_by_name(llmobs_events, "b3") + + _assert_span_link(None, graph_span, "input", "input") + _assert_span_link(b3_span, b_span, "output", "output") + _assert_span_link(graph_span, a_span, "input", "input") + _assert_span_link(a_span, b_span, "output", "input") + _assert_span_link(b3_span, b_span, "output", "output") + _assert_span_link(b_span, b1_span, "input", "input") + _assert_span_link(b1_span, b2_span, "output", "input") + _assert_span_link(b2_span, b3_span, "output", "input") def test_fanning_graph(self, llmobs_events, fanning_graph): fanning_graph.invoke({"a_list": [], "which": ""}) - graph_span = llmobs_events[4] - a_span = llmobs_events[0] - b_span = llmobs_events[1] - c_span = llmobs_events[2] - d_span = llmobs_events[3] - - assert graph_span["name"] == "LangGraph" - _assert_span_link(graph_span, None, "input", "input") - _assert_span_link(graph_span, d_span, "output", "output") - assert a_span["name"] == "a" - _assert_span_link(a_span, graph_span, "input", "input") - assert b_span["name"] == "b" - _assert_span_link(b_span, a_span, "output", "input") - assert c_span["name"] == "c" - _assert_span_link(c_span, a_span, "output", "input") - assert d_span["name"] == "d" - _assert_span_link(d_span, b_span, "output", "input") - _assert_span_link(d_span, c_span, "output", "input") + graph_span = _find_span_by_name(llmobs_events, "LangGraph") + a_span = _find_span_by_name(llmobs_events, "a") + b_span = _find_span_by_name(llmobs_events, "b") + c_span = _find_span_by_name(llmobs_events, "c") + d_span = _find_span_by_name(llmobs_events, "d") + + _assert_span_link(None, graph_span, "input", "input") + _assert_span_link(d_span, graph_span, "output", "output") + _assert_span_link(graph_span, a_span, "input", "input") + _assert_span_link(a_span, b_span, "output", "input") + _assert_span_link(a_span, c_span, "output", "input") + _assert_span_link(b_span, d_span, "output", "input") + _assert_span_link(c_span, d_span, "output", "input") async def test_fanning_graph_async(self, llmobs_events, fanning_graph): await fanning_graph.ainvoke({"a_list": [], "which": ""}) - graph_span = llmobs_events[4] - a_span = llmobs_events[0] - b_span = llmobs_events[1] - c_span = llmobs_events[2] - d_span = llmobs_events[3] + graph_span = _find_span_by_name(llmobs_events, "LangGraph") + a_span = _find_span_by_name(llmobs_events, "a") + b_span = _find_span_by_name(llmobs_events, "b") + c_span = _find_span_by_name(llmobs_events, "c") + d_span = _find_span_by_name(llmobs_events, "d") assert graph_span["name"] == "LangGraph" - _assert_span_link(graph_span, None, "input", "input") - _assert_span_link(graph_span, d_span, "output", "output") - assert a_span["name"] == "a" - _assert_span_link(a_span, graph_span, "input", "input") - assert b_span["name"] == "b" - _assert_span_link(b_span, a_span, "output", "input") - assert c_span["name"] == "c" - _assert_span_link(c_span, a_span, "output", "input") - assert d_span["name"] == "d" - _assert_span_link(d_span, b_span, "output", "input") - _assert_span_link(d_span, c_span, "output", "input") + _assert_span_link(None, graph_span, "input", "input") + _assert_span_link(d_span, graph_span, "output", "output") + _assert_span_link(graph_span, a_span, "input", "input") + _assert_span_link(a_span, b_span, "output", "input") + _assert_span_link(a_span, c_span, "output", "input") + _assert_span_link(b_span, d_span, "output", "input") + _assert_span_link(c_span, d_span, "output", "input")