96 lines
2.4 KiB
Python
96 lines
2.4 KiB
Python
"""
|
||
LLM 流式调用测试
|
||
"""
|
||
|
||
from src.Module.llm.llm import StreamingLLM, chat
|
||
|
||
|
||
def test_streaming_chat():
|
||
"""测试流式对话"""
|
||
print("=" * 60)
|
||
print(" LLM 流式对话测试")
|
||
print("=" * 60)
|
||
|
||
llm = StreamingLLM(
|
||
model='deepseek-v3.2',
|
||
enable_thinking=False, # 不开启思考模式
|
||
)
|
||
|
||
message = "用一句话介绍你自己"
|
||
print(f"\n用户: {message}\n")
|
||
print("助手: ", end="", flush=True)
|
||
|
||
full_response = ""
|
||
for chunk in llm.chat(message):
|
||
if chunk.error:
|
||
print(f"\n[错误] {chunk.error}")
|
||
return False
|
||
if chunk.content:
|
||
print(chunk.content, end="", flush=True)
|
||
full_response += chunk.content
|
||
|
||
print("\n")
|
||
print(f"完整响应长度: {len(full_response)} 字符")
|
||
|
||
return len(full_response) > 0
|
||
|
||
|
||
def test_thinking_mode():
|
||
"""测试思考模式"""
|
||
print("\n" + "=" * 60)
|
||
print("思考模式测试")
|
||
print("=" * 60)
|
||
|
||
llm = StreamingLLM(
|
||
model='deepseek-v3.2',
|
||
enable_thinking=True, # 开启思考模式
|
||
)
|
||
|
||
message = "1+1等于几？简单回答"
|
||
print(f"\n用户: {message}\n")
|
||
|
||
reasoning = ""
|
||
content = ""
|
||
is_answering = False
|
||
|
||
for chunk in llm.chat(message):
|
||
if chunk.error:
|
||
print(f"\n[错误] {chunk.error}")
|
||
return False
|
||
|
||
if chunk.reasoning_content:
|
||
if not is_answering:
|
||
if not reasoning:
|
||
print("思考过程: ", end="", flush=True)
|
||
print(chunk.reasoning_content, end="", flush=True)
|
||
reasoning += chunk.reasoning_content
|
||
|
||
if chunk.content:
|
||
if not is_answering:
|
||
print(f"\n\n回复: ", end="", flush=True)
|
||
is_answering = True
|
||
print(chunk.content, end="", flush=True)
|
||
content += chunk.content
|
||
|
||
print("\n")
|
||
print(f"思考过程: {len(reasoning)} 字符")
|
||
print(f"回复内容: {len(content)} 字符")
|
||
|
||
return len(content) > 0
|
||
|
||
|
||
if __name__ == '__main__':
|
||
results = []
|
||
|
||
success1 = test_streaming_chat()
|
||
results.append(("流式对话", success1))
|
||
|
||
success2 = test_thinking_mode()
|
||
results.append(("思考模式", success2))
|
||
|
||
print("\n" + "=" * 60)
|
||
print("测试结果:")
|
||
for name, success in results:
|
||
status = "✓ 通过" if success else "✗ 失败"
|
||
print(f" {name}: {status}")
|