wehub-resource-sync
380 lines
12 KiB
Python
380 lines
12 KiB
Python
"""Async parity: every sync API must work the same
|
|
way on AsyncResponse and AsyncConversation.
|
|
|
|
Uses the llm-echo plugin (sync ``Echo`` + async ``EchoAsync``) so both
|
|
paths exercise real registered models with identical behaviour.
|
|
"""
|
|
|
|
import json
|
|
|
|
import llm
|
|
import pytest
|
|
|
|
# ---- basic sanity: both variants are registered --------------------
|
|
|
|
|
|
def test_echo_registered_for_both():
|
|
assert isinstance(llm.get_model("echo"), llm.Model)
|
|
assert isinstance(llm.get_async_model("echo"), llm.AsyncModel)
|
|
|
|
|
|
# ---- AsyncResponse.to_dict / from_dict -----------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_to_dict_captures_chain_and_output():
|
|
model = llm.get_async_model("echo")
|
|
r = model.prompt("hello")
|
|
await r.text()
|
|
|
|
d = r.to_dict()
|
|
assert d["model"] == "echo"
|
|
assert d["prompt"]["messages"] == [llm.user("hello").to_dict()]
|
|
# Echo's output is JSON describing the input; it's the assistant's text.
|
|
assert len(d["messages"]) == 1
|
|
assert d["messages"][0]["role"] == "assistant"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_to_dict_raises_before_awaited():
|
|
model = llm.get_async_model("echo")
|
|
r = model.prompt("hello")
|
|
with pytest.raises(ValueError):
|
|
r.to_dict()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_from_dict_rehydrates():
|
|
model = llm.get_async_model("echo")
|
|
r = model.prompt("hello")
|
|
await r.text()
|
|
|
|
payload = json.dumps(r.to_dict())
|
|
restored = llm.AsyncResponse.from_dict(json.loads(payload))
|
|
|
|
assert restored._done
|
|
# text_or_raise should match (same text as original)
|
|
assert restored.text_or_raise() == r.text_or_raise()
|
|
# messages structure preserved
|
|
assert await restored.messages() == await r.messages()
|
|
# prompt.messages (the chain that was sent) preserved
|
|
assert restored.prompt.messages == r.prompt.messages
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_from_dict_then_reply_continues():
|
|
"""Persist an async response across process
|
|
boundary (via JSON), rehydrate, continue with reply()."""
|
|
model = llm.get_async_model("echo")
|
|
r1 = model.prompt("q1")
|
|
await r1.text()
|
|
|
|
payload = json.dumps(r1.to_dict())
|
|
restored = llm.AsyncResponse.from_dict(json.loads(payload))
|
|
|
|
r2 = await restored.reply("q2")
|
|
await r2.text()
|
|
|
|
# r2 was sent the full chain including r1's output.
|
|
chain_roles = [m.role for m in r2.prompt.messages]
|
|
assert chain_roles == ["user", "assistant", "user"]
|
|
assert r2.prompt.messages[0].parts[0].text == "q1"
|
|
assert r2.prompt.messages[-1].parts[0].text == "q2"
|
|
|
|
|
|
# ---- AsyncResponse rehydrated via from_row (SQLite path) -----------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_from_row_response_messages_synthesized(tmp_path):
|
|
"""SQLite rehydrate for async responses must populate
|
|
response.messages from _chunks+_tool_calls so follow-up chains
|
|
don't silently drop the assistant turn."""
|
|
import sqlite_utils
|
|
from llm.migrations import migrate
|
|
|
|
model = llm.get_async_model("echo")
|
|
r = model.prompt("hello")
|
|
await r.text()
|
|
|
|
db = sqlite_utils.Database(str(tmp_path / "logs.db"))
|
|
migrate(db)
|
|
# to_sync_response is what log_to_db uses for async.
|
|
sync_r = await r.to_sync_response()
|
|
sync_r.log_to_db(db)
|
|
|
|
row = next(db["responses"].rows)
|
|
rehydrated = llm.AsyncResponse.from_row(db, row)
|
|
|
|
assert rehydrated._stream_events == []
|
|
# response.messages falls back to _chunks — must not be empty.
|
|
msgs = await rehydrated.messages()
|
|
assert len(msgs) == 1
|
|
assert msgs[0].role == "assistant"
|
|
assert isinstance(msgs[0].parts[0], llm.parts.TextPart)
|
|
|
|
|
|
# ---- AsyncConversation follow-up via load_conversation -------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_load_conversation_follow_up_preserves_chain(tmp_path):
|
|
"""Async equivalent of the llm -c regression: after log_to_db +
|
|
load_conversation, a follow-up turn's prompt.messages is the full
|
|
[user, assistant, user] chain — not missing the assistant."""
|
|
import sqlite_utils
|
|
from llm.cli import load_conversation
|
|
from llm.migrations import migrate
|
|
|
|
model = llm.get_async_model("echo")
|
|
r1 = model.prompt("q1")
|
|
await r1.text()
|
|
|
|
db_path = tmp_path / "logs.db"
|
|
db = sqlite_utils.Database(str(db_path))
|
|
migrate(db)
|
|
(await r1.to_sync_response()).log_to_db(db)
|
|
|
|
conv = load_conversation(None, async_=True, database=str(db_path))
|
|
r2 = conv.prompt("q2")
|
|
await r2.text()
|
|
|
|
chain = r2.prompt.messages
|
|
assert [m.role for m in chain] == ["user", "assistant", "user"]
|
|
assert chain[0].parts[0].text == "q1"
|
|
assert chain[-1].parts[0].text == "q2"
|
|
|
|
|
|
# ---- Sync/async semantic parity for reply()+to_dict() --------------
|
|
|
|
|
|
def _capture_sync(model):
|
|
r1 = model.prompt("ping")
|
|
r1.text()
|
|
payload1 = json.dumps(r1.to_dict())
|
|
restored = llm.Response.from_dict(json.loads(payload1))
|
|
r2 = restored.reply("pong")
|
|
r2.text()
|
|
return r2.prompt.messages
|
|
|
|
|
|
async def _capture_async(model):
|
|
r1 = model.prompt("ping")
|
|
await r1.text()
|
|
payload1 = json.dumps(r1.to_dict())
|
|
restored = llm.AsyncResponse.from_dict(json.loads(payload1))
|
|
r2 = await restored.reply("pong")
|
|
await r2.text()
|
|
return r2.prompt.messages
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sync_and_async_produce_identical_chain():
|
|
"""Run the full save → restore → reply loop against sync Echo and
|
|
async EchoAsync. The chain sent on the second turn must be
|
|
structurally identical."""
|
|
sync_chain = _capture_sync(llm.get_model("echo"))
|
|
async_chain = await _capture_async(llm.get_async_model("echo"))
|
|
|
|
# Echo's assistant output differs between invocations only in
|
|
# the "previous" field — but for the first turn both see empty
|
|
# previous, so outputs match.
|
|
sync_dicts = [m.to_dict() for m in sync_chain]
|
|
async_dicts = [m.to_dict() for m in async_chain]
|
|
assert sync_dicts == async_dicts
|
|
|
|
|
|
# ---- AsyncChainResponse tool-result turn pre-bakes chain -----------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_chain_tool_result_turn_has_full_chain():
|
|
"""AsyncChainResponse must pre-bake the full chain on tool-result
|
|
turns, same as sync ChainResponse."""
|
|
|
|
async def my_tool(x: int) -> int:
|
|
"Double the input."
|
|
return x * 2
|
|
|
|
model = llm.get_async_model("echo")
|
|
# Drive a one-iteration chain by asking echo to emit a tool call
|
|
# (echo's JSON-prompt syntax).
|
|
chain = model.chain(
|
|
json.dumps(
|
|
{
|
|
"tool_calls": [{"name": "my_tool", "arguments": {"x": 5}}],
|
|
"prompt": "prompt",
|
|
}
|
|
),
|
|
tools=[llm.Tool.function(my_tool, name="my_tool")],
|
|
)
|
|
|
|
responses = []
|
|
async for response in chain.responses():
|
|
responses.append(response)
|
|
|
|
# Two responses: the tool-call turn and the tool-result turn.
|
|
assert len(responses) == 2
|
|
second = responses[1]
|
|
# Second turn's prompt.messages includes the prior turn (user +
|
|
# assistant with tool call) plus a tool-role message with the result.
|
|
chain_roles = [m.role for m in second.prompt.messages]
|
|
assert "tool" in chain_roles
|
|
assert chain_roles[0] == "user"
|
|
|
|
|
|
# ---- astream_events() parity with stream_events() ------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_astream_events_matches_stream_events_for_text_only():
|
|
"""Echo yields plain str (legacy plugin). Both sync and async
|
|
paths should wrap those into StreamEvent(type='text') with the
|
|
same shape."""
|
|
sync_model = llm.get_model("echo")
|
|
async_model = llm.get_async_model("echo")
|
|
|
|
sync_r = sync_model.prompt("hello")
|
|
sync_events = list(sync_r.stream_events())
|
|
|
|
async_r = async_model.prompt("hello")
|
|
async_events = []
|
|
async for ev in async_r.astream_events():
|
|
async_events.append(ev)
|
|
|
|
# Same event types, same payload.
|
|
assert [e.type for e in sync_events] == [e.type for e in async_events]
|
|
assert all(e.type == "text" for e in sync_events)
|
|
assert "".join(e.chunk for e in sync_events) == "".join(
|
|
e.chunk for e in async_events
|
|
)
|
|
|
|
|
|
# ---- Async reply chaining --------------------------------------------
|
|
|
|
|
|
# ---- Additional edge cases ----------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_from_dict_model_override():
|
|
model = llm.get_async_model("echo")
|
|
r = model.prompt("hi")
|
|
await r.text()
|
|
payload = json.dumps(r.to_dict())
|
|
|
|
# Pass model explicitly to override whatever's in the payload.
|
|
alt = llm.get_async_model("echo")
|
|
restored = llm.AsyncResponse.from_dict(json.loads(payload), model=alt)
|
|
assert restored.model is alt
|
|
|
|
|
|
def test_sync_from_dict_model_override():
|
|
model = llm.get_model("echo")
|
|
r = model.prompt("hi")
|
|
r.text()
|
|
payload = json.dumps(r.to_dict())
|
|
|
|
alt = llm.get_model("echo")
|
|
restored = llm.Response.from_dict(json.loads(payload), model=alt)
|
|
assert restored.model is alt
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_to_dict_preserves_datetime():
|
|
model = llm.get_async_model("echo")
|
|
r = model.prompt("hi")
|
|
await r.text()
|
|
d = r.to_dict()
|
|
assert "datetime_utc" in d
|
|
assert isinstance(d["datetime_utc"], str)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_to_dict_preserves_usage_when_set(async_mock_model):
|
|
"""When a plugin calls response.set_usage, to_dict captures it.
|
|
async_mock_model does set usage; llm-echo's async variant doesn't."""
|
|
async_mock_model.enqueue(["ok"])
|
|
r = async_mock_model.prompt("hi")
|
|
await r.text()
|
|
d = r.to_dict()
|
|
assert "usage" in d
|
|
assert d["usage"]["input"] is not None
|
|
assert d["usage"]["output"] is not None
|
|
|
|
# And it round-trips.
|
|
restored = llm.AsyncResponse.from_dict(d, model=async_mock_model)
|
|
assert restored.input_tokens == d["usage"]["input"]
|
|
assert restored.output_tokens == d["usage"]["output"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_reply_messages_kwarg_appends():
|
|
"""AsyncResponse.reply(messages=[...]) appends extra messages onto
|
|
the chain in place of a trailing user string (mirrors sync test)."""
|
|
model = llm.get_async_model("echo")
|
|
r1 = model.prompt("q1")
|
|
await r1.text()
|
|
r2 = await r1.reply(messages=[llm.user("extra")])
|
|
await r2.text()
|
|
assert [m.role for m in r2.prompt.messages] == ["user", "assistant", "user"]
|
|
assert r2.prompt.messages[-1].parts[0].text == "extra"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_full_chain_to_dict_round_trip_three_turns():
|
|
"""Serialize on turn 3 — chain must include q1, a1, q2, a2, q3 on
|
|
round-trip."""
|
|
model = llm.get_async_model("echo")
|
|
r1 = model.prompt("q1")
|
|
await r1.text()
|
|
r2 = await r1.reply("q2")
|
|
await r2.text()
|
|
r3 = await r2.reply("q3")
|
|
await r3.text()
|
|
|
|
payload = json.dumps(r3.to_dict())
|
|
restored = llm.AsyncResponse.from_dict(json.loads(payload))
|
|
assert [m.role for m in restored.prompt.messages] == [
|
|
"user",
|
|
"assistant",
|
|
"user",
|
|
"assistant",
|
|
"user",
|
|
]
|
|
texts = [m.parts[0].text for m in restored.prompt.messages if m.parts]
|
|
assert texts[0] == "q1"
|
|
assert texts[2] == "q2"
|
|
assert texts[4] == "q3"
|
|
|
|
# And continuing from the restored response extends the chain.
|
|
r4 = await restored.reply("q4")
|
|
await r4.text()
|
|
assert [m.role for m in r4.prompt.messages] == [
|
|
"user",
|
|
"assistant",
|
|
"user",
|
|
"assistant",
|
|
"user",
|
|
"assistant",
|
|
"user",
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_reply_chains_three_turns():
|
|
model = llm.get_async_model("echo")
|
|
r1 = model.prompt("q1")
|
|
await r1.text()
|
|
r2 = await r1.reply("q2")
|
|
await r2.text()
|
|
r3 = await r2.reply("q3")
|
|
await r3.text()
|
|
|
|
chain = r3.prompt.messages
|
|
assert [m.role for m in chain] == ["user", "assistant", "user", "assistant", "user"]
|
|
texts = [m.parts[0].text for m in chain if m.parts]
|
|
assert texts[0] == "q1"
|
|
assert texts[2] == "q2"
|
|
assert texts[4] == "q3"
|