wehub-resource-sync
614 lines
22 KiB
Python
614 lines
22 KiB
Python
import json
|
|
|
|
import pytest
|
|
from pytest_httpx import IteratorStream
|
|
|
|
import llm
|
|
from llm.default_plugins.openai_models import Chat
|
|
from llm.models import Prompt
|
|
|
|
API_KEY = "badkey"
|
|
|
|
|
|
def _sse(delta, finish_reason=None, usage=None, tool_calls=None):
|
|
chunk = {
|
|
"id": "c1",
|
|
"object": "chat.completion.chunk",
|
|
"created": 1700000000,
|
|
"model": "gpt-4o-mini",
|
|
"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
|
|
}
|
|
if tool_calls is not None:
|
|
chunk["choices"][0]["delta"]["tool_calls"] = tool_calls
|
|
if usage is not None:
|
|
chunk["usage"] = usage
|
|
return f"data: {json.dumps(chunk)}\n\n".encode("utf-8")
|
|
|
|
|
|
def _text_stream():
|
|
yield _sse({"role": "assistant", "content": ""})
|
|
yield _sse({"content": "Hel"})
|
|
yield _sse({"content": "lo"})
|
|
yield _sse({}, finish_reason="stop")
|
|
yield b"data: [DONE]\n\n"
|
|
|
|
|
|
def _tool_call_stream():
|
|
"""Mimic an OpenAI stream with a tool call (no preceding text)."""
|
|
yield _sse({"role": "assistant", "content": None})
|
|
yield _sse(
|
|
{},
|
|
tool_calls=[
|
|
{
|
|
"index": 0,
|
|
"id": "call_1",
|
|
"type": "function",
|
|
"function": {"name": "get_weather", "arguments": ""},
|
|
}
|
|
],
|
|
)
|
|
yield _sse(
|
|
{},
|
|
tool_calls=[
|
|
{
|
|
"index": 0,
|
|
"function": {"arguments": '{"city":'},
|
|
}
|
|
],
|
|
)
|
|
yield _sse(
|
|
{},
|
|
tool_calls=[
|
|
{
|
|
"index": 0,
|
|
"function": {"arguments": '"Paris"}'},
|
|
}
|
|
],
|
|
)
|
|
yield _sse({}, finish_reason="tool_calls")
|
|
yield b"data: [DONE]\n\n"
|
|
|
|
|
|
def _text_then_tool_call_stream():
|
|
"""Text arrives first, then a tool call — the tool call must get
|
|
a part_index past the text so assembly doesn't mix families."""
|
|
yield _sse({"role": "assistant", "content": ""})
|
|
yield _sse({"content": "Looking up"})
|
|
yield _sse(
|
|
{},
|
|
tool_calls=[
|
|
{
|
|
"index": 0,
|
|
"id": "call_1",
|
|
"type": "function",
|
|
"function": {"name": "get_weather", "arguments": '{"c":1}'},
|
|
}
|
|
],
|
|
)
|
|
yield _sse({}, finish_reason="tool_calls")
|
|
yield b"data: [DONE]\n\n"
|
|
|
|
|
|
@pytest.fixture
|
|
def chat_model():
|
|
# A plain Chat instance with vision and tools enabled — enough
|
|
# capabilities for the Part subtypes we translate.
|
|
return Chat("gpt-4o-mini", vision=True, supports_tools=True)
|
|
|
|
|
|
class TestBuildMessagesFromExplicitMessages:
|
|
def test_single_user_message(self, chat_model):
|
|
prompt = Prompt(None, model=chat_model, messages=[llm.user("hi")])
|
|
result = chat_model.build_messages(prompt, None)
|
|
assert result == [{"role": "user", "content": "hi"}]
|
|
|
|
def test_system_plus_user(self, chat_model):
|
|
prompt = Prompt(
|
|
None,
|
|
model=chat_model,
|
|
messages=[llm.system("be brief"), llm.user("hi")],
|
|
)
|
|
result = chat_model.build_messages(prompt, None)
|
|
assert result == [
|
|
{"role": "system", "content": "be brief"},
|
|
{"role": "user", "content": "hi"},
|
|
]
|
|
|
|
def test_user_with_attachment(self, chat_model):
|
|
att = llm.Attachment(type="image/jpeg", url="http://example.com/cat.jpg")
|
|
prompt = Prompt(
|
|
None,
|
|
model=chat_model,
|
|
messages=[llm.user("describe", att)],
|
|
)
|
|
result = chat_model.build_messages(prompt, None)
|
|
assert result == [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "describe"},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": "http://example.com/cat.jpg"},
|
|
},
|
|
],
|
|
}
|
|
]
|
|
|
|
def test_assistant_with_tool_call(self, chat_model):
|
|
tool_call = llm.parts.ToolCallPart(
|
|
name="search",
|
|
arguments={"q": "weather"},
|
|
tool_call_id="c1",
|
|
)
|
|
prompt = Prompt(
|
|
None,
|
|
model=chat_model,
|
|
messages=[
|
|
llm.user("search weather"),
|
|
llm.assistant("on it", tool_call),
|
|
],
|
|
)
|
|
result = chat_model.build_messages(prompt, None)
|
|
assert result == [
|
|
{"role": "user", "content": "search weather"},
|
|
{
|
|
"role": "assistant",
|
|
"content": "on it",
|
|
"tool_calls": [
|
|
{
|
|
"type": "function",
|
|
"id": "c1",
|
|
"function": {
|
|
"name": "search",
|
|
"arguments": json.dumps({"q": "weather"}),
|
|
},
|
|
}
|
|
],
|
|
},
|
|
]
|
|
|
|
def test_assistant_tool_call_only_no_text(self, chat_model):
|
|
"""When an assistant message has tool_calls but no text, OpenAI
|
|
expects content=null."""
|
|
tool_call = llm.parts.ToolCallPart(
|
|
name="search", arguments={"q": "x"}, tool_call_id="c1"
|
|
)
|
|
prompt = Prompt(
|
|
None,
|
|
model=chat_model,
|
|
messages=[llm.user("q"), llm.assistant(tool_call)],
|
|
)
|
|
result = chat_model.build_messages(prompt, None)
|
|
assert result[1] == {
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"type": "function",
|
|
"id": "c1",
|
|
"function": {
|
|
"name": "search",
|
|
"arguments": json.dumps({"q": "x"}),
|
|
},
|
|
}
|
|
],
|
|
}
|
|
|
|
def test_tool_role_message_with_tool_result(self, chat_model):
|
|
tr = llm.parts.ToolResultPart(name="search", output="sunny", tool_call_id="c1")
|
|
prompt = Prompt(
|
|
None,
|
|
model=chat_model,
|
|
messages=[
|
|
llm.user("q"),
|
|
llm.tool_message(tr),
|
|
],
|
|
)
|
|
result = chat_model.build_messages(prompt, None)
|
|
assert result == [
|
|
{"role": "user", "content": "q"},
|
|
{"role": "tool", "tool_call_id": "c1", "content": "sunny"},
|
|
]
|
|
|
|
def test_multiple_tool_results_emit_multiple_messages(self, chat_model):
|
|
"""Parallel tool results: one OpenAI 'tool' message per result."""
|
|
a = llm.parts.ToolResultPart(name="t", output="A", tool_call_id="c1")
|
|
b = llm.parts.ToolResultPart(name="t", output="B", tool_call_id="c2")
|
|
prompt = Prompt(
|
|
None,
|
|
model=chat_model,
|
|
messages=[llm.user("q"), llm.tool_message(a, b)],
|
|
)
|
|
result = chat_model.build_messages(prompt, None)
|
|
assert result == [
|
|
{"role": "user", "content": "q"},
|
|
{"role": "tool", "tool_call_id": "c1", "content": "A"},
|
|
{"role": "tool", "tool_call_id": "c2", "content": "B"},
|
|
]
|
|
|
|
|
|
class TestBuildMessagesLegacyFieldsStillWork:
|
|
"""prompt=, system=, attachments= keep working — they synthesize
|
|
messages via Prompt.messages before build_messages sees them."""
|
|
|
|
def test_prompt_only(self, chat_model):
|
|
prompt = Prompt("hi", model=chat_model)
|
|
result = chat_model.build_messages(prompt, None)
|
|
assert result == [{"role": "user", "content": "hi"}]
|
|
|
|
def test_system_and_prompt(self, chat_model):
|
|
prompt = Prompt("hi", model=chat_model, system="be brief")
|
|
result = chat_model.build_messages(prompt, None)
|
|
assert result == [
|
|
{"role": "system", "content": "be brief"},
|
|
{"role": "user", "content": "hi"},
|
|
]
|
|
|
|
def test_attachments(self, chat_model):
|
|
att = llm.Attachment(type="image/jpeg", url="http://example.com/a.jpg")
|
|
prompt = Prompt("look", model=chat_model, attachments=[att])
|
|
result = chat_model.build_messages(prompt, None)
|
|
assert result == [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "look"},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": "http://example.com/a.jpg"},
|
|
},
|
|
],
|
|
}
|
|
]
|
|
|
|
|
|
class TestBuildMessagesSystemDedup:
|
|
"""Explicit messages with repeated system messages dedupe
|
|
repeated unchanged systems; OpenAI accepts one."""
|
|
|
|
def test_same_system_not_repeated(self, chat_model):
|
|
prompt = Prompt(
|
|
None,
|
|
model=chat_model,
|
|
messages=[
|
|
llm.system("be brief"),
|
|
llm.user("q1"),
|
|
llm.assistant("a1"),
|
|
llm.system("be brief"),
|
|
llm.user("q2"),
|
|
],
|
|
)
|
|
result = chat_model.build_messages(prompt, None)
|
|
system_msgs = [m for m in result if m["role"] == "system"]
|
|
assert len(system_msgs) == 1
|
|
assert system_msgs[0]["content"] == "be brief"
|
|
|
|
def test_system_change_emitted(self, chat_model):
|
|
prompt = Prompt(
|
|
None,
|
|
model=chat_model,
|
|
messages=[
|
|
llm.system("be brief"),
|
|
llm.user("q1"),
|
|
llm.assistant("a1"),
|
|
llm.system("be expansive"),
|
|
llm.user("q2"),
|
|
],
|
|
)
|
|
result = chat_model.build_messages(prompt, None)
|
|
system_msgs = [m for m in result if m["role"] == "system"]
|
|
assert [m["content"] for m in system_msgs] == [
|
|
"be brief",
|
|
"be expansive",
|
|
]
|
|
|
|
|
|
class TestBuildMessagesConversationHistory:
|
|
def test_prior_turn_text_plus_current_user(self, chat_model):
|
|
new_prompt = Prompt(
|
|
None,
|
|
model=chat_model,
|
|
messages=[
|
|
llm.user("what's 1+1?"),
|
|
llm.assistant("2"),
|
|
llm.user("what about 2+2?"),
|
|
],
|
|
)
|
|
result = chat_model.build_messages(new_prompt, None)
|
|
assert result == [
|
|
{"role": "user", "content": "what's 1+1?"},
|
|
{"role": "assistant", "content": "2"},
|
|
{"role": "user", "content": "what about 2+2?"},
|
|
]
|
|
|
|
def test_no_double_emission_from_conversation_prompt_flow(
|
|
self, chat_model, httpx_mock
|
|
):
|
|
# Two staged responses so conv.prompt twice can complete.
|
|
httpx_mock.add_response(
|
|
method="POST",
|
|
url="https://api.openai.com/v1/chat/completions",
|
|
json={
|
|
"model": "gpt-4o-mini",
|
|
"usage": {
|
|
"prompt_tokens": 1,
|
|
"completion_tokens": 1,
|
|
"total_tokens": 2,
|
|
},
|
|
"choices": [
|
|
{
|
|
"message": {"role": "assistant", "content": "A1"},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
},
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
httpx_mock.add_response(
|
|
method="POST",
|
|
url="https://api.openai.com/v1/chat/completions",
|
|
json={
|
|
"model": "gpt-4o-mini",
|
|
"usage": {
|
|
"prompt_tokens": 1,
|
|
"completion_tokens": 1,
|
|
"total_tokens": 2,
|
|
},
|
|
"choices": [
|
|
{
|
|
"message": {"role": "assistant", "content": "A2"},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
},
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
|
|
model = llm.get_model("gpt-4o-mini")
|
|
conv = model.conversation()
|
|
r1 = conv.prompt("Q1", key=API_KEY, stream=False)
|
|
r1.text()
|
|
r2 = conv.prompt("Q2", key=API_KEY, stream=False)
|
|
r2.text()
|
|
|
|
# Inspect what was sent on the SECOND turn.
|
|
sent_body = json.loads(httpx_mock.get_requests()[-1].content)
|
|
sent_messages = sent_body["messages"]
|
|
# Exactly three: user(Q1), assistant(A1), user(Q2).
|
|
assert sent_messages == [
|
|
{"role": "user", "content": "Q1"},
|
|
{"role": "assistant", "content": "A1"},
|
|
{"role": "user", "content": "Q2"},
|
|
]
|
|
|
|
|
|
class TestStreamingExecuteYieldsStreamEvents:
|
|
def test_text_stream_yields_text_events(self, httpx_mock):
|
|
httpx_mock.add_response(
|
|
method="POST",
|
|
url="https://api.openai.com/v1/chat/completions",
|
|
stream=IteratorStream(_text_stream()),
|
|
headers={"Content-Type": "text/event-stream"},
|
|
)
|
|
model = llm.get_model("gpt-4o-mini")
|
|
response = model.prompt("hi", key=API_KEY)
|
|
events = list(response.stream_events())
|
|
# At least one StreamEvent, all text, all at part_index=0.
|
|
assert events, "expected stream events"
|
|
assert all(isinstance(e, llm.parts.StreamEvent) for e in events)
|
|
assert all(e.type == "text" for e in events)
|
|
assert all(e.part_index == 0 for e in events)
|
|
# Text chunks concatenate to the expected full text.
|
|
assert "".join(e.chunk for e in events) == "Hello"
|
|
|
|
def test_text_stream_plain_iteration_still_returns_strings(self, httpx_mock):
|
|
httpx_mock.add_response(
|
|
method="POST",
|
|
url="https://api.openai.com/v1/chat/completions",
|
|
stream=IteratorStream(_text_stream()),
|
|
headers={"Content-Type": "text/event-stream"},
|
|
)
|
|
model = llm.get_model("gpt-4o-mini")
|
|
response = model.prompt("hi", key=API_KEY)
|
|
chunks = list(response)
|
|
assert all(isinstance(c, str) for c in chunks)
|
|
assert "".join(chunks) == "Hello"
|
|
|
|
def test_text_stream_messages_assembled(self, httpx_mock):
|
|
httpx_mock.add_response(
|
|
method="POST",
|
|
url="https://api.openai.com/v1/chat/completions",
|
|
stream=IteratorStream(_text_stream()),
|
|
headers={"Content-Type": "text/event-stream"},
|
|
)
|
|
model = llm.get_model("gpt-4o-mini")
|
|
response = model.prompt("hi", key=API_KEY)
|
|
response.text()
|
|
assert response.messages() == [
|
|
llm.Message(role="assistant", parts=[llm.parts.TextPart(text="Hello")])
|
|
]
|
|
|
|
def test_tool_call_stream_yields_name_and_args_events(self, httpx_mock):
|
|
httpx_mock.add_response(
|
|
method="POST",
|
|
url="https://api.openai.com/v1/chat/completions",
|
|
stream=IteratorStream(_tool_call_stream()),
|
|
headers={"Content-Type": "text/event-stream"},
|
|
)
|
|
|
|
def get_weather(city: str) -> str:
|
|
"Look up the weather."
|
|
return "sunny"
|
|
|
|
model = llm.get_model("gpt-4o-mini")
|
|
response = model.prompt("weather?", tools=[get_weather], key=API_KEY)
|
|
events = list(response.stream_events())
|
|
types = [e.type for e in events]
|
|
assert "tool_call_name" in types
|
|
assert "tool_call_args" in types
|
|
# Name event carries the tool_call_id and name.
|
|
name_ev = next(e for e in events if e.type == "tool_call_name")
|
|
assert name_ev.tool_call_id == "call_1"
|
|
assert name_ev.chunk == "get_weather"
|
|
# Args events share the same part_index and concatenate to
|
|
# valid JSON.
|
|
args_events = [e for e in events if e.type == "tool_call_args"]
|
|
assert all(e.part_index == name_ev.part_index for e in args_events)
|
|
assert json.loads("".join(e.chunk for e in args_events)) == {"city": "Paris"}
|
|
|
|
def test_tool_call_registered_via_add_tool_call(self, httpx_mock):
|
|
"""response.tool_calls() still works — chain/execute relies on it."""
|
|
httpx_mock.add_response(
|
|
method="POST",
|
|
url="https://api.openai.com/v1/chat/completions",
|
|
stream=IteratorStream(_tool_call_stream()),
|
|
headers={"Content-Type": "text/event-stream"},
|
|
)
|
|
|
|
def get_weather(city: str) -> str:
|
|
"Look up the weather."
|
|
return "sunny"
|
|
|
|
model = llm.get_model("gpt-4o-mini")
|
|
response = model.prompt("weather?", tools=[get_weather], key=API_KEY)
|
|
response.text()
|
|
tcs = response.tool_calls()
|
|
assert len(tcs) == 1
|
|
assert tcs[0].name == "get_weather"
|
|
assert tcs[0].arguments == {"city": "Paris"}
|
|
assert tcs[0].tool_call_id == "call_1"
|
|
|
|
def test_text_then_tool_call_part_index_advances(self, httpx_mock):
|
|
httpx_mock.add_response(
|
|
method="POST",
|
|
url="https://api.openai.com/v1/chat/completions",
|
|
stream=IteratorStream(_text_then_tool_call_stream()),
|
|
headers={"Content-Type": "text/event-stream"},
|
|
)
|
|
|
|
def get_weather(c: int) -> str:
|
|
"Weather."
|
|
return "sunny"
|
|
|
|
model = llm.get_model("gpt-4o-mini")
|
|
response = model.prompt("q", tools=[get_weather], key=API_KEY)
|
|
response.text()
|
|
# After streaming, messages has both a TextPart and a ToolCallPart.
|
|
parts = response.messages()[0].parts
|
|
assert any(isinstance(p, llm.parts.TextPart) for p in parts)
|
|
assert any(isinstance(p, llm.parts.ToolCallPart) for p in parts)
|
|
text_part = next(p for p in parts if isinstance(p, llm.parts.TextPart))
|
|
tc_part = next(p for p in parts if isinstance(p, llm.parts.ToolCallPart))
|
|
assert text_part.text == "Looking up"
|
|
assert tc_part.name == "get_weather"
|
|
assert tc_part.arguments == {"c": 1}
|
|
|
|
|
|
class TestAsyncStreamingExecuteYieldsStreamEvents:
|
|
@pytest.mark.asyncio
|
|
async def test_text_stream_yields_text_events(self, httpx_mock):
|
|
httpx_mock.add_response(
|
|
method="POST",
|
|
url="https://api.openai.com/v1/chat/completions",
|
|
stream=IteratorStream(_text_stream()),
|
|
headers={"Content-Type": "text/event-stream"},
|
|
)
|
|
model = llm.get_async_model("gpt-4o-mini")
|
|
response = model.prompt("hi", key=API_KEY)
|
|
events = []
|
|
async for event in response.astream_events():
|
|
events.append(event)
|
|
assert all(isinstance(e, llm.parts.StreamEvent) for e in events)
|
|
assert [e.type for e in events] == ["text"] * len(events)
|
|
assert "".join(e.chunk for e in events) == "Hello"
|
|
|
|
|
|
def _text_stream_with_reasoning_usage(reasoning_tokens):
|
|
"""Stream with usage in the final chunk reporting reasoning_tokens."""
|
|
yield _sse({"role": "assistant", "content": ""})
|
|
yield _sse({"content": "Hel"})
|
|
yield _sse({"content": "lo"})
|
|
yield _sse({}, finish_reason="stop")
|
|
# Final chunk with usage — OpenAI streams usage once at the end
|
|
# when stream_options.include_usage=True.
|
|
yield _sse(
|
|
{},
|
|
usage={
|
|
"prompt_tokens": 5,
|
|
"completion_tokens": 2,
|
|
"total_tokens": 7,
|
|
"completion_tokens_details": {"reasoning_tokens": reasoning_tokens},
|
|
},
|
|
)
|
|
yield b"data: [DONE]\n\n"
|
|
|
|
|
|
class TestReasoningTokenCount:
|
|
def test_redacted_reasoning_part_emitted_when_count_present(self, httpx_mock):
|
|
httpx_mock.add_response(
|
|
method="POST",
|
|
url="https://api.openai.com/v1/chat/completions",
|
|
stream=IteratorStream(_text_stream_with_reasoning_usage(150)),
|
|
headers={"Content-Type": "text/event-stream"},
|
|
)
|
|
model = llm.get_model("gpt-4o-mini")
|
|
response = model.prompt("hi", key=API_KEY)
|
|
response.text()
|
|
assert response.messages() == [
|
|
llm.Message(
|
|
role="assistant",
|
|
parts=[
|
|
llm.parts.ReasoningPart(text="", redacted=True),
|
|
llm.parts.TextPart(text="Hello"),
|
|
],
|
|
)
|
|
]
|
|
|
|
def test_no_reasoning_part_when_zero_or_absent(self, httpx_mock):
|
|
httpx_mock.add_response(
|
|
method="POST",
|
|
url="https://api.openai.com/v1/chat/completions",
|
|
stream=IteratorStream(_text_stream_with_reasoning_usage(0)),
|
|
headers={"Content-Type": "text/event-stream"},
|
|
)
|
|
model = llm.get_model("gpt-4o-mini")
|
|
response = model.prompt("hi", key=API_KEY)
|
|
response.text()
|
|
parts = response.messages()[0].parts
|
|
assert not any(
|
|
isinstance(p, llm.parts.ReasoningPart) for p in parts
|
|
), "should not add a redacted reasoning part when count=0"
|
|
|
|
|
|
class TestNonStreamingExecuteYieldsStreamEvents:
|
|
def test_non_streaming_text_yields_single_event(self, httpx_mock):
|
|
httpx_mock.add_response(
|
|
method="POST",
|
|
url="https://api.openai.com/v1/chat/completions",
|
|
json={
|
|
"model": "gpt-4o-mini",
|
|
"usage": {
|
|
"prompt_tokens": 1,
|
|
"completion_tokens": 1,
|
|
"total_tokens": 2,
|
|
},
|
|
"choices": [
|
|
{
|
|
"message": {"role": "assistant", "content": "Hello"},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
},
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
model = llm.get_model("gpt-4o-mini")
|
|
response = model.prompt("hi", key=API_KEY, stream=False)
|
|
events = list(response.stream_events())
|
|
assert events == [
|
|
llm.parts.StreamEvent(type="text", chunk="Hello", part_index=0)
|
|
]
|
|
assert response.messages() == [
|
|
llm.Message(role="assistant", parts=[llm.parts.TextPart(text="Hello")])
|
|
]
|