diff --git a/packages/langchain-agents/pyproject.toml b/packages/langchain-agents/pyproject.toml index 9dc7c3e1..edcf496b 100644 --- a/packages/langchain-agents/pyproject.toml +++ b/packages/langchain-agents/pyproject.toml @@ -9,6 +9,7 @@ dependencies = [ "langgraph>=1.2", "langgraph-prebuilt>=1.1", "langchain-openai>=1.3", + "pydantic>=2.13", ] description = "LangChain/LangGraph agent handler for LaunchDarkly AI SDK" readme = "README.md" diff --git a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/native_graph.py b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/native_graph.py index 8bc1a375..cdae0eb8 100644 --- a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/native_graph.py +++ b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/native_graph.py @@ -6,12 +6,16 @@ from __future__ import annotations import inspect +import json import re import time import uuid from collections.abc import Awaitable, Callable, Mapping from typing import Annotated, Any, TypedDict +from langchain_core.tools import InjectedToolCallId +from pydantic import BaseModel + from launchdarkly_ai_server import ( GraphDefinition, GraphNode, @@ -54,6 +58,82 @@ _HAS_LANGGRAPH = False +_HANDOFF_IGNORED = "Ignored: only the first transfer in a turn routes." + + +class _HandoffArgs(BaseModel): + """Hand-off tool schema: the model passes no arguments, and LangChain injects + the call id (a dict schema would skip the injection).""" + + tool_call_id: Annotated[str, InjectedToolCallId] + + +def _content_text(content: Any) -> str: + """The text parts of a message's content, joined with no separator.""" + if isinstance(content, str): + return content + if isinstance(content, list): + return "".join( + part.get("text") or "" + for part in content + if isinstance(part, dict) and part.get("type") == "text" + ) + return "" + + +# Message `type` values, chunks included (an AIMessageChunk's type is not "ai"). +_AI_TYPES = frozenset({"ai", "AIMessageChunk"}) +_TOOL_TYPES = frozenset({"tool", "ToolMessageChunk"}) + + +def _to_neutral_history( + messages: list[Any], ai_message: Any, human_message: Any +) -> list[Any]: + """Rewrite the history a node finds when it is entered as provider-neutral + text, so a target on another provider never sees the previous node's tool + or reasoning blocks (spec §2.x.3). Human and system messages pass through.""" + out: list[Any] = [] + calls: list[dict[str, Any]] = [] + results: list[Any] = [] + + def _flush() -> None: + if not results: + return + order = {c.get("id"): i for i, c in enumerate(calls)} + lines = [] + for r in sorted(results, key=lambda r: order.get(r.tool_call_id, len(calls))): + i = order.get(r.tool_call_id) + name = calls[i]["name"] if i is not None else getattr(r, "name", "") or "" + lines.append(f"[Tool result] {name}: {_content_text(r.content)}") + out.append(human_message("\n".join(lines))) + results.clear() + + for m in messages: + if m.type in _TOOL_TYPES: + results.append(m) + continue + _flush() + if m.type not in _AI_TYPES: + out.append(m) + continue + calls = list(getattr(m, "tool_calls", None) or []) + text = _content_text(m.content) + lines = [text] if text else [] + lines += [ + "[Tool call] {} {}".format( + c["name"], + json.dumps( + c.get("args") or {}, separators=(",", ":"), ensure_ascii=False + ), + ) + for c in calls + ] + if lines: + out.append(ai_message("\n".join(lines))) + _flush() + return out + + def _sanitize_name(key: str) -> str: return re.sub(r"[^a-z0-9_-]", "_", key, flags=re.IGNORECASE) @@ -173,6 +253,8 @@ async def invoke( HumanMessage = lc_msgs.HumanMessage SystemMessage = lc_msgs.SystemMessage + ToolMessage = lc_msgs.ToolMessage + AIMessage = lc_msgs.AIMessage vs = variables or {} def_obj: GraphDefinition = await def_promise @@ -257,26 +339,55 @@ async def _traverse_node(node: GraphNode) -> None: langgraph_types = importlib.import_module("langgraph.types") Command = langgraph_types.Command + # The node's last model call sets "first" (its first hand-off call + # id) and "call_ids"; the tools step sets "handed_off" when that + # call routes. "start" is where the node's current visit began. + routing: dict[str, Any] = { + "first": None, + "handed_off": False, + "call_ids": set(), + "start": 0, + } + handoff_tools = [] for edge in outgoing: target_key = _sanitize_name(edge.target_key) async def _handoff_exec( + tool_call_id: Annotated[str, InjectedToolCallId], _target: str = target_key, + _target_raw: str = edge.target_key, _node: GraphNode = node, ) -> Any: + # Only the first hand-off in a model message routes, as in graph(). + if tool_call_id != routing["first"]: + return _HANDOFF_IGNORED + routing["handed_off"] = True if ld_context: td = make_track_data(_node, def_obj.key, run_id) get_client().track( "$ld:ai:graph:handoff_success", ld_context, td, 1 ) - return Command(goto=_target) + # The ToolMessage answers the call, so the next model call + # never sees an unanswered tool call. + return Command( + goto=_target, + update={ + "messages": [ + ToolMessage( + content=f"Handed off to {_target_raw}.", + tool_call_id=tool_call_id, + name=f"transfer_to_{_target}", + ) + ] + }, + ) ht = tool_fn( f"transfer_to_{_sanitize_name(edge.target_key)}", _handoff_exec, description=f"Transfer control to the {edge.target_key} agent", - args_schema={}, + args_schema=_HandoffArgs, ) handoff_tools.append(ht) @@ -299,25 +410,48 @@ async def _node_fn( node_start = time.monotonic() system_prompt = _build_system_prompt(_node, vs) - conv_messages: list[Any] = state.get("messages", []) + state_messages: list[Any] = state.get("messages", []) + last = state_messages[-1] if state_messages else None + # A tool result for this node's own last call continues its + # tool loop; anything else means the node was just entered. + if not ( + last is not None + and last.type in _TOOL_TYPES + and last.tool_call_id in routing["call_ids"] + ): + routing["start"] = len(state_messages) + conv_messages = [ + *_to_neutral_history( + state_messages[: routing["start"]], AIMessage, HumanMessage + ), + *state_messages[routing["start"] :], + ] full_messages = ( [SystemMessage(system_prompt), *conv_messages] if system_prompt else list(conv_messages) ) + # No parallel_tool_calls: Bedrock Converse rejects it. The + # hand-off tools keep only the first hand-off instead. bound = ( - chat_model.bind_tools( - all_tools, - **( - {"parallel_tool_calls": False} if is_multi_child else {} - ), - ) - if all_tools - else chat_model + chat_model.bind_tools(all_tools) if all_tools else chat_model ) result_msg = await bound.ainvoke(full_messages) + routing["first"] = next( + ( + tc.get("id") + for tc in getattr(result_msg, "tool_calls", None) or [] + if tc.get("name", "").startswith("transfer_to_") + ), + None, + ) + routing["handed_off"] = False + routing["call_ids"] = { + tc.get("id") + for tc in getattr(result_msg, "tool_calls", None) or [] + } usage = _extract_usage(result_msg) total_usage["input"] += usage["input"] total_usage["output"] += usage["output"] @@ -363,27 +497,28 @@ async def _node_fn( builder.add_edge(f"{node_key}_tools", node_key) else: builder.add_edge(node_key, END) - elif is_multi_child: - if all_tools: - builder.add_conditional_edges( - node_key, - tools_condition, - {"tools": f"{node_key}_tools", "__end__": END}, - ) - builder.add_edge(f"{node_key}_tools", node_key) - else: - builder.add_edge(node_key, END) else: - child_key = _sanitize_name(outgoing[0].target_key) - if all_tools: - builder.add_conditional_edges( - node_key, - tools_condition, - {"tools": f"{node_key}_tools", "__end__": child_key}, - ) - builder.add_edge(f"{node_key}_tools", node_key) - else: - builder.add_edge(node_key, child_key) + # A non-terminal node always has hand-off tools. Ending the turn + # without one ends the run (multi-child) or moves to the child. + builder.add_conditional_edges( + node_key, + tools_condition, + { + "tools": f"{node_key}_tools", + "__end__": END + if is_multi_child + else _sanitize_name(outgoing[0].target_key), + }, + ) + # No static edge back to the node: after a hand-off the Command + # routes to the target and the node must not run again. + builder.add_conditional_edges( + f"{node_key}_tools", + lambda _s, _r=routing, _k=node_key: ( + END if _r["handed_off"] else _k + ), + [node_key, END], + ) # Pre-order traversal (root first, to mirror TS traverse) pre_visited: set[str] = set() @@ -424,18 +559,7 @@ async def _pre_visit(node_key: str) -> None: last_msg = result_messages[-1] if result_messages else None def _content_str(msg: Any) -> str: - if msg is None: - return "" - c = msg.content - if isinstance(c, str): - return c - if isinstance(c, list): - return "".join( - part.get("text", "") if isinstance(part, dict) else "" - for part in c - if isinstance(part, dict) and part.get("type") == "text" - ) - return "" + return "" if msg is None else _content_text(msg.content) final_output = _content_str(last_msg) diff --git a/packages/langchain-agents/tests/test_native_graph.py b/packages/langchain-agents/tests/test_native_graph.py index 67d68622..ad835137 100644 --- a/packages/langchain-agents/tests/test_native_graph.py +++ b/packages/langchain-agents/tests/test_native_graph.py @@ -17,6 +17,7 @@ from launchdarkly_ai_langchain_agents.native_graph import ( _build_node_tools, _extract_usage, + _to_neutral_history, to_lang_graph, ) from launchdarkly_ai_server import GraphDefinition, GraphEdge, GraphNode @@ -1353,3 +1354,372 @@ async def _visit(fn: Any, ctx: Any = None) -> None: expected = LANGCHAIN_CHAT_OPENAI candidates = candidate_keys(_REAL_CHAT_OPENAI.model_fields, expected) assert await probe_forwarded_keys(candidates, call) == expected + + +# --------------------------------------------------------------------------- +# Hand-offs through a real StateGraph (AIC-3592) +# --------------------------------------------------------------------------- + + +def _handoff_graph_def( + children: list[str], root_tools: dict[str, Any] | None = None +) -> GraphDefinition: + def _node( + key: str, edges: list[str], tools: dict[str, Any] | None = None + ) -> dict[str, Any]: + config: dict[str, Any] = { + "model": {"name": "m"}, + "instructions": f"you are {key}", + } + if tools: + config["tools"] = tools + return { + "key": key, + "config": config, + "meta": {"variationKey": "v1", "version": 1}, + "edges": [{"source_key": key, "target_key": c} for c in edges], + "is_terminal": not edges, + } + + nodes = {"router": _node("router", children, root_tools)} + nodes.update({c: _node(c, []) for c in children}) + return _make_graph_def( + nodes=nodes, + edges=[{"source_key": "router", "target_key": c} for c in children], + root_key="router", + ) + + +def _scripted_models(replies: dict[str, list[Any]]) -> tuple[Any, dict[str, Any]]: + """One mock chat model per node; each ainvoke pops that node's next reply. + bind_tools is a MagicMock, so a test can assert on its kwargs.""" + models: dict[str, Any] = {} + for key, msgs in replies.items(): + m = MagicMock() + m.ainvoke = AsyncMock(side_effect=list(msgs)) + m.bind_tools = MagicMock(return_value=m) + models[key] = m + return (lambda node: models[node.key]), models + + +def _tool_call_msg(*calls: tuple[Any, ...], content: Any = "") -> Any: + """An AI message with tool calls; each call is (name, id) or (name, id, args).""" + from langchain_core.messages import AIMessage + + return AIMessage( + content=content, + tool_calls=[ + { + "name": c[0], + "args": c[2] if len(c) > 2 else {}, + "id": c[1], + "type": "tool_call", + } + for c in calls + ], + ) + + +async def _run_handoff( + graph_def: GraphDefinition, + factory: Any, + tool_handlers: dict[str, Any] | None = None, +) -> tuple[dict[str, Any], list[str]]: + from langchain_core.messages import AIMessage # noqa: F401 + + tracked: list[str] = [] + client = MagicMock() + client.track = MagicMock(side_effect=lambda evt, *_a: tracked.append(evt)) + with ( + patch("launchdarkly_ai_langchain_agents.native_graph._HAS_OTEL", False), + patch( + "launchdarkly_ai_langchain_agents.native_graph.get_client", + return_value=client, + ), + ): + result = await to_lang_graph( + _make_def_promise(graph_def), + { + "model_factory": factory, + "tool_handlers": tool_handlers or {}, + "context": {"kind": "user", "key": "u"}, + }, + ).invoke("who did it?") + return result, tracked + + +def _sent(model: Any, call: int = 0) -> list[tuple[str, Any]]: + """(type, content) of each message a node's model received on one call.""" + return [(m.type, m.content) for m in model.ainvoke.call_args_list[call].args[0]] + + +def _assert_neutral(model: Any) -> None: + """No tool role, tool calls or non-text blocks reach a node that was handed to.""" + for m in model.ainvoke.call_args.args[0]: + assert m.type != "tool" + assert not getattr(m, "tool_calls", None) + assert isinstance(m.content, str) + + +# Same literal as the TypeScript tests' REASONING, so the two rewrite tests share one input. +_REASONING = {"type": "reasoning_content", "reasoning_content": {"text": "hmm"}} + + +class TestNativeHandoffRealStateGraph: + @pytest.mark.asyncio + async def test_multi_child_router_hands_off_once(self) -> None: + from langchain_core.messages import AIMessage + + factory, models = _scripted_models( + { + "router": [ + _tool_call_msg( + ("transfer_to_manor_butler", "call_1"), + content=[_REASONING, {"type": "text", "text": "Routing."}], + ) + ], + "manor.butler": [AIMessage(content="the butler did it")], + "cook": [], + } + ) + result, tracked = await _run_handoff( + _handoff_graph_def(["manor.butler", "cook"]), factory + ) + + assert result["response"] == "the butler did it" + assert models["router"].ainvoke.call_count == 1 # not re-run after the hand-off + assert models["cook"].ainvoke.call_count == 0 + # A target with no tools gets the router's turn as plain text: reasoning + # dropped, the call and its result as lines. The result text uses the raw + # target key; the tool name uses the sanitized one. + assert _sent(models["manor.butler"]) == [ + ("system", "you are manor.butler"), + ("human", "who did it?"), + ("ai", "Routing.\n[Tool call] transfer_to_manor_butler {}"), + ( + "human", + "[Tool result] transfer_to_manor_butler: Handed off to manor.butler.", + ), + ] + _assert_neutral(models["manor.butler"]) + assert tracked.count("$ld:ai:graph:handoff_success") == 1 + + @pytest.mark.asyncio + async def test_bind_tools_gets_no_parallel_tool_calls(self) -> None: + """Bedrock Converse rejects parallel_tool_calls, so it is never passed.""" + from langchain_core.messages import AIMessage + + factory, models = _scripted_models( + { + "router": [_tool_call_msg(("transfer_to_butler", "call_1"))], + "butler": [AIMessage(content="done")], + "cook": [], + } + ) + await _run_handoff(_handoff_graph_def(["butler", "cook"]), factory) + + assert models["router"].bind_tools.call_args.kwargs == {} + + @pytest.mark.asyncio + async def test_normal_tool_and_handoff_in_one_message_both_run(self) -> None: + from langchain_core.messages import AIMessage + + lookups: list[Any] = [] + factory, models = _scripted_models( + { + "router": [ + _tool_call_msg( + ("lookup", "call_a", {"room": "library"}), + ("transfer_to_butler", "call_b"), + ) + ], + "butler": [AIMessage(content="done")], + "cook": [], + } + ) + result, _ = await _run_handoff( + _handoff_graph_def( + ["butler", "cook"], + root_tools={ + "lookup": { + "description": "look up", + "parameters": { + "type": "object", + "properties": {"room": {"type": "string"}}, + }, + } + }, + ), + factory, + tool_handlers={"lookup": lambda args: lookups.append(args) or "clue"}, + ) + + assert result["response"] == "done" + assert lookups == [{"room": "library"}] + assert models["router"].ainvoke.call_count == 1 + # Both results become one user message, in the call order. + assert _sent(models["butler"])[2:] == [ + ( + "ai", + '[Tool call] lookup {"room":"library"}\n[Tool call] transfer_to_butler {}', + ), + ( + "human", + "[Tool result] lookup: clue\n" + "[Tool result] transfer_to_butler: Handed off to butler.", + ), + ] + _assert_neutral(models["butler"]) + + @pytest.mark.asyncio + async def test_first_handoff_wins(self) -> None: + from langchain_core.messages import AIMessage + + factory, models = _scripted_models( + { + "router": [ + _tool_call_msg( + ("transfer_to_butler", "call_1"), ("transfer_to_cook", "call_2") + ) + ], + "butler": [AIMessage(content="butler answer")], + "cook": [], + } + ) + result, tracked = await _run_handoff( + _handoff_graph_def(["butler", "cook"]), factory + ) + + assert result["response"] == "butler answer" + assert models["cook"].ainvoke.call_count == 0 + assert _sent(models["butler"])[-1] == ( + "human", + "[Tool result] transfer_to_butler: Handed off to butler.\n" + "[Tool result] transfer_to_cook: Ignored: only the first transfer in a turn routes.", + ) + assert tracked.count("$ld:ai:graph:handoff_success") == 1 + + @pytest.mark.asyncio + async def test_single_child_hands_off_once(self) -> None: + from langchain_core.messages import AIMessage + + factory, models = _scripted_models( + { + "router": [_tool_call_msg(("transfer_to_butler", "call_1"))], + "butler": [AIMessage(content="the butler did it")], + } + ) + result, _ = await _run_handoff(_handoff_graph_def(["butler"]), factory) + + assert result["response"] == "the butler did it" + assert models["router"].ainvoke.call_count == 1 + assert _sent(models["butler"])[2:] == [ + ("ai", "[Tool call] transfer_to_butler {}"), + ("human", "[Tool result] transfer_to_butler: Handed off to butler."), + ] + _assert_neutral(models["butler"]) + + @pytest.mark.asyncio + async def test_own_tool_loop_stays_native(self) -> None: + """Only the history from before a node was entered is rewritten: a node's + own tool call and its result reach its next model call unchanged.""" + from langchain_core.messages import AIMessage + + factory, models = _scripted_models( + { + "router": [ + _tool_call_msg(("lookup", "call_x"), content=[_REASONING]), + AIMessage(content="answer"), + ], + } + ) + result, _ = await _run_handoff( + _handoff_graph_def( + [], + root_tools={ + "lookup": { + "description": "look up", + "parameters": {"type": "object", "properties": {}}, + } + }, + ), + factory, + tool_handlers={"lookup": lambda args: "clue"}, + ) + + assert result["response"] == "answer" + second = models["router"].ainvoke.call_args_list[1].args[0] + assert [m.type for m in second] == ["system", "human", "ai", "tool"] + assert second[2].tool_calls[0]["id"] == "call_x" + assert second[2].content == [_REASONING] + assert second[3].tool_call_id == "call_x" + + +class TestToNeutralHistory: + """The rewrite rules in spec §2.x.3, on one input that hits each of them. + The TypeScript test uses the same input and expects the same output.""" + + def test_rewrites_every_case(self) -> None: + from langchain_core.messages import ( + AIMessage, + AIMessageChunk, + HumanMessage, + SystemMessage, + ToolMessage, + ToolMessageChunk, + ) + + image = [ + {"type": "text", "text": "look"}, + {"type": "image_url", "image_url": {"url": "data:x"}}, + ] + history = [ + SystemMessage("sys"), + HumanMessage(image), + AIMessage( + content=[ + _REASONING, + {"type": "text", "text": "a"}, + {"type": "text", "text": None}, + {"type": "text", "text": "b"}, + ], + tool_calls=[ + { + "name": "first", + "args": {"q": "café", "n": {"x": [1, 2]}}, + "id": "c1", + }, + {"name": "second", "args": {}, "id": "c2"}, + ], + ), + ToolMessage("orphan", tool_call_id="zz", name="lost"), + ToolMessage([{"type": "text", "text": "two"}], tool_call_id="c2"), + ToolMessage("one", tool_call_id="c1"), + AIMessage(content=[_REASONING]), + AIMessageChunk( + content="chunk", tool_calls=[{"name": "t", "args": {}, "id": "c3"}] + ), + ToolMessageChunk(content="three", tool_call_id="c3"), + ] + + out = _to_neutral_history(history, AIMessage, HumanMessage) + + assert [(m.type, m.content) for m in out] == [ + ("system", "sys"), + ("human", image), + ( + "ai", + "ab\n" + '[Tool call] first {"q":"café","n":{"x":[1,2]}}\n' + "[Tool call] second {}", + ), + ( + "human", + "[Tool result] first: one\n" + "[Tool result] second: two\n" + "[Tool result] lost: orphan", + ), + ("ai", "chunk\n[Tool call] t {}"), + ("human", "[Tool result] t: three"), + ] + assert all(not getattr(m, "tool_calls", None) for m in out) diff --git a/uv.lock b/uv.lock index f3053e7f..87ca4867 100644 --- a/uv.lock +++ b/uv.lock @@ -1367,6 +1367,7 @@ dependencies = [ { name = "langgraph-prebuilt" }, { name = "launchdarkly-ai-server" }, { name = "opentelemetry-api" }, + { name = "pydantic" }, ] [package.optional-dependencies] @@ -1387,6 +1388,7 @@ requires-dist = [ { name = "langgraph-prebuilt", specifier = ">=1.1" }, { name = "launchdarkly-ai-server", editable = "packages/client" }, { name = "opentelemetry-api", specifier = ">=1.25" }, + { name = "pydantic", specifier = ">=2.13" }, ] provides-extras = ["anthropic", "aws"]