diff --git a/.github/workflows/ci_check.yml b/.github/workflows/ci_check.yml index 0c652e7fa..7f33e9ca0 100644 --- a/.github/workflows/ci_check.yml +++ b/.github/workflows/ci_check.yml @@ -202,7 +202,7 @@ jobs: uv tool install pre-commit==${{ steps.ci-config.outputs.pre_commit_version }} FLOW_CI_UV_SYNC_EXTRA_ARGS=() if [[ "$PYTHON_INTEGRATION_LANGCHAIN" == "true" ]]; then - FLOW_CI_UV_SYNC_EXTRA_ARGS+=(--extra langchain --extra langgraph --extra deepagents) + FLOW_CI_UV_SYNC_EXTRA_ARGS+=(--extra langchain --extra langchain-nvidia --extra langgraph --extra deepagents) fi uv sync --inexact --no-install-project --no-install-package nemo-relay "${FLOW_CI_UV_SYNC_EXTRA_ARGS[@]}" diff --git a/justfile b/justfile index 1e731ae21..1d5c1c948 100644 --- a/justfile +++ b/justfile @@ -1442,7 +1442,7 @@ test-python-langchain: {{ bash_helpers }} pytest_cmd=(pytest) cd "$NEMO_RELAY_REPO_ROOT" - uv sync --inexact --no-install-project --no-install-package nemo-relay --extra langchain --extra langgraph --extra deepagents + uv sync --inexact --no-install-project --no-install-package nemo-relay --extra langchain --extra langchain-nvidia --extra langgraph --extra deepagents activate_project_venv export_uv_python_runtime python_executable="$(project_python_executable)" diff --git a/python/nemo_relay/integrations/langchain/_serialization.py b/python/nemo_relay/integrations/langchain/_serialization.py index 4c08ed8b5..82219af5a 100644 --- a/python/nemo_relay/integrations/langchain/_serialization.py +++ b/python/nemo_relay/integrations/langchain/_serialization.py @@ -112,6 +112,25 @@ def _annotated_tool_calls_to_langchain(cls, tool_calls: Any) -> list[dict[str, A return langchain_tool_calls or None + @classmethod + def _annotated_tool_calls_to_provider(cls, tool_calls: Any) -> list[dict[str, Any]] | None: + """Return the OpenAI-compatible representation required by some providers.""" + langchain_tool_calls = cls._annotated_tool_calls_to_langchain(tool_calls) + if langchain_tool_calls is None: + return None + + return [ + { + "id": tool_call["id"], + "type": "function", + "function": { + "name": tool_call["name"], + "arguments": json.dumps(tool_call["args"], separators=(",", ":")), + }, + } + for tool_call in langchain_tool_calls + ] + @classmethod def _langchain_message_to_annotated(cls, message: BaseMessage) -> list[dict[str, Any]]: content = message.content @@ -124,7 +143,7 @@ def _langchain_message_to_annotated(cls, message: BaseMessage) -> list[dict[str, role = _LC_TO_RELAY_MESSAGE_ROLE.get(message.type, message.type) messages = [] - for msg in content: + for content_index, msg in enumerate(content): relay_message: dict[str, Any] = {"role": role} if isinstance(msg, str): relay_message["content"] = msg @@ -140,7 +159,7 @@ def _langchain_message_to_annotated(cls, message: BaseMessage) -> list[dict[str, # Using getattr as we are inferring subclasses of BaseMessage based upon the role if role == "assistant": - tool_calls = getattr(message, "tool_calls", []) + tool_calls = getattr(message, "tool_calls", []) if content_index == 0 else [] relay_message["tool_calls"] = cls._langchain_tool_calls_to_annotated(tool_calls) elif role == "tool": relay_message["tool_call_id"] = getattr(message, "tool_call_id", "") @@ -150,7 +169,11 @@ def _langchain_message_to_annotated(cls, message: BaseMessage) -> list[dict[str, return messages @classmethod - def _annotated_message_to_langchain(cls, message: dict[str, Any]) -> BaseMessage: + def _annotated_message_to_langchain( + cls, + message: dict[str, Any], + provider_tool_calls: list[dict[str, Any]] | None = None, + ) -> BaseMessage: role = message.get("role") content = message.get("content", "") name = message.get("name") @@ -161,11 +184,34 @@ def _annotated_message_to_langchain(cls, message: dict[str, Any]) -> BaseMessage return HumanMessage(content=content, name=name) if role == "assistant": tool_calls = cls._annotated_tool_calls_to_langchain(message.get("tool_calls")) - return AIMessage(content=content, name=name, tool_calls=tool_calls or []) + additional_kwargs = {"tool_calls": provider_tool_calls} if provider_tool_calls is not None else {} + return AIMessage( + content=content, name=name, tool_calls=tool_calls or [], additional_kwargs=additional_kwargs + ) if role == "tool": return ToolMessage(content=content, name=name, tool_call_id=str(message.get("tool_call_id") or "")) raise ValueError(f"Unsupported annotated LangChain message role: {role!r}") + @classmethod + def _original_provider_tool_calls(cls, original: LLMRequest) -> list[tuple[dict[str, Any], list[dict[str, Any]]]]: + """Return assistant messages associated with their provider tool calls.""" + raw_messages = original.content.get("messages") + if not isinstance(raw_messages, list): + return [] + + provider_tool_calls: list[tuple[dict[str, Any], list[dict[str, Any]]]] = [] + for message in messages_from_dict(raw_messages): + raw_tool_calls = None + if isinstance(message, AIMessage): + candidate = message.additional_kwargs.get("tool_calls") + if isinstance(candidate, list): + raw_tool_calls = candidate + if raw_tool_calls is not None: + annotated_messages = cls._langchain_message_to_annotated(message) + if annotated_messages: + provider_tool_calls.append((annotated_messages[0], raw_tool_calls)) + return provider_tool_calls + def decode(self, request: LLMRequest) -> AnnotatedLLMRequest: """Decode a LangChain-shaped request payload into an annotated request.""" payload = request.content @@ -192,9 +238,37 @@ def encode(self, annotated: AnnotatedLLMRequest, original: LLMRequest) -> LLMReq """Encode annotated request edits back into a LangChain-shaped payload.""" payload = dict(original.content) payload.update(annotated.extra) - payload["messages"] = messages_to_dict( - [self._annotated_message_to_langchain(message) for message in annotated.messages] - ) + original_provider_tool_calls = self._original_provider_tool_calls(original) + matched_original_messages: set[int] = set() + encoded_messages = [] + for message in annotated.messages: + provider_tool_calls = None + if message.get("role") == "assistant": + for index, (original_message, original_tool_calls) in enumerate(original_provider_tool_calls): + if index in matched_original_messages or message != original_message: + continue + matched_original_messages.add(index) + provider_tool_calls = original_tool_calls + break + + if message.get("role") == "assistant" and provider_tool_calls is None: + message_without_tool_calls = {key: value for key, value in message.items() if key != "tool_calls"} + for index, (original_message, original_tool_calls) in enumerate(original_provider_tool_calls): + original_without_tool_calls = { + key: value for key, value in original_message.items() if key != "tool_calls" + } + if index in matched_original_messages or message_without_tool_calls != original_without_tool_calls: + continue + matched_original_messages.add(index) + if message.get("tool_calls") == original_message.get("tool_calls"): + provider_tool_calls = original_tool_calls + else: + provider_tool_calls = self._annotated_tool_calls_to_provider(message.get("tool_calls")) + break + if message.get("role") == "assistant" and provider_tool_calls is None: + provider_tool_calls = self._annotated_tool_calls_to_provider(message.get("tool_calls")) + encoded_messages.append(self._annotated_message_to_langchain(message, provider_tool_calls)) + payload["messages"] = messages_to_dict(encoded_messages) if annotated.model is not None: payload["model"] = annotated.model if annotated.tools is not None: diff --git a/python/tests/integrations/langchain_tests/test_middleware.py b/python/tests/integrations/langchain_tests/test_middleware.py index 7789b7411..2e9a071a8 100644 --- a/python/tests/integrations/langchain_tests/test_middleware.py +++ b/python/tests/integrations/langchain_tests/test_middleware.py @@ -244,6 +244,386 @@ def test_langchain_model_request_codec_round_trips_messages(model_request: Model assert round_tripped.messages[0].content == "hello from intercept" +def test_langchain_request_codec_preserves_provider_tool_calls(): + from langchain_core.messages import AIMessage, messages_from_dict, messages_to_dict + + from nemo_relay.integrations.langchain._serialization import LangChainCodec + + provider_tool_calls = [ + { + "id": "call-weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"SF"}'}, + } + ] + request = nemo_relay.LLMRequest( + {}, + { + "messages": messages_to_dict( + [ + AIMessage( + content="", + tool_calls=[ + { + "id": "call-weather", + "name": "get_weather", + "args": {"city": "SF"}, + "type": "tool_call", + } + ], + additional_kwargs={"tool_calls": provider_tool_calls}, + ) + ] + ) + }, + ) + + codec = LangChainCodec() + encoded = codec.encode(codec.decode(request), request) + rebuilt = messages_from_dict(cast(list[dict[str, Any]], encoded.content["messages"]))[0] + assert isinstance(rebuilt, AIMessage) + + assert rebuilt.tool_calls == [ + { + "id": "call-weather", + "name": "get_weather", + "args": {"city": "SF"}, + "type": "tool_call", + } + ] + assert rebuilt.additional_kwargs["tool_calls"] == provider_tool_calls + + +def test_langchain_request_codec_preserves_chat_nvidia_tool_call_payload_after_prepending_message(): + from langchain_core.messages import AIMessage, messages_from_dict, messages_to_dict + + from nemo_relay.integrations.langchain._serialization import LangChainCodec + + convert_message_to_dict = pytest.importorskip("langchain_nvidia_ai_endpoints._utils").convert_message_to_dict + provider_tool_calls = [ + { + "id": "call-weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"SF"}'}, + } + ] + request = nemo_relay.LLMRequest( + {}, + { + "messages": messages_to_dict( + [ + AIMessage( + content="", + tool_calls=[ + { + "id": "call-weather", + "name": "get_weather", + "args": {"city": "SF"}, + "type": "tool_call", + } + ], + additional_kwargs={"tool_calls": provider_tool_calls}, + ) + ] + ) + }, + ) + + codec = LangChainCodec() + annotated = codec.decode(request) + annotated.messages = [{"role": "user", "content": "Prepended by an interceptor"}, *annotated.messages] + encoded = codec.encode(annotated, request) + rebuilt = messages_from_dict(cast(list[dict[str, Any]], encoded.content["messages"]))[1] + + assert convert_message_to_dict(rebuilt) == { + "role": "assistant", + "content": None, + "tool_calls": provider_tool_calls, + } + + +def test_langchain_request_codec_preserves_reordered_provider_tool_calls(): + from langchain_core.messages import AIMessage, messages_from_dict, messages_to_dict + + from nemo_relay.integrations.langchain._serialization import LangChainCodec + + first_provider_tool_calls = [ + { + "id": "call-one", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}, + "provider_field": "first", + } + ] + second_provider_tool_calls = [ + { + "id": "call-two", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "NY"}'}, + "provider_field": "second", + } + ] + request = nemo_relay.LLMRequest( + {}, + { + "messages": messages_to_dict( + [ + AIMessage( + content="", + tool_calls=[ + { + "id": "call-one", + "name": "get_weather", + "args": {"city": "SF"}, + "type": "tool_call", + } + ], + additional_kwargs={"tool_calls": first_provider_tool_calls}, + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "call-two", + "name": "get_weather", + "args": {"city": "NY"}, + "type": "tool_call", + } + ], + additional_kwargs={"tool_calls": second_provider_tool_calls}, + ), + ] + ) + }, + ) + + codec = LangChainCodec() + annotated = codec.decode(request) + annotated.messages = list(reversed(annotated.messages)) + rebuilt = messages_from_dict(cast(list[dict[str, Any]], codec.encode(annotated, request).content["messages"])) + + assert [message.additional_kwargs["tool_calls"] for message in rebuilt] == [ + second_provider_tool_calls, + first_provider_tool_calls, + ] + + +def test_langchain_request_codec_rebuilds_provider_tool_calls_after_content_edit(): + from langchain_core.messages import AIMessage, messages_from_dict, messages_to_dict + + from nemo_relay.integrations.langchain._serialization import LangChainCodec + + request = nemo_relay.LLMRequest( + {}, + { + "messages": messages_to_dict( + [ + AIMessage( + content="sensitive content", + tool_calls=[ + { + "id": "call-weather", + "name": "get_weather", + "args": {"city": "SF"}, + "type": "tool_call", + } + ], + additional_kwargs={ + "tool_calls": [ + { + "id": "call-weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}, + } + ] + }, + ) + ] + ) + }, + ) + + codec = LangChainCodec() + annotated = codec.decode(request) + annotated.messages = [{**annotated.messages[0], "content": "[redacted]"}] + rebuilt = messages_from_dict(cast(list[dict[str, Any]], codec.encode(annotated, request).content["messages"]))[0] + + assert isinstance(rebuilt, AIMessage) + assert rebuilt.additional_kwargs["tool_calls"] == [ + { + "id": "call-weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"SF"}'}, + } + ] + + +def test_langchain_request_codec_builds_provider_tool_calls_for_new_assistant(): + from langchain_core.messages import AIMessage, HumanMessage, messages_from_dict, messages_to_dict + + from nemo_relay.integrations.langchain._serialization import LangChainCodec + + request = nemo_relay.LLMRequest({}, {"messages": messages_to_dict([HumanMessage(content="hello")])}) + + codec = LangChainCodec() + annotated = codec.decode(request) + annotated.messages = [ + *annotated.messages, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call-weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"SF"}'}, + } + ], + }, + ] + rebuilt = messages_from_dict(cast(list[dict[str, Any]], codec.encode(annotated, request).content["messages"]))[1] + + assert isinstance(rebuilt, AIMessage) + assert rebuilt.additional_kwargs["tool_calls"] == [ + { + "id": "call-weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"SF"}'}, + } + ] + + +def test_langchain_request_codec_keeps_multi_block_tool_calls_on_first_assistant_fragment(): + from langchain_core.messages import AIMessage, messages_from_dict, messages_to_dict + + from nemo_relay.integrations.langchain._serialization import LangChainCodec + + request = nemo_relay.LLMRequest( + {}, + { + "messages": messages_to_dict( + [ + AIMessage( + content=[{"type": "text", "text": "first"}, {"type": "text", "text": "second"}], + tool_calls=[ + { + "id": "call-weather", + "name": "get_weather", + "args": {"city": "SF"}, + "type": "tool_call", + } + ], + additional_kwargs={ + "tool_calls": [ + { + "id": "call-weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}, + } + ] + }, + ) + ] + ) + }, + ) + + codec = LangChainCodec() + rebuilt = messages_from_dict( + cast(list[dict[str, Any]], codec.encode(codec.decode(request), request).content["messages"]) + ) + assert all(isinstance(message, AIMessage) for message in rebuilt) + rebuilt_assistants = cast(list[AIMessage], rebuilt) + + assert [message.tool_calls for message in rebuilt_assistants] == [ + [{"id": "call-weather", "name": "get_weather", "args": {"city": "SF"}, "type": "tool_call"}], + [], + ] + assert [message.additional_kwargs for message in rebuilt_assistants] == [ + { + "tool_calls": [ + { + "id": "call-weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"SF"}'}, + } + ] + }, + {}, + ] + + +def test_model_call_intercept_rebuilds_provider_tool_calls( + nemo_relay_middleware: NemoRelayMiddleware, + model_request: ModelRequest[Any], + model_request_handler: tuple[Callable[[ModelRequest[Any]], ModelResponse[Any]], dict[str, ModelRequest[Any]]], +): + from langchain_core.messages import AIMessage + + provider_tool_calls = [ + { + "id": "call-weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"SF"}'}, + } + ] + original = model_request.override( + messages=[ + AIMessage( + content="", + tool_calls=[ + { + "id": "call-weather", + "name": "get_weather", + "args": {"city": "SF"}, + "type": "tool_call", + } + ], + additional_kwargs={"tool_calls": provider_tool_calls}, + ) + ] + ) + + def change_tool_call(_: str, request: nemo_relay.LLMRequest, annotated: Any): + assert annotated is not None + annotated.messages = [ + { + **message, + "tool_calls": [ + { + **message["tool_calls"][0], + "function": { + **message["tool_calls"][0]["function"], + "arguments": '{"city":"San Jose"}', + }, + } + ], + } + if message.get("role") == "assistant" + else message + for message in annotated.messages + ] + return nemo_relay.LLMRequestInterceptOutcome(request, annotated) + + nemo_relay.intercepts.register_llm_request("test_langchain_change_tool_call", 1, False, change_tool_call) + try: + (handler, seen_request) = model_request_handler + nemo_relay_middleware.wrap_model_call(original, handler) + finally: + nemo_relay.intercepts.deregister_llm_request("test_langchain_change_tool_call") + + rebuilt = next(message for message in seen_request["request"].messages if message.type == "ai") + assert isinstance(rebuilt, AIMessage) + assert rebuilt.tool_calls[0]["args"] == {"city": "San Jose"} + assert rebuilt.additional_kwargs["tool_calls"] == [ + { + "id": "call-weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"San Jose"}'}, + } + ] + + def test_payload_to_model_request_moves_relay_headers_to_chat_nvidia_transport(model_request: ModelRequest[Any]): from nemo_relay.integrations.langchain._serialization import ( model_request_to_payload,