322 lines
12 KiB
Python
322 lines
12 KiB
Python
from __future__ import annotations
|
||
|
||
import io
|
||
import json
|
||
import unittest
|
||
from types import SimpleNamespace
|
||
from unittest import mock
|
||
|
||
import grok_direct_chat
|
||
|
||
|
||
class _Manager:
|
||
def __init__(self, *, backend: str = "dify"):
|
||
self.backend = backend
|
||
|
||
def load_ai_settings(self):
|
||
return {
|
||
"GROK_MODEL_ENABLED": True,
|
||
"GROK_API_KEY": "app-test-key",
|
||
"GROK_CUSTOMER_SERVICE_TIMEOUT": 180,
|
||
"GROK_DIFY_INPUTS": {"tenant": "test"},
|
||
}
|
||
|
||
def model_profile(self, _settings):
|
||
return SimpleNamespace(
|
||
compatible=True,
|
||
reason="",
|
||
model="self-model",
|
||
base_url="https://model.example/v1",
|
||
api_backend=self.backend,
|
||
auth_scheme="bearer",
|
||
temperature=0.2,
|
||
max_completion_tokens=512,
|
||
)
|
||
|
||
|
||
class _StreamResponse(io.BytesIO):
|
||
def __init__(self, events: list[object], *, content_type: str = "text/event-stream"):
|
||
body_parts: list[bytes] = []
|
||
for event in events:
|
||
if isinstance(event, bytes):
|
||
body_parts.append(event)
|
||
elif event == "[DONE]":
|
||
body_parts.append(b"data: [DONE]\n\n")
|
||
else:
|
||
body_parts.append(
|
||
(
|
||
"data: "
|
||
+ json.dumps(event, ensure_ascii=False)
|
||
+ "\n\n"
|
||
).encode("utf-8")
|
||
)
|
||
super().__init__(b"".join(body_parts))
|
||
self.status = 200
|
||
self.headers = {"Content-Type": content_type}
|
||
|
||
|
||
class ChatRouteTests(unittest.TestCase):
|
||
def test_normal_language_defaults_to_direct(self):
|
||
for text in (
|
||
"你好",
|
||
"你是什么模型",
|
||
"解释一下量子纠缠",
|
||
"帮我润色这段文字",
|
||
"?",
|
||
):
|
||
with self.subTest(text=text):
|
||
self.assertEqual(grok_direct_chat.classify_chat_route(text), "direct")
|
||
|
||
def test_live_or_executable_work_uses_agent(self):
|
||
for text in (
|
||
"帮我查询郑州天气",
|
||
"查天气",
|
||
"查一下最新新闻",
|
||
"打开这个网页 https://example.com",
|
||
"运行项目里的测试脚本",
|
||
"修改这个文件",
|
||
"给企业微信联系人张三发送消息",
|
||
"使用 Agent 调用 MCP 工具",
|
||
"@agent 处理这个任务",
|
||
):
|
||
with self.subTest(text=text):
|
||
self.assertEqual(grok_direct_chat.classify_chat_route(text), "agent")
|
||
|
||
def test_discussion_about_tools_does_not_start_agent(self):
|
||
self.assertEqual(
|
||
grok_direct_chat.classify_chat_route("MCP 是什么意思？"),
|
||
"direct",
|
||
)
|
||
self.assertEqual(
|
||
grok_direct_chat.classify_chat_route("解释一下这段代码"),
|
||
"direct",
|
||
)
|
||
|
||
def test_short_follow_up_inherits_agent_route(self):
|
||
self.assertEqual(
|
||
grok_direct_chat.classify_chat_route("明天呢？", last_route="agent"),
|
||
"agent",
|
||
)
|
||
self.assertEqual(
|
||
grok_direct_chat.classify_chat_route("？", last_route="direct"),
|
||
"direct",
|
||
)
|
||
|
||
def test_identity_reply_uses_actual_configured_model(self):
|
||
result = grok_direct_chat.direct_chat(
|
||
"你是哪个模型",
|
||
manager=_Manager(),
|
||
)
|
||
self.assertIn("self-model", result.text)
|
||
self.assertNotIn("xAI", result.text)
|
||
|
||
@mock.patch("grok_direct_chat._post_json")
|
||
def test_dify_direct_chat_reuses_conversation(self, post):
|
||
post.return_value = {
|
||
"answer": "直接回复",
|
||
"conversation_id": "dify-conv-1",
|
||
}
|
||
result = grok_direct_chat.direct_chat(
|
||
"继续说明",
|
||
history=[{"role": "assistant", "content": "上一轮"}],
|
||
conversation_id="dify-conv-old",
|
||
manager=_Manager(),
|
||
)
|
||
self.assertEqual(result.text, "直接回复")
|
||
self.assertEqual(result.conversation_id, "dify-conv-1")
|
||
args, kwargs = post.call_args
|
||
self.assertEqual(args[0], "https://model.example/v1/chat-messages")
|
||
self.assertEqual(kwargs["payload"]["conversation_id"], "dify-conv-old")
|
||
self.assertEqual(kwargs["payload"]["response_mode"], "blocking")
|
||
self.assertLessEqual(kwargs["timeout"], 90)
|
||
|
||
@mock.patch("grok_direct_chat._post_json")
|
||
def test_chat_completions_sends_bounded_history(self, post):
|
||
post.return_value = {
|
||
"choices": [{"message": {"content": "普通回复"}}]
|
||
}
|
||
result = grok_direct_chat.direct_chat(
|
||
"当前问题",
|
||
history=[{"role": "user", "content": "历史问题"}],
|
||
manager=_Manager(backend="chat_completions"),
|
||
)
|
||
self.assertEqual(result.text, "普通回复")
|
||
payload = post.call_args.kwargs["payload"]
|
||
self.assertEqual(payload["messages"][-1]["content"], "当前问题")
|
||
self.assertEqual(payload["messages"][-2]["content"], "历史问题")
|
||
self.assertFalse(payload["stream"])
|
||
|
||
|
||
class DirectChatStreamingTests(unittest.TestCase):
|
||
def _stream(self, backend: str, events: list[object]):
|
||
updates: list[tuple[str, bool]] = []
|
||
response = _StreamResponse(events)
|
||
with mock.patch.object(
|
||
grok_direct_chat._HTTP_OPENER,
|
||
"open",
|
||
return_value=response,
|
||
) as opened:
|
||
result = grok_direct_chat.stream_direct_chat(
|
||
"请回答",
|
||
on_update=lambda text, replace: updates.append((text, replace)),
|
||
manager=_Manager(backend=backend),
|
||
)
|
||
request = opened.call_args.args[0]
|
||
payload = json.loads(request.data.decode("utf-8"))
|
||
return result, updates, payload, request
|
||
|
||
def test_dify_streams_deltas_and_conversation_id(self):
|
||
result, updates, payload, request = self._stream(
|
||
"dify",
|
||
[
|
||
{"event": "message", "answer": "你", "conversation_id": "c1"},
|
||
{"event": "message", "answer": "好", "conversation_id": "c1"},
|
||
{"event": "message_end", "conversation_id": "c1"},
|
||
],
|
||
)
|
||
self.assertEqual(result.text, "你好")
|
||
self.assertEqual(result.conversation_id, "c1")
|
||
self.assertEqual(updates, [("你", False), ("好", False)])
|
||
self.assertEqual(payload["response_mode"], "streaming")
|
||
self.assertEqual(request.get_header("Accept"), "text/event-stream")
|
||
|
||
def test_dify_message_replace_replaces_full_text(self):
|
||
result, updates, _payload, _request = self._stream(
|
||
"dify",
|
||
[
|
||
{"event": "agent_message", "answer": "草稿"},
|
||
{"event": "message_replace", "answer": "最终答案"},
|
||
{"event": "message_end"},
|
||
],
|
||
)
|
||
self.assertEqual(result.text, "最终答案")
|
||
self.assertEqual(updates[-1], ("最终答案", True))
|
||
|
||
def test_chat_completions_streams_until_done(self):
|
||
result, updates, payload, _request = self._stream(
|
||
"chat_completions",
|
||
[
|
||
{"choices": [{"delta": {"role": "assistant"}, "finish_reason": None}]},
|
||
{"choices": [{"delta": {"content": "A"}, "finish_reason": None}]},
|
||
{"choices": [{"delta": {"content": "B"}, "finish_reason": "stop"}]},
|
||
"[DONE]",
|
||
],
|
||
)
|
||
self.assertEqual(result.text, "AB")
|
||
self.assertEqual(updates, [("A", False), ("B", False)])
|
||
self.assertTrue(payload["stream"])
|
||
|
||
def test_responses_streams_until_completed(self):
|
||
result, updates, payload, _request = self._stream(
|
||
"responses",
|
||
[
|
||
{"type": "response.output_text.delta", "delta": "甲"},
|
||
{"type": "response.output_text.delta", "delta": "乙"},
|
||
{"type": "response.completed", "response": {"status": "completed"}},
|
||
],
|
||
)
|
||
self.assertEqual(result.text, "甲乙")
|
||
self.assertEqual(updates, [("甲", False), ("乙", False)])
|
||
self.assertTrue(payload["stream"])
|
||
|
||
def test_anthropic_streams_text_until_message_stop(self):
|
||
result, updates, payload, request = self._stream(
|
||
"messages",
|
||
[
|
||
{"type": "message_start", "message": {"content": []}},
|
||
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}},
|
||
{"type": "content_block_stop", "index": 0},
|
||
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}},
|
||
{"type": "message_stop"},
|
||
],
|
||
)
|
||
self.assertEqual(result.text, "Hello")
|
||
self.assertEqual(updates, [("Hello", False)])
|
||
self.assertTrue(payload["stream"])
|
||
self.assertEqual(request.headers["Anthropic-version"], "2023-06-01")
|
||
|
||
def test_partial_eof_is_rejected(self):
|
||
response = _StreamResponse(
|
||
[{"type": "response.output_text.delta", "delta": "残缺"}]
|
||
)
|
||
with mock.patch.object(
|
||
grok_direct_chat._HTTP_OPENER,
|
||
"open",
|
||
return_value=response,
|
||
):
|
||
with self.assertRaisesRegex(
|
||
grok_direct_chat.DirectChatError,
|
||
"未正常结束",
|
||
):
|
||
grok_direct_chat.stream_direct_chat(
|
||
"请回答",
|
||
on_update=lambda _text, _replace: None,
|
||
manager=_Manager(backend="responses"),
|
||
)
|
||
|
||
def test_chat_completions_requires_finish_reason(self):
|
||
response = _StreamResponse(
|
||
[
|
||
{"choices": [{"delta": {"content": "文本"}, "finish_reason": None}]},
|
||
"[DONE]",
|
||
]
|
||
)
|
||
with mock.patch.object(
|
||
grok_direct_chat._HTTP_OPENER,
|
||
"open",
|
||
return_value=response,
|
||
):
|
||
with self.assertRaisesRegex(
|
||
grok_direct_chat.DirectChatError,
|
||
"finish_reason",
|
||
):
|
||
grok_direct_chat.stream_direct_chat(
|
||
"请回答",
|
||
on_update=lambda _text, _replace: None,
|
||
manager=_Manager(backend="chat_completions"),
|
||
)
|
||
|
||
def test_cancel_before_open(self):
|
||
cancellation = grok_direct_chat.DirectChatCancellation()
|
||
cancellation.cancel()
|
||
with mock.patch.object(grok_direct_chat._HTTP_OPENER, "open") as opened:
|
||
with self.assertRaises(grok_direct_chat.DirectChatCancelled):
|
||
grok_direct_chat.stream_direct_chat(
|
||
"请回答",
|
||
on_update=lambda _text, _replace: None,
|
||
manager=_Manager(backend="responses"),
|
||
cancellation=cancellation,
|
||
)
|
||
opened.assert_not_called()
|
||
|
||
def test_cancel_during_stream_closes_response(self):
|
||
cancellation = grok_direct_chat.DirectChatCancellation()
|
||
response = _StreamResponse(
|
||
[
|
||
{"type": "response.output_text.delta", "delta": "第一段"},
|
||
{"type": "response.completed", "response": {"status": "completed"}},
|
||
]
|
||
)
|
||
|
||
def cancel_after_first(_text, _replace):
|
||
cancellation.cancel()
|
||
|
||
with mock.patch.object(
|
||
grok_direct_chat._HTTP_OPENER,
|
||
"open",
|
||
return_value=response,
|
||
):
|
||
with self.assertRaises(grok_direct_chat.DirectChatCancelled):
|
||
grok_direct_chat.stream_direct_chat(
|
||
"请回答",
|
||
on_update=cancel_after_first,
|
||
manager=_Manager(backend="responses"),
|
||
cancellation=cancellation,
|
||
)
|
||
self.assertTrue(response.closed)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|