102 lines
5.5 KiB
Python
102 lines
5.5 KiB
Python
"""Exercise the real OpenAI adapter and LangGraph streaming against a local transport."""
|
|
|
|
import asyncio
|
|
import json
|
|
from types import SimpleNamespace
|
|
|
|
import httpx
|
|
import pytest
|
|
from langchain.agents import create_agent
|
|
from langchain.agents.middleware import AgentMiddleware
|
|
from langchain_core.messages import HumanMessage
|
|
|
|
from deerflow.agents.middlewares.context_window_guard_middleware import ContextWindowGuardMiddleware
|
|
from deerflow.agents.middlewares.llm_error_handling_middleware import LLMErrorHandlingMiddleware
|
|
from deerflow.models.patched_openai import PatchedChatOpenAI
|
|
|
|
|
|
class InjectOncePerCall(AgentMiddleware):
|
|
calls: int = 0
|
|
|
|
def wrap_model_call(self, request, handler):
|
|
self.calls += 1
|
|
return handler(request.override(messages=[*request.messages, HumanMessage(content="背景" * 10000, additional_kwargs={"hide_from_ui": True})]))
|
|
|
|
async def awrap_model_call(self, request, handler):
|
|
self.calls += 1
|
|
return await handler(request.override(messages=[*request.messages, HumanMessage(content="背景" * 10000, additional_kwargs={"hide_from_ui": True})]))
|
|
|
|
|
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
|
@pytest.mark.parametrize("streaming", [False, True])
|
|
def test_actual_openai_transport_and_agent_recover_without_reinjection(asynchronous, streaming):
|
|
payloads = []
|
|
|
|
def transport(request):
|
|
body = json.loads(request.content)
|
|
payloads.append(body)
|
|
size = sum(len(str(m.get("content", "")).encode("utf-8")) for m in body["messages"])
|
|
if len(payloads) == 1 or size > 2000:
|
|
return httpx.Response(400, json={"error": {"message": "maximum context length exceeded", "type": "invalid_request_error", "code": "context_length_exceeded"}})
|
|
if body.get("stream"):
|
|
chunks = [
|
|
{"id": "recovered", "object": "chat.completion.chunk", "created": 0, "model": "test", "choices": [{"index": 0, "delta": {"role": "assistant", "content": "正常回答"}, "finish_reason": None}]},
|
|
{"id": "recovered", "object": "chat.completion.chunk", "created": 0, "model": "test", "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
|
|
]
|
|
data = "".join("data: " + json.dumps(chunk, ensure_ascii=False) + "\n\n" for chunk in chunks) + "data: [DONE]\n\n"
|
|
return httpx.Response(200, content=data.encode(), headers={"content-type": "text/event-stream"})
|
|
return httpx.Response(200, json={"id": "recovered", "object": "chat.completion", "created": 0, "model": "test", "choices": [{"index": 0, "message": {"role": "assistant", "content": "正常回答"}, "finish_reason": "stop"}]})
|
|
|
|
model = PatchedChatOpenAI(
|
|
model="context-test",
|
|
api_key="synthetic-test-key",
|
|
base_url="https://context.invalid/v1",
|
|
max_retries=0,
|
|
http_client=httpx.Client(transport=httpx.MockTransport(transport)),
|
|
http_async_client=httpx.AsyncClient(transport=httpx.MockTransport(transport)),
|
|
)
|
|
injection = InjectOncePerCall()
|
|
errors = LLMErrorHandlingMiddleware(app_config=SimpleNamespace(circuit_breaker=SimpleNamespace(failure_threshold=3, recovery_timeout_sec=30)))
|
|
agent = create_agent(model=model, middleware=[errors, injection, ContextWindowGuardMiddleware()])
|
|
inputs = {"messages": [HumanMessage(content="请概括材料：" + "长材料" * 40000)]}
|
|
original = inputs["messages"][0].content
|
|
if streaming:
|
|
|
|
async def collect():
|
|
return [event async for event in agent.astream(inputs, stream_mode=["messages", "updates"])]
|
|
|
|
events = asyncio.run(collect()) if asynchronous else list(agent.stream(inputs, stream_mode=["messages", "updates"]))
|
|
serialized = str(events)
|
|
assert "正常回答" in serialized
|
|
assert "context_length_exceeded" not in serialized
|
|
assert "maximum context length" not in serialized
|
|
visible = [payload[0].content for mode, payload in events if mode == "messages" and getattr(payload[0], "content", "")]
|
|
assert "".join(visible) == "正常回答"
|
|
else:
|
|
result = asyncio.run(agent.ainvoke(inputs)) if asynchronous else agent.invoke(inputs)
|
|
assert result["messages"][-1].content == "正常回答"
|
|
assert result["messages"][0].content == original
|
|
assert 2 <= len(payloads) <= 5
|
|
assert injection.calls == 1
|
|
assert inputs["messages"][0].content == original
|
|
assert errors._circuit_failure_count == 0
|
|
|
|
|
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
|
def test_summary_internal_stream_does_not_leak_into_answer(asynchronous):
|
|
from langchain_core.language_models.fake_chat_models import FakeListChatModel
|
|
from langchain_core.messages import AIMessage
|
|
|
|
from deerflow.agents.middlewares.summarization_middleware import DeerFlowSummarizationMiddleware
|
|
|
|
summary = DeerFlowSummarizationMiddleware(FakeListChatModel(responses=["PRIVATE_INTERNAL_SUMMARY"]), trigger=("messages", 3), keep=("messages", 1))
|
|
agent = create_agent(model=FakeListChatModel(responses=["PUBLIC_ANSWER"]), middleware=[summary, ContextWindowGuardMiddleware()])
|
|
inputs = {"messages": [HumanMessage(content="old material" * 1000), AIMessage(content="old answer"), HumanMessage(content="continue")]}
|
|
|
|
async def collect():
|
|
return [event async for event in agent.astream(inputs, stream_mode="messages")]
|
|
|
|
events = asyncio.run(collect()) if asynchronous else list(agent.stream(inputs, stream_mode="messages"))
|
|
visible = "".join(message.content for message, metadata in events if isinstance(message.content, str))
|
|
assert visible == "PUBLIC_ANSWER"
|