351 lines
14 KiB
Python
351 lines
14 KiB
Python
"""Reject unfinished Responses output before parsers or tool consumers see it."""
|
|
|
|
import asyncio
|
|
import json
|
|
import hashlib
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import httpx2 as httpx
|
|
from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, ToolMessage
|
|
from langchain_core.runnables import RunnableLambda
|
|
from pydantic import BaseModel, SecretStr
|
|
import pytest
|
|
import pytest_asyncio
|
|
|
|
from artemis.llm.reliability import FailureCategory, LLMExhaustedError
|
|
from artemis.llm.response_guard import ResponseIntegrityError, ensure_response_complete
|
|
from artemis.llm.router import ModelEndpoint, ModelFactory
|
|
from artemis.llm.diagnostics import RequestAttempt
|
|
from artemis.services import llm as llm_service
|
|
from artemis.services.llm import RobustChatModelWrapper, acomplete
|
|
|
|
|
|
def response_body(status, arguments, name):
|
|
return {
|
|
"id": "resp_test",
|
|
"object": "response",
|
|
"created_at": 1,
|
|
"status": status,
|
|
"model": "test-model",
|
|
"output": [
|
|
{
|
|
"type": "function_call",
|
|
"id": "fc_test",
|
|
"call_id": "call_test",
|
|
"name": name,
|
|
"arguments": arguments,
|
|
"status": "completed",
|
|
}
|
|
],
|
|
"incomplete_details": {"reason": "max_output_tokens"} if status == "incomplete" else None,
|
|
"usage": {"input_tokens": 10, "output_tokens": 3, "total_tokens": 13},
|
|
}
|
|
|
|
|
|
def response_stream(status, arguments, name):
|
|
body = response_body(status or "in_progress", arguments, name)
|
|
events = [
|
|
{"type": "response.created", "response": {**body, "status": "in_progress", "output": []}},
|
|
{
|
|
"type": "response.output_item.added",
|
|
"output_index": 0,
|
|
"item": {**body["output"][0], "arguments": "", "status": "in_progress"},
|
|
},
|
|
{
|
|
"type": "response.function_call_arguments.delta",
|
|
"output_index": 0,
|
|
"item_id": "fc_test",
|
|
"delta": arguments,
|
|
},
|
|
]
|
|
if status:
|
|
events.append({"type": "response." + status, "response": body})
|
|
return (
|
|
"".join("data: " + json.dumps(event, ensure_ascii=False) + "\n\n" for event in events)
|
|
+ "data: [DONE]\n\n"
|
|
)
|
|
|
|
|
|
class FragmentedEventStream(httpx.AsyncByteStream):
|
|
"""Split SSE frames and multibyte UTF-8 characters at transport boundaries."""
|
|
|
|
def __init__(self, text):
|
|
self.data = text.encode("utf-8")
|
|
|
|
async def __aiter__(self):
|
|
for offset in range(0, len(self.data), 4093):
|
|
yield self.data[offset : offset + 4093]
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def transport_case(monkeypatch, tmp_path, request):
|
|
from langchain_openai import ChatOpenAI
|
|
import openai
|
|
|
|
requests, outputs, clients, events = [], [], [], []
|
|
|
|
def handle(http_request):
|
|
payload = json.loads(http_request.content)
|
|
requests.append(payload)
|
|
status, arguments = outputs.pop(0)
|
|
name = payload.get("tools", [{"name": "finish"}])[0]["name"]
|
|
if payload.get("stream"):
|
|
stream = response_stream(status, arguments, name)
|
|
body = (
|
|
{"stream": FragmentedEventStream(stream)}
|
|
if getattr(request, "param", None) == "fragmented"
|
|
else {"text": stream}
|
|
)
|
|
return httpx.Response(
|
|
200,
|
|
headers={"content-type": "text/event-stream"},
|
|
**body,
|
|
)
|
|
return httpx.Response(200, json=response_body(status, arguments, name))
|
|
|
|
def http_client(**kwargs):
|
|
client = httpx.AsyncClient(transport=httpx.MockTransport(handle), **kwargs)
|
|
clients.append(client)
|
|
return client
|
|
|
|
monkeypatch.setattr(openai, "DefaultAsyncHttpxClient", http_client)
|
|
monkeypatch.setattr(llm_service.settings, "OPENAI_API_KEY", SecretStr("test-secret"))
|
|
monkeypatch.setattr(llm_service.settings, "OPENAI_BASE_URL", "https://responses.example/v1")
|
|
monkeypatch.setattr(llm_service.settings, "OPENAI_MAX_RETRIES", 5)
|
|
monkeypatch.setattr(llm_service.settings, "OPENAI_RETRY_INTERVAL_SECONDS", 0)
|
|
monkeypatch.delenv("ARTEMIS_FAKE_LLM", raising=False)
|
|
monkeypatch.setattr(llm_service, "PAUSE_FILE", tmp_path / ".artemis_paused")
|
|
monkeypatch.setattr(llm_service, "_wait_for_resume", AsyncMock(return_value=False))
|
|
monkeypatch.setattr(llm_service, "_get_current_data_engine", lambda: None)
|
|
monkeypatch.setattr(
|
|
llm_service, "_record_llm_event", lambda *args, **kwargs: events.append(args)
|
|
)
|
|
llm_service._ENDPOINT_BREAKER._states.clear()
|
|
llm_service._NON_STREAMING_ENDPOINTS.clear()
|
|
endpoint = ModelEndpoint(model_name="test-model", transport="http", max_retries=0)
|
|
model = ModelFactory.create_model(endpoint)
|
|
assert isinstance(model, ChatOpenAI)
|
|
yield SimpleNamespace(
|
|
model=model, endpoint=endpoint, requests=requests, outputs=outputs, events=events
|
|
)
|
|
for client in clients:
|
|
await client.aclose()
|
|
llm_service._ENDPOINT_BREAKER._states.clear()
|
|
llm_service._NON_STREAMING_ENDPOINTS.clear()
|
|
|
|
|
|
def wrapped(case):
|
|
return RobustChatModelWrapper(case.model, endpoint=case.endpoint).bind_tools(
|
|
[
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "finish",
|
|
"description": "Record the verdict",
|
|
"parameters": {"type": "object", "properties": {"passed": {"type": "boolean"}}},
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("status", [None, "incomplete", "failed"])
|
|
@pytest.mark.parametrize("arguments", ['{"passed":true}', '{"passed":true'])
|
|
async def test_unfinished_stream_never_returns_a_verdict_or_tool_call(
|
|
transport_case, status, arguments
|
|
):
|
|
case = transport_case
|
|
case.outputs.append((status, arguments))
|
|
with pytest.raises(LLMExhaustedError) as failure:
|
|
await wrapped(case).complete([HumanMessage(content="Verify the task")])
|
|
assert failure.value.failure.category is FailureCategory.INVALID_OUTPUT
|
|
assert len(case.requests) == 1
|
|
resets = [payload for event, payload in case.events if event == "llm_stream_reset"]
|
|
assert len(resets) == 1 and resets[0]["action"] == "discard"
|
|
assert not llm_service._NON_STREAMING_ENDPOINTS
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("arguments", ['{"passed":true', ""])
|
|
async def test_completed_event_cannot_repair_invalid_tool_json(transport_case, arguments):
|
|
case = transport_case
|
|
case.outputs.append(("completed", arguments))
|
|
with pytest.raises(LLMExhaustedError) as failure:
|
|
await wrapped(case).complete([])
|
|
assert failure.value.failure.category is FailureCategory.INVALID_OUTPUT
|
|
assert len(case.requests) == 1
|
|
|
|
|
|
@pytest.mark.parametrize("arguments", [None, "[]", "null", '{"passed":true'])
|
|
def test_original_function_call_arguments_must_be_a_complete_json_object(arguments):
|
|
message = AIMessage(
|
|
content=[{"type": "function_call", "arguments": arguments}],
|
|
response_metadata={"model_provider": "openai", "status": "completed"},
|
|
)
|
|
with pytest.raises(ResponseIntegrityError, match="invalid_tool_arguments"):
|
|
ensure_response_complete(message)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_completed_stream_returns_original_valid_arguments(transport_case):
|
|
case = transport_case
|
|
case.outputs.append(("completed", '{"passed":true}'))
|
|
result = await wrapped(case).complete([])
|
|
assert result.response_metadata["status"] == "completed"
|
|
assert result.tool_calls[0]["args"] == {"passed": True}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("transport_case", ["fragmented"], indirect=True)
|
|
async def test_real_sdk_fingerprint_counts_received_arguments_once(transport_case):
|
|
case = transport_case
|
|
arguments = json.dumps({"passed": True, "evidence": "中文🙂" * 1000}, ensure_ascii=False)
|
|
case.outputs.append(("completed", arguments))
|
|
observation = RequestAttempt(
|
|
None, request_id="sdk-test", attempt=1, timeout=300, max_retries=0, metadata={}
|
|
)
|
|
merged = None
|
|
async for chunk in case.model.astream([]):
|
|
observation.chunk(chunk)
|
|
merged = chunk if merged is None else merged + chunk
|
|
observation.finish("completed", response=merged)
|
|
assert len(observation.data["tools"]) == 1
|
|
tool = observation.data["tools"][0]
|
|
assert tool["call_id"] == "call_test"
|
|
assert tool["argument_chars"] == len(arguments)
|
|
assert tool["argument_bytes"] == len(arguments.encode())
|
|
assert tool["argument_sha256"] == hashlib.sha256(arguments.encode()).hexdigest()
|
|
assert tool["argument_prefixes"][0] == {
|
|
"bytes": 4096,
|
|
"sha256": hashlib.sha256(arguments.encode()[:4096]).hexdigest(),
|
|
}
|
|
assert tool["argument_deltas"] == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("transport_case", ["fragmented"], indirect=True)
|
|
@pytest.mark.parametrize("status", ["completed", None])
|
|
async def test_large_unicode_tool_stream_requires_completion(transport_case, status):
|
|
case = transport_case
|
|
arguments = {"passed": True, "evidence": "截图已核验。\n" * 65536}
|
|
encoded = json.dumps(arguments, ensure_ascii=False)
|
|
assert len(encoded.encode("utf-8")) > 1024 * 1024
|
|
case.outputs.append((status, encoded))
|
|
|
|
if status == "completed":
|
|
result = await wrapped(case).complete([])
|
|
assert result.response_metadata["status"] == "completed"
|
|
assert result.tool_calls[0]["args"] == arguments
|
|
else:
|
|
with pytest.raises(LLMExhaustedError) as failure:
|
|
await wrapped(case).complete([])
|
|
assert failure.value.failure.category is FailureCategory.INVALID_OUTPUT
|
|
assert len(case.requests) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("recovers", [False, True])
|
|
@pytest.mark.parametrize("initial_status", [None, "failed"])
|
|
async def test_integrity_retry_is_bounded_and_keeps_existing_tool_evidence(
|
|
transport_case, recovers, initial_status
|
|
):
|
|
case = transport_case
|
|
case.endpoint.max_retries = 5
|
|
case.outputs.extend(
|
|
[
|
|
(initial_status, '{"passed":false}'),
|
|
("completed" if recovers else initial_status, '{"passed":true}'),
|
|
]
|
|
)
|
|
messages = [
|
|
HumanMessage(content="Inspect then verify"),
|
|
AIMessage(content="", tool_calls=[{"name": "inspect", "args": {}, "id": "prior_call"}]),
|
|
ToolMessage(content="Existing screenshot evidence", tool_call_id="prior_call"),
|
|
]
|
|
if recovers:
|
|
result = await wrapped(case).complete(messages)
|
|
assert result.tool_calls[0]["args"] == {"passed": True}
|
|
else:
|
|
with pytest.raises(LLMExhaustedError):
|
|
await wrapped(case).complete(messages)
|
|
assert len(case.requests) == 2
|
|
assert case.requests[0]["input"] == case.requests[1]["input"]
|
|
assert any(
|
|
item.get("call_id") == "prior_call" and item.get("output") == "Existing screenshot evidence"
|
|
for item in case.requests[1]["input"]
|
|
)
|
|
assert isinstance(messages[-1], ToolMessage)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("status", ["incomplete", "in_progress", "queued", "cancelled"])
|
|
async def test_non_streaming_response_also_requires_completed(transport_case, status):
|
|
case = transport_case
|
|
case.outputs.append((status, '{"passed":true}'))
|
|
with pytest.raises(LLMExhaustedError) as failure:
|
|
await wrapped(case).ainvoke([])
|
|
assert failure.value.failure.category is FailureCategory.INVALID_OUTPUT
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("include_raw", [False, True])
|
|
async def test_structured_output_is_checked_before_returning_parsed_verdict(
|
|
transport_case, include_raw
|
|
):
|
|
class Verdict(BaseModel):
|
|
passed: bool
|
|
|
|
case = transport_case
|
|
case.outputs.append(("incomplete", '{"passed":true}'))
|
|
model = RobustChatModelWrapper(case.model, endpoint=case.endpoint)
|
|
with pytest.raises(LLMExhaustedError):
|
|
await model.with_structured_output(Verdict, include_raw=include_raw).ainvoke("Verify")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_downstream_parser_never_receives_incomplete_response(transport_case):
|
|
case = transport_case
|
|
case.outputs.append(("incomplete", '{"passed":true}'))
|
|
parsed = []
|
|
chain = wrapped(case) | RunnableLambda(lambda value: parsed.append(value))
|
|
with pytest.raises(LLMExhaustedError):
|
|
await chain.ainvoke([])
|
|
assert parsed == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_raw_model_compatibility_path_does_not_accept_missing_terminal_event(transport_case):
|
|
case = transport_case
|
|
case.outputs.append((None, '{"passed":true}'))
|
|
with pytest.raises(ValueError):
|
|
await acomplete(case.model, [])
|
|
assert len(case.requests) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancellation_still_closes_stream_without_retry(transport_case):
|
|
started, closed = asyncio.Event(), []
|
|
|
|
class StreamingModel:
|
|
async def astream(self, *args, **kwargs):
|
|
try:
|
|
yield AIMessageChunk(
|
|
content="partial", response_metadata={"model_provider": "openai"}
|
|
)
|
|
started.set()
|
|
await asyncio.Event().wait()
|
|
finally:
|
|
closed.append(True)
|
|
|
|
task = asyncio.create_task(RobustChatModelWrapper(StreamingModel()).complete([]))
|
|
await asyncio.wait_for(started.wait(), timeout=2)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert closed == [True]
|
|
assert not transport_case.requests
|
|
resets = [payload for event, payload in transport_case.events if event == "llm_stream_reset"]
|
|
assert len(resets) == 1 and resets[0]["reason"] == "incomplete_response"
|