121 lines
4.3 KiB
Python
121 lines
4.3 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""tool_result_messages: one helper decides how a tool result — with or
|
|
without images — enters the conversation, per LLM provider."""
|
|
|
|
from enum import StrEnum
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from langchain_core.messages import HumanMessage, ToolMessage
|
|
|
|
from artemis.tools.tool_wrapper import (
|
|
resolve_image_carrier,
|
|
split_multimodal_result,
|
|
tool_result_messages,
|
|
)
|
|
|
|
IMAGE = {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,AAAA"}}
|
|
SHOT = [{"type": "text", "text": "Screenshot of step 7 (pre-action) is attached."}, IMAGE]
|
|
|
|
|
|
class _Provider(StrEnum):
|
|
GOOGLE = "google"
|
|
VERTEX_AI = "vertexai"
|
|
|
|
|
|
def _llm(provider):
|
|
return SimpleNamespace(endpoint=SimpleNamespace(provider=provider))
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("provider", "carrier"),
|
|
[
|
|
("google", "tool"),
|
|
("anthropic", "tool"),
|
|
("vertexai", "human"),
|
|
("openai", "human"),
|
|
("openrouter", "human"),
|
|
("xai", "human"),
|
|
("ollama", "human"),
|
|
(None, "human"),
|
|
],
|
|
)
|
|
def test_resolve_image_carrier_by_provider_name(provider, carrier):
|
|
assert resolve_image_carrier(provider) == carrier
|
|
assert resolve_image_carrier(_llm(provider)) == carrier
|
|
|
|
|
|
def test_resolve_image_carrier_accepts_enums_and_endpoint_less_llms():
|
|
assert resolve_image_carrier(_Provider.GOOGLE) == "tool"
|
|
assert resolve_image_carrier(_llm(_Provider.VERTEX_AI)) == "human"
|
|
assert resolve_image_carrier(SimpleNamespace()) == "human"
|
|
|
|
|
|
def test_split_multimodal_result():
|
|
assert split_multimodal_result("plain") == ("plain", [])
|
|
assert split_multimodal_result(None) == ("", [])
|
|
assert split_multimodal_result(ToolMessage(content="x", tool_call_id="t")) == ("x", [])
|
|
text, images = split_multimodal_result(SHOT)
|
|
assert text == "Screenshot of step 7 (pre-action) is attached."
|
|
assert images == [IMAGE]
|
|
|
|
|
|
def test_text_only_results_stay_a_single_tool_message():
|
|
msgs = tool_result_messages("tc-1", "Step 3 replay", name="replay_steps", status="error")
|
|
assert len(msgs) == 1
|
|
assert isinstance(msgs[0], ToolMessage)
|
|
assert msgs[0].content == "Step 3 replay"
|
|
assert msgs[0].tool_call_id == "tc-1"
|
|
assert msgs[0].name == "replay_steps"
|
|
assert msgs[0].status == "error"
|
|
|
|
# Text-block lists pass through unchanged; other objects are stringified.
|
|
blocks = [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}]
|
|
assert tool_result_messages("tc", blocks)[0].content == blocks
|
|
assert tool_result_messages("tc", {"k": 1})[0].content == "{'k': 1}"
|
|
|
|
|
|
def test_tool_carrier_puts_the_image_inside_the_tool_message():
|
|
msgs = tool_result_messages("tc-2", SHOT, name="get_step_screenshot", llm=_llm("google"))
|
|
assert len(msgs) == 1
|
|
assert isinstance(msgs[0], ToolMessage)
|
|
assert msgs[0].content == SHOT
|
|
assert msgs[0].name == "get_step_screenshot"
|
|
|
|
|
|
def test_human_carrier_splits_text_and_image_with_a_caption():
|
|
msgs = tool_result_messages(
|
|
"tc-3", SHOT, name="get_step_screenshot", status="success", llm=_llm("openai")
|
|
)
|
|
assert len(msgs) == 2
|
|
tool_msg, human_msg = msgs
|
|
assert isinstance(tool_msg, ToolMessage)
|
|
assert tool_msg.content == "Screenshot of step 7 (pre-action) is attached."
|
|
assert tool_msg.tool_call_id == "tc-3"
|
|
assert isinstance(human_msg, HumanMessage)
|
|
assert human_msg.content[0] == {
|
|
"type": "text",
|
|
"text": "[Screenshot returned by get_step_screenshot for step 7]",
|
|
}
|
|
assert human_msg.content[1] == IMAGE
|
|
|
|
|
|
def test_explicit_carrier_wins_over_the_llm():
|
|
msgs = tool_result_messages("tc", SHOT, image_carrier="human", llm=_llm("google"))
|
|
assert len(msgs) == 2
|
|
msgs = tool_result_messages("tc", SHOT, image_carrier="tool", llm=_llm("openai"))
|
|
assert len(msgs) == 1
|