278 lines
14 KiB
Python
278 lines
14 KiB
Python
"""Model-specific administrator instructions: persistence, authorization and calls."""
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
from langchain_core.language_models.fake_chat_models import FakeListChatModel
|
|
from langchain_core.messages import HumanMessage, SystemMessage
|
|
from pydantic import Field
|
|
|
|
from app.gateway.routers import system_settings as router
|
|
from deerflow.config import system_settings as ss
|
|
from deerflow.models.prompt_suffix import append_model_prompt_suffix, with_model_prompt_suffix
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def runtime(monkeypatch, tmp_path):
|
|
monkeypatch.setattr(ss, "runtime_home", lambda: tmp_path)
|
|
|
|
|
|
def save(prompt="请直接回答问题", enabled=True):
|
|
settings = ss.load_system_settings()
|
|
settings.model_prompt_suffixes["provider/model-a"] = ss.ModelPromptSuffix(enabled=enabled, prompt=prompt)
|
|
ss.save_system_settings(settings)
|
|
|
|
|
|
def test_matching_order_and_no_state_mutation():
|
|
save()
|
|
messages = [SystemMessage(content="基础要求"), HumanMessage(content="问题"), SystemMessage(content="补充上下文")]
|
|
assert append_model_prompt_suffix(messages, "model-b") is messages
|
|
result = append_model_prompt_suffix(messages, "provider/model-a")
|
|
assert result[0].content == "基础要求\n\n---\n\n补充上下文\n\n请直接回答问题"
|
|
assert result[1:] == [messages[1]]
|
|
assert messages[0].content == "基础要求"
|
|
assert len(messages) == 3
|
|
assert append_model_prompt_suffix(result, "provider/model-a") == result
|
|
|
|
|
|
def test_blocks_no_system_and_disabled():
|
|
save()
|
|
original = SystemMessage(content=[{"type": "text", "text": "原始要求"}])
|
|
result = append_model_prompt_suffix([original], "provider/model-a")
|
|
assert result[0].content[-1]["text"] == "\n\n请直接回答问题"
|
|
assert len(original.content) == 1
|
|
assert append_model_prompt_suffix(result, "provider/model-a") == result
|
|
assert append_model_prompt_suffix([HumanMessage(content="问题")], "provider/model-a")[0].content == "请直接回答问题"
|
|
save(enabled=False)
|
|
assert append_model_prompt_suffix([original], "provider/model-a") == [original]
|
|
save(prompt=" ")
|
|
assert append_model_prompt_suffix([original], "provider/model-a") == [original]
|
|
|
|
|
|
class RecordingModel(FakeListChatModel):
|
|
seen: list = Field(default_factory=list)
|
|
|
|
def _generate(self, messages, *args, **kwargs):
|
|
self.seen.append(messages)
|
|
return super()._generate(messages, *args, **kwargs)
|
|
|
|
def _stream(self, messages, *args, **kwargs):
|
|
self.seen.append(messages)
|
|
yield from super()._stream(messages, *args, **kwargs)
|
|
|
|
async def _astream(self, messages, *args, **kwargs):
|
|
self.seen.append(messages)
|
|
async for chunk in super()._astream(messages, *args, **kwargs):
|
|
yield chunk
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_all_invocation_paths_and_live_updates():
|
|
model = with_model_prompt_suffix(RecordingModel)(responses=["ok"], deerflow_prompt_model_name="provider/model-a")
|
|
messages = [SystemMessage(content="base"), HumanMessage(content="question")]
|
|
save()
|
|
model.invoke(messages)
|
|
await model.ainvoke(messages)
|
|
list(model.stream(messages))
|
|
_ = [chunk async for chunk in model.astream(messages)]
|
|
model.generate([messages])
|
|
await model.agenerate([messages])
|
|
model.bind(temperature=0).invoke(messages)
|
|
assert len(model.seen) == 7
|
|
assert all(batch[0].content == "base\n\n请直接回答问题" for batch in model.seen)
|
|
save(prompt="新的要求")
|
|
model.invoke(messages)
|
|
assert model.seen[-1][0].content == "base\n\n新的要求"
|
|
save(enabled=False)
|
|
model.invoke(messages)
|
|
assert model.seen[-1][0].content == "base"
|
|
assert messages[0].content == "base"
|
|
assert "deerflow_prompt_model_name" not in model.model_dump()
|
|
|
|
|
|
def test_per_model_instance_switch_skips_suffix():
|
|
save()
|
|
model = with_model_prompt_suffix(RecordingModel)(
|
|
responses=["ok"],
|
|
deerflow_prompt_model_name="provider/model-a",
|
|
deerflow_prompt_suffix_enabled=False,
|
|
)
|
|
|
|
model.invoke([SystemMessage(content="base"), HumanMessage(content="question")])
|
|
|
|
assert model.seen[-1][0].content == "base"
|
|
assert "deerflow_prompt_suffix_enabled" not in model.model_dump()
|
|
|
|
|
|
def test_factory_preserves_provider_and_tool_binding(monkeypatch):
|
|
from deerflow.config.app_config import AppConfig
|
|
from deerflow.config.model_config import ModelConfig
|
|
from deerflow.models.factory import create_chat_model
|
|
from deerflow.models.patched_openai import PatchedChatOpenAI
|
|
|
|
config = AppConfig(sandbox={"use": "deerflow.sandbox.local:LocalSandboxProvider"}, models=[ModelConfig(name="provider/model-a", model="test", use="langchain_openai:ChatOpenAI", api_key="test", base_url="http://localhost:9999/v1")])
|
|
model = create_chat_model(app_config=config)
|
|
assert isinstance(model, PatchedChatOpenAI)
|
|
save()
|
|
bound = model.bind_tools([{"type": "function", "function": {"name": "test", "parameters": {"type": "object", "properties": {}}}}])
|
|
payload = bound.bound._get_request_payload([HumanMessage(content="hello")], **bound.kwargs)
|
|
assert payload["messages"][0]["content"] == "请直接回答问题"
|
|
assert payload["tools"][0]["function"]["name"] == "test"
|
|
assert "deerflow_prompt_model_name" not in payload
|
|
responses_model = create_chat_model(app_config=config, use_responses_api=True)
|
|
responses_payload = responses_model._get_request_payload([SystemMessage(content="base"), HumanMessage(content="hello")])
|
|
system_content = responses_payload["input"][0]["content"]
|
|
assert "请直接回答问题" in str(system_content)
|
|
|
|
|
|
def test_admin_api_validation_and_persistence(monkeypatch):
|
|
app = FastAPI()
|
|
app.include_router(router.router)
|
|
client = TestClient(app)
|
|
user = SimpleNamespace(system_role="user")
|
|
|
|
async def current_user(_request):
|
|
return user
|
|
|
|
monkeypatch.setattr(router, "get_optional_user_from_request", current_user)
|
|
monkeypatch.setattr(router, "get_app_config", lambda: SimpleNamespace(get_model_config=lambda name: object() if name == "provider/model-a" else None))
|
|
endpoint = "/api/system-settings/model-prompt-suffixes"
|
|
url = endpoint + "/provider%2Fmodel-a"
|
|
assert client.get(endpoint).status_code == 403
|
|
assert client.put(url, json={"enabled": True, "prompt": "要求"}).status_code == 403
|
|
user.system_role = "admin"
|
|
assert client.put(url, json={"enabled": True, "prompt": " "}).status_code == 422
|
|
assert client.put(url, json={"enabled": True, "prompt": "x" * 20001}).status_code == 422
|
|
assert client.put(endpoint + "/unknown", json={"enabled": True, "prompt": "要求"}).status_code == 404
|
|
settings = ss.load_system_settings()
|
|
settings.conversation_context.max_conversation_rounds = 12
|
|
ss.save_system_settings(settings)
|
|
assert client.put(url, json={"enabled": True, "prompt": "要求"}).status_code == 200
|
|
assert client.get(endpoint).json()["provider/model-a"] == {"enabled": True, "prompt": "要求"}
|
|
assert ss.load_system_settings().conversation_context.max_conversation_rounds == 12
|
|
user = None
|
|
assert client.get(endpoint).status_code == 401
|
|
assert client.put(url, json={"enabled": False, "prompt": ""}).status_code == 401
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("custom_agent", [False, True], ids=["ordinary-chat", "agent-chat"])
|
|
@pytest.mark.parametrize("agent_prompt_suffix_enabled", [True, False], ids=["agent-suffix-on", "agent-suffix-off"])
|
|
@pytest.mark.parametrize("selected", [None, "provider/model-a", "model-b"], ids=["default-model", "selected-target", "selected-other"])
|
|
@pytest.mark.parametrize("streaming", [False, True], ids=["invoke", "stream"])
|
|
async def test_chat_graph_final_http_payload(monkeypatch, custom_agent, agent_prompt_suffix_enabled, selected, streaming):
|
|
"""Run both lead-agent branches through LangGraph to the actual HTTP boundary.
|
|
|
|
Replace external skill/sandbox middleware with a late context injector, and
|
|
use distinguishable base prompt fixtures. Model selection, graph execution,
|
|
factory, provider serialization and HTTP transport remain real.
|
|
"""
|
|
import json
|
|
|
|
import httpx
|
|
from langchain.agents.middleware import AgentMiddleware
|
|
|
|
from deerflow.agents.lead_agent import agent as lead
|
|
from deerflow.agents.lead_agent import prompt as prompts
|
|
from deerflow.config.app_config import AppConfig
|
|
from deerflow.config.model_config import ModelConfig
|
|
from deerflow.models.factory import create_chat_model
|
|
|
|
save()
|
|
captured = []
|
|
|
|
def respond(request):
|
|
payload = json.loads(request.content)
|
|
captured.append(payload)
|
|
if payload.get("stream"):
|
|
chunks = [
|
|
{"id": "chat-test", "object": "chat.completion.chunk", "model": "test", "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": None}]},
|
|
{"id": "chat-test", "object": "chat.completion.chunk", "model": "test", "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
|
|
]
|
|
content = "".join("data: " + json.dumps(chunk) + "\n\n" for chunk in chunks) + "data: [DONE]\n\n"
|
|
return httpx.Response(200, headers={"Content-Type": "text/event-stream"}, content=content)
|
|
return httpx.Response(200, json={"id": "chat-test", "object": "chat.completion", "model": "test", "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}]})
|
|
|
|
class LateContext(AgentMiddleware):
|
|
async def awrap_model_call(self, request, handler):
|
|
return await handler(
|
|
request.override(
|
|
system_message=SystemMessage(content=request.system_message.content + "\n最后一个中间件的要求"),
|
|
messages=[*request.messages, SystemMessage(content="运行时检索上下文")],
|
|
)
|
|
)
|
|
|
|
config = AppConfig(
|
|
sandbox={"use": "deerflow.sandbox.local:LocalSandboxProvider"},
|
|
models=[ModelConfig(name=name, model="test", use="langchain_openai:ChatOpenAI", api_key="test", base_url="http://model.test/v1") for name in ["model-b", "provider/model-a"]],
|
|
)
|
|
agent_config = SimpleNamespace(name="测试智能体", model="provider/model-a", model_prompt_suffix_enabled=agent_prompt_suffix_enabled, skills=[], tool_groups=None) if custom_agent else None
|
|
monkeypatch.setattr(lead, "load_agent_config", lambda _: agent_config)
|
|
monkeypatch.setattr(lead, "load_agent_soul", lambda _: "SOUL.md 的要求" if custom_agent else "")
|
|
monkeypatch.setattr(lead, "apply_prompt_template", lambda **_: "普通问答模板")
|
|
monkeypatch.setattr(prompts, "apply_agent_only_prompt_template", lambda **_: "SOUL.md 的要求")
|
|
monkeypatch.setattr(lead, "_build_middlewares", lambda *args, **kwargs: [LateContext()])
|
|
runtime = {"disable_tools": True, "thinking_enabled": False, "agent_id": "test-agent" if custom_agent else "lead"}
|
|
if selected:
|
|
runtime["model_name"] = selected
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
|
monkeypatch.setattr(lead, "create_chat_model", lambda **kwargs: create_chat_model(**kwargs, http_async_client=client))
|
|
graph = lead._make_lead_agent({"configurable": runtime}, app_config=config)
|
|
inputs = {"messages": [HumanMessage(content="用户问题")]}
|
|
if streaming:
|
|
_ = [chunk async for chunk in graph.astream(inputs, stream_mode="messages")]
|
|
else:
|
|
await graph.ainvoke(inputs)
|
|
|
|
assert len(captured) == 1
|
|
messages = captured[0]["messages"]
|
|
systems = [message for message in messages if message["role"] in {"system", "developer"}]
|
|
assert len(systems) == 1
|
|
text = systems[0]["content"]
|
|
assert ("SOUL.md 的要求" if custom_agent else "普通问答模板") in text
|
|
assert "最后一个中间件的要求" in text
|
|
assert "运行时检索上下文" in text
|
|
effective_model = selected or ("provider/model-a" if custom_agent else "model-b")
|
|
if effective_model == "provider/model-a" and (not custom_agent or agent_prompt_suffix_enabled):
|
|
assert text.endswith("运行时检索上下文\n\n请直接回答问题")
|
|
assert text.count("请直接回答问题") == 1
|
|
else:
|
|
assert "请直接回答问题" not in text
|
|
assert [message["content"] for message in messages if message["role"] == "user"] == ["用户问题"]
|
|
assert inputs["messages"][0].content == "用户问题"
|
|
|
|
|
|
@pytest.mark.parametrize("transport", ["text_bridge", "buffered_gateway", "artifact_renderer"])
|
|
def test_suffix_follows_file_transport_instructions(transport):
|
|
from deerflow.models import patched_openai as provider
|
|
|
|
save()
|
|
model = with_model_prompt_suffix(provider.PatchedChatOpenAI)(
|
|
model="test",
|
|
api_key="test",
|
|
deerflow_prompt_model_name="provider/model-a",
|
|
)
|
|
tool = {"type": "function", "function": {"name": "write_file", "description": "Write a file", "parameters": {"type": "object", "properties": {"path": {"type": "string"}, "content": {"type": "string"}}, "required": ["path", "content"]}}}
|
|
tool["function"]["parameters"]["properties"].update({"description": {"type": "string"}, "append": {"type": "boolean"}})
|
|
begin_tool = {
|
|
"type": "function",
|
|
"function": {"name": "begin_write_file", "description": "Prepare a file", "parameters": {"type": "object", "properties": {"path": {"type": "string"}, "description": {"type": "string"}, "append": {"type": "boolean"}}}},
|
|
}
|
|
token = provider._ACTIVE_WRITE_FILE_BRIDGE.set(provider._WriteFileBridgeStreamState(transport=transport))
|
|
try:
|
|
payload = model._get_request_payload([SystemMessage(content="SOUL 要求"), HumanMessage(content="写文件")], tools=[tool, begin_tool])
|
|
finally:
|
|
provider._ACTIVE_WRITE_FILE_BRIDGE.reset(token)
|
|
system = payload["messages"][0]["content"]
|
|
instruction = {
|
|
"text_bridge": provider._TEXT_BRIDGE_SELECTION_INSTRUCTION,
|
|
"buffered_gateway": provider._BUFFERED_GATEWAY_SELECTION_INSTRUCTION,
|
|
"artifact_renderer": provider._ARTIFACT_RENDERER_SELECTION_INSTRUCTION,
|
|
}[transport]
|
|
assert instruction in system
|
|
assert system.startswith("SOUL 要求")
|
|
assert system.endswith("\n\n请直接回答问题")
|
|
assert system.count("请直接回答问题") == 1
|