diff --git a/packages/client/tests/test_agui.py b/packages/client/tests/test_agui.py index 915e6944..68613a08 100644 --- a/packages/client/tests/test_agui.py +++ b/packages/client/tests/test_agui.py @@ -149,6 +149,16 @@ def test_fail_open_masks_a_call_that_never_finished(self) -> None: assert events[-1]["content"] == "" assert events[-1]["metadata"]["launchdarkly"]["toolFailed"] is True + def test_tool_end_extends_a_streamed_argument_prefix(self) -> None: + emitter = TurnEmitter(capture_content=True) + events = emitter.tool_start("a", "search") + events.extend(emitter.tool_args("a", '{"q":')) + events.extend(emitter.tool_end("a", arguments={"q": "x"})) + deltas = [ + event["delta"] for event in events if event["type"] == "TOOL_CALL_ARGS" + ] + assert "".join(deltas) == '{"q":"x"}' + def test_a_second_start_for_the_same_id_is_dropped(self) -> None: emitter = TurnEmitter(capture_content=True) assert emitter.tool_start("a", "search") diff --git a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py index 9bdab685..4f14b1ae 100644 --- a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py +++ b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py @@ -31,6 +31,7 @@ set_input_content_attributes, set_output_content_attributes, ) +from launchdarkly_ai_server.agui import TurnEmitter from launchdarkly_ai_server.lifecycle import effective_capture_content from ._version import PACKAGE_NAME, __version__ @@ -474,6 +475,7 @@ async def _stream_gen( # stream; a CancelledError means something cancelled the run, usually a timeout, and the # consumer chose nothing. The blocking path already tells these apart. cancelled = False + emitter = TurnEmitter(capture_content=capture_content) try: # Inside the guard, because serialising the prompt raises on anything that is not # JSON-serialisable. A raise out here would leave the root open with the `finally` never @@ -519,8 +521,36 @@ async def _stream_gen( if getattr(msg, "type", None) == "ai": text = lang_chain_content_text(msg.content) if text: - yield {"type": "chunk", "text": text} + for rich in emitter.text(text): + yield rich full_output = text + for tc in getattr(msg, "tool_calls", None) or []: + if not isinstance(tc, dict): + continue + call_id = str(tc.get("id") or tc.get("name") or "") + name = str(tc.get("name") or "") + for rich in emitter.tool_start(call_id, name): + yield rich + for rich in emitter.tool_end( + call_id, arguments=tc.get("args") + ): + yield rich + elif getattr(msg, "type", None) == "tool": + call_id = str(getattr(msg, "tool_call_id", "") or "") + if call_id: + for rich in emitter.tool_result( + call_id, + getattr(msg, "content", None), + failed=getattr(msg, "status", None) == "error", + ): + yield rich + for rich in emitter.end_text(): + yield rich + + # A step can start a call the graph never answers with a tool message. + # Close it before the run finishes, or the client keeps it open. + for rich in emitter.fail_open(): + yield rich # The same reconciliation the blocking path does, for the same reason. This walk reads # usage_metadata off the astream payloads, and sees nothing at all unless the graph is @@ -562,6 +592,8 @@ async def _stream_gen( if span_callbacks.run_usage.reported: finish_root_span(span, config, span_callbacks.run_usage.total) fail_span(span, exc, ended) + for rich in emitter.fail_open(): + yield rich raise finally: # A no-op on the success and failure paths, because both already ended their spans through diff --git a/packages/langchain-agents/tests/test_handler.py b/packages/langchain-agents/tests/test_handler.py index 15d59ef2..07bf45a2 100644 --- a/packages/langchain-agents/tests/test_handler.py +++ b/packages/langchain-agents/tests/test_handler.py @@ -1427,7 +1427,11 @@ async def _mock_astream(*a: Any, **kw: Any) -> AsyncIterator[Any]: llm=MagicMock() ).stream(_make_config(), "q") ] - assert {"type": "chunk", "text": "visible answer"} in events + assert any( + event.get("type") == "TEXT_MESSAGE_CONTENT" + and event.get("delta") == "visible answer" + for event in events + ) assert events[-1]["output"] == "visible answer" @pytest.mark.asyncio @@ -1447,6 +1451,111 @@ async def _bad_astream(*a: Any, **kw: Any) -> AsyncIterator[Any]: async for _ in await h.stream(_make_config(), "hi"): pass + @pytest.mark.asyncio + async def test_a_failed_run_closes_started_tool_calls(self) -> None: + mocks = _make_langchain_mock() + msg = _make_ai_msg("looking") + msg.tool_calls = [ + {"name": "search", "id": "c1", "args": {"q": "x"}}, + {"name": "other", "id": "d1", "args": {"n": 1}}, + ] + + async def _astream(*_a: Any, **_kw: Any) -> AsyncIterator[Any]: + yield {"agent": {"messages": [msg]}} + raise RuntimeError("tool failed") + + mocks["_agent"].astream = _astream + seen: list[dict[str, Any]] = [] + with ( + _patch_lc(mocks), + patch.object(spans_mod, "_HAS_OTEL", False), + pytest.raises(RuntimeError, match="tool failed"), + ): + handler = create_langchain_agents_handler(llm=MagicMock()) + async for event in await handler.stream(_make_config(), "q"): + seen.append(event) + results = { + event["toolCallId"] for event in seen if event["type"] == "TOOL_CALL_RESULT" + } + assert results == {"c1", "d1"} + + @pytest.mark.asyncio + async def test_tool_messages_close_parallel_calls(self) -> None: + mocks = _make_langchain_mock() + ai = _make_ai_msg("looking") + ai.tool_calls = [ + {"name": "search", "id": "c1", "args": {"q": "x"}}, + {"name": "other", "id": "d1", "args": {"n": 1}}, + ] + + def tool(call_id: str, content: str) -> Any: + message = MagicMock() + message.type = "tool" + message.tool_call_id = call_id + message.content = content + message.status = "success" + message.usage_metadata = None + return message + + async def _astream(*_a: Any, **_kw: Any) -> AsyncIterator[Any]: + yield {"agent": {"messages": [ai]}} + yield {"tools": {"messages": [tool("c1", "found"), tool("d1", "also")]}} + + mocks["_agent"].astream = _astream + with _patch_lc(mocks), patch.object(spans_mod, "_HAS_OTEL", False): + events = [ + event + async for event in await create_langchain_agents_handler( + llm=MagicMock(), capture_content=True + ).stream(_make_config(), "q") + ] + by_id: dict[str, list[str]] = {} + for event in events: + if not str(event.get("type", "")).startswith("TOOL_CALL"): + continue + by_id.setdefault(event["toolCallId"], []).append(event["type"]) + assert by_id["c1"] == [ + "TOOL_CALL_START", + "TOOL_CALL_ARGS", + "TOOL_CALL_END", + "TOOL_CALL_RESULT", + ] + assert by_id["d1"] == [ + "TOOL_CALL_START", + "TOOL_CALL_ARGS", + "TOOL_CALL_END", + "TOOL_CALL_RESULT", + ] + contents = { + event["toolCallId"]: event["content"] + for event in events + if event["type"] == "TOOL_CALL_RESULT" + } + assert contents == {"c1": "found", "d1": "also"} + + @pytest.mark.asyncio + async def test_a_successful_run_closes_a_call_with_no_tool_message(self) -> None: + mocks = _make_langchain_mock() + ai = _make_ai_msg("looking") + ai.tool_calls = [{"name": "search", "id": "c1", "args": {"q": "x"}}] + + async def _astream(*_a: Any, **_kw: Any) -> AsyncIterator[Any]: + yield {"agent": {"messages": [ai]}} + + mocks["_agent"].astream = _astream + with _patch_lc(mocks), patch.object(spans_mod, "_HAS_OTEL", False): + events = [ + event + async for event in await create_langchain_agents_handler( + llm=MagicMock() + ).stream(_make_config(), "q") + ] + result = next(event for event in events if event["type"] == "TOOL_CALL_RESULT") + assert result["toolCallId"] == "c1" + assert result["metadata"]["launchdarkly"]["toolFailed"] is True + assert events[-1]["type"] == "done" + assert events.index(result) < len(events) - 1 + # --------------------------------------------------------------------------- # §1.5 Streaming telemetry (Appendix A.5 — do not patch _HAS_OTEL=False) @@ -2586,5 +2695,6 @@ def factory(config: Any) -> Any: ] assert seen[0]["model"]["parameters"] == {"temperature": 0.2} assert any( - e.get("type") == "chunk" and e.get("text") == "streamed" for e in events + e.get("type") == "TEXT_MESSAGE_CONTENT" and e.get("delta") == "streamed" + for e in events ) diff --git a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py index 95390dfa..47997a32 100644 --- a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py +++ b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py @@ -32,6 +32,7 @@ set_output_content_attributes, set_tool_call_content_attributes, ) +from launchdarkly_ai_server.agui import TurnEmitter from launchdarkly_ai_server.lifecycle import effective_capture_content from ._version import PACKAGE_NAME, __version__ @@ -151,6 +152,21 @@ def append_turn(turn: dict[str, Any]) -> None: return messages +def _merge_stream_chunk(gathered: Any, chunk: Any) -> Any: + """Fold one streamed chunk into the turn the way LangChain does. + + ``AIMessageChunk.tool_calls`` is parsed from that chunk alone. Keeping only + the last chunk drops every earlier parallel call and every argument slice + that chunk did not repeat. Addition merges chunks by index. + """ + if gathered is None: + return chunk + try: + return gathered + chunk + except TypeError: + return chunk + + def _assistant_output_messages( content: Any, tool_calls: list[Any] | None ) -> list[SpanMessage]: @@ -761,6 +777,7 @@ async def _stream_gen( conversation_messages = list(initial_messages) full_output: Any = "" steps = 0 + emitter = TurnEmitter(capture_content=capture_content) ended: set[int] = set() open_model_span: Any = None @@ -797,7 +814,9 @@ async def _stream_gen( open_model_span = model_span accumulated_content = "" - accumulated_tool_calls: list[Any] = [] + gathered: Any = None + tool_chunk_ids: dict[int, str] = {} + seen_call_ids: list[str] = [] # The cache breakdown has to be accumulated alongside the scalars. LangChain reports it # per chunk in `usage_metadata.input_token_details`, and synthesizing a usage bag # without it would make the streaming span emit `cache_read = 0` where the blocking path @@ -830,8 +849,60 @@ async def _stream_gen( async for chunk in chunk_stream: text = lang_chain_content_text(chunk.content) if text: - yield {"type": "chunk", "text": text} + for rich in emitter.text(text): + yield rich accumulated_content += text + for piece in getattr(chunk, "tool_call_chunks", None) or []: + piece_name = ( + piece.get("name") + if isinstance(piece, dict) + else getattr(piece, "name", None) + ) + piece_id = ( + piece.get("id") + if isinstance(piece, dict) + else getattr(piece, "id", None) + ) + piece_index = ( + piece.get("index") + if isinstance(piece, dict) + else getattr(piece, "index", None) + ) + piece_args = ( + piece.get("args") + if isinstance(piece, dict) + else getattr(piece, "args", None) + ) + # Later chunks often omit id. They carry an index, or, + # when only one call is open, just the args slice. + if isinstance(piece_id, str) and piece_id: + if piece_id not in seen_call_ids: + seen_call_ids.append(piece_id) + if isinstance(piece_index, int): + tool_chunk_ids[piece_index] = piece_id + elif ( + isinstance(piece_index, int) + and piece_index in tool_chunk_ids + ): + piece_id = tool_chunk_ids[piece_index] + elif len(seen_call_ids) == 1: + piece_id = seen_call_ids[0] + if ( + isinstance(piece_name, str) + and piece_name + and isinstance(piece_id, str) + and piece_id + ): + for rich in emitter.tool_start(piece_id, piece_name): + yield rich + if ( + isinstance(piece_id, str) + and piece_id + and isinstance(piece_args, str) + and piece_args + ): + for rich in emitter.tool_args(piece_id, piece_args): + yield rich usage = getattr(chunk, "usage_metadata", None) if usage: usage_reported = True @@ -844,9 +915,7 @@ async def _stream_gen( turn_usage.cache_creation += number_or_zero( details.get("cache_creation") ) - chunk_tools = getattr(chunk, "tool_calls", None) or [] - if chunk_tools: - accumulated_tool_calls = chunk_tools + gathered = _merge_stream_chunk(gathered, chunk) finish_reasons = lang_chain_finish_reasons(chunk) or finish_reasons open_chunk_stream = None except Exception as exc: @@ -871,6 +940,11 @@ async def _stream_gen( # billed these tokens: a content failure is our problem and must not report the run as # having spent less than it did. run_usage.add(turn_usage if usage_reported else None) + accumulated_tool_calls = [ + tc + for tc in (getattr(gathered, "tool_calls", None) or []) + if isinstance(tc, dict) and tc.get("name") + ] # Inside a guard for the same reason as the blocking path: a raise while serialising # completion content would otherwise leave this span for `finally` to end as abandoned, @@ -894,6 +968,13 @@ async def _stream_gen( raise if not accumulated_tool_calls: + for rich in emitter.end_text(): + yield rich + # A chunk can open a call that never becomes a complete tool + # call. Close it before the turn ends, or the client keeps it + # open across RUN_FINISHED. + for rich in emitter.fail_open(): + yield rich full_output = (full_output or "") + accumulated_content break @@ -904,6 +985,14 @@ async def _stream_gen( steps += 1 full_output = (full_output or "") + accumulated_content + for tc in accumulated_tool_calls: + call_id = str(tc.get("id") or tc["name"]) + for rich in emitter.tool_start(call_id, str(tc["name"])): + yield rich + for rich in emitter.tool_end(call_id, arguments=tc.get("args")): + yield rich + for rich in emitter.end_text(): + yield rich assistant_msg = AIMessage( content=accumulated_content, tool_calls=accumulated_tool_calls ) @@ -944,11 +1033,19 @@ async def _stream_gen( # Cleared on both paths that end the span, and deliberately not in a `finally`: # a `finally` would also clear it for a BaseException, which is the one case where # the span is still open and the outer `finally` is the only thing left to close it. + for rich in emitter.tool_result( + str(tc.get("id") or tc["name"]), result_val + ): + yield rich tool_results.append( ToolMessage( tool_call_id=tc.get("id") or tc["name"], content=str(result_val) ) ) + # Chunks can start a call the gathered message then omits. The + # calls we ran are already closed; this finishes the rest. + for rich in emitter.fail_open(): + yield rich conversation_messages.extend(tool_results) full_output_str = ( @@ -989,6 +1086,10 @@ async def _stream_gen( if run_usage.reported: finish_root_span(span, config, run_usage.total) fail_span(span, exc, ended) + for rich in emitter.end_text(): + yield rich + for rich in emitter.fail_open(): + yield rich raise finally: # A no-op on the success and failure paths, because both already ended their spans through diff --git a/packages/langchain-messages/tests/test_handler.py b/packages/langchain-messages/tests/test_handler.py index 32768dff..57c6a06a 100644 --- a/packages/langchain-messages/tests/test_handler.py +++ b/packages/langchain-messages/tests/test_handler.py @@ -1178,10 +1178,231 @@ async def test_yields_chunk_events(self) -> None: llm = self._make_streaming_llm(["hello ", "world"]) h = create_langchain_messages_handler(llm=llm) events = [e async for e in await h.stream(CONFIG, "q")] - chunks = [e for e in events if e.get("type") == "chunk"] - assert len(chunks) == 2 - assert chunks[0]["text"] == "hello " - assert chunks[1]["text"] == "world" + deltas = [e["delta"] for e in events if e.get("type") == "TEXT_MESSAGE_CONTENT"] + assert deltas == ["hello ", "world"] + + async def test_later_tool_arg_chunks_keep_the_first_id(self) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + first = FakeAIMessage("") + first.tool_call_chunks = [ + {"name": "search", "id": "c1", "index": 0, "args": '{"q":'} + ] + second = FakeAIMessage("") + second.tool_call_chunks = [{"index": 0, "args": '"x"}'}] + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + + async def _astream(_msgs: Any) -> AsyncGenerator[Any, None]: + yield first + yield second + + llm.astream = _astream + llm.ainvoke = AsyncMock(return_value=FakeAIMessage("")) + handler = create_langchain_messages_handler(llm=llm, capture_content=True) + events = [event async for event in await handler.stream(CONFIG, "q")] + args = "".join( + event["delta"] for event in events if event["type"] == "TOOL_CALL_ARGS" + ) + assert args == '{"q":"x"}' + + async def test_a_later_arg_chunk_without_an_index_keeps_the_only_id(self) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + first = FakeAIMessage("") + first.tool_call_chunks = [{"name": "search", "id": "c1", "args": '{"q":'}] + second = FakeAIMessage("") + second.tool_call_chunks = [{"args": '"x"}'}] + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + + async def _astream(_msgs: Any) -> AsyncGenerator[Any, None]: + yield first + yield second + + llm.astream = _astream + llm.ainvoke = AsyncMock(return_value=FakeAIMessage("")) + handler = create_langchain_messages_handler(llm=llm, capture_content=True) + events = [event async for event in await handler.stream(CONFIG, "q")] + args = "".join( + event["delta"] for event in events if event["type"] == "TOOL_CALL_ARGS" + ) + assert args == '{"q":"x"}' + result = next(event for event in events if event["type"] == "TOOL_CALL_RESULT") + assert result["toolCallId"] == "c1" + assert result["metadata"]["launchdarkly"]["toolFailed"] is True + assert events.index(result) < len(events) - 1 + assert events[-1]["type"] == "done" + + async def test_a_gathered_call_list_that_omits_a_started_call_closes_it( + self, + ) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + first = FakeAIMessage("") + first.tool_call_chunks = [ + {"name": "search", "id": "c1", "args": "{}"}, + {"name": "other", "id": "c2", "args": "{}"}, + ] + second = FakeAIMessage( + "", + tool_calls=[{"name": "search", "id": "c1", "args": {}}], + ) + done = FakeAIMessage("ok") + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + turns = [[first, second], [done]] + + async def _astream(_msgs: Any) -> AsyncGenerator[Any, None]: + for chunk in turns.pop(0): + yield chunk + + llm.astream = _astream + handler = create_langchain_messages_handler(llm=llm, capture_content=True) + events = [ + event + async for event in await handler.stream( + CONFIG, "q", {"search": lambda _args: "found"} + ) + ] + results = { + event["toolCallId"]: event + for event in events + if event["type"] == "TOOL_CALL_RESULT" + } + assert results["c1"]["content"] == "found" + assert "metadata" not in results["c1"] + assert results["c2"]["metadata"]["launchdarkly"]["toolFailed"] is True + ok = next( + event + for event in events + if event["type"] == "TEXT_MESSAGE_CONTENT" and event["delta"] == "ok" + ) + assert events.index(results["c2"]) < events.index(ok) + + async def test_a_provider_error_closes_the_open_text(self) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + + async def _astream(_msgs: Any) -> AsyncGenerator[Any, None]: + yield FakeAIMessage("Hello") + raise RuntimeError("stream broke") + + llm.astream = _astream + seen: list[dict[str, Any]] = [] + handler = create_langchain_messages_handler(llm=llm) + with pytest.raises(RuntimeError, match="stream broke"): + async for event in await handler.stream(CONFIG, "q"): + seen.append(event) + start = next(event for event in seen if event["type"] == "TEXT_MESSAGE_START") + end = next(event for event in seen if event["type"] == "TEXT_MESSAGE_END") + assert end["messageId"] == start["messageId"] + + async def test_tool_end_completes_a_truncated_argument_prefix(self) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + first = FakeAIMessage("") + first.tool_call_chunks = [{"name": "search", "id": "c1", "args": '{"q":'}] + second = FakeAIMessage( + "", + tool_calls=[{"name": "search", "id": "c1", "args": {"q": "x"}}], + ) + done = FakeAIMessage("ok") + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + turns = [first, second] + + async def _astream(_msgs: Any) -> AsyncGenerator[Any, None]: + if turns: + yield turns.pop(0) + if turns: + yield turns.pop(0) + return + yield done + + llm.astream = _astream + llm.ainvoke = AsyncMock(return_value=FakeAIMessage("")) + handler = create_langchain_messages_handler(llm=llm, capture_content=True) + events = [ + event + async for event in await handler.stream( + CONFIG, "q", {"search": lambda _args: "found"} + ) + ] + args = "".join( + event["delta"] for event in events if event["type"] == "TOOL_CALL_ARGS" + ) + assert args == '{"q":"x"}' + + async def test_parallel_chunks_use_the_gathered_tool_calls(self) -> None: + from langchain_core.messages import AIMessageChunk + + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + first = AIMessageChunk( + content="", + tool_call_chunks=[ + {"name": "search", "args": "", "id": "call_1", "index": 0} + ], + ) + second = AIMessageChunk( + content="", + tool_call_chunks=[ + {"name": None, "args": '{"q":', "id": None, "index": 0}, + {"name": "other", "args": "", "id": "call_2", "index": 1}, + ], + ) + third = AIMessageChunk( + content="", + tool_call_chunks=[ + {"name": None, "args": '"x"}', "id": None, "index": 0}, + {"name": None, "args": "{}", "id": None, "index": 1}, + ], + ) + seen: list[dict[str, Any]] = [] + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + first_turn = True + + async def _astream(_msgs: Any) -> AsyncGenerator[Any, None]: + nonlocal first_turn + if first_turn: + first_turn = False + yield first + yield second + yield third + return + yield FakeAIMessage("ok") + + llm.astream = _astream + llm.ainvoke = AsyncMock(return_value=FakeAIMessage("")) + + def search(args: dict[str, Any]) -> str: + seen.append(args) + return "found" + + def other(args: dict[str, Any]) -> str: + seen.append(args) + return "also" + + handler = create_langchain_messages_handler(llm=llm, capture_content=True) + events = [ + event + async for event in await handler.stream( + CONFIG, "q", {"search": search, "other": other} + ) + ] + starts = [ + event["toolCallId"] + for event in events + if event["type"] == "TOOL_CALL_START" + ] + assert starts == ["call_1", "call_2"] + assert seen == [{"q": "x"}, {}] + assert events[-1]["type"] == "done" + assert "ok" in events[-1]["output"] async def test_streams_text_while_ignoring_thinking_blocks(self) -> None: from launchdarkly_ai_langchain_messages import create_langchain_messages_handler @@ -1200,7 +1421,13 @@ async def test_streams_text_while_ignoring_thinking_blocks(self) -> None: CONFIG, "q" ) ] - assert {"type": "chunk", "text": "visible answer"} in events + assert { + "type": "TEXT_MESSAGE_CONTENT", + "messageId": events[1]["messageId"] + if events[0]["type"] == "TEXT_MESSAGE_START" + else events[0]["messageId"], + "delta": "visible answer", + } in events assert events[-1]["output"] == "visible answer" async def test_yields_exactly_one_done_event(self) -> None: @@ -2608,5 +2835,6 @@ def factory(config: Any) -> Any: ] assert seen[0]["model"]["parameters"] == {"temperature": 0.2} assert any( - e.get("type") == "chunk" and e.get("text") == "streamed" for e in events + e.get("type") == "TEXT_MESSAGE_CONTENT" and e.get("delta") == "streamed" + for e in events )