118 lines
4.0 KiB
Python
118 lines
4.0 KiB
Python
"""tool.llm 自检：本地 Mock OpenAI 兼容服务，验证双模式 URL、重试、备用模型降级、embedding。
|
||
|
||
运行：python test_llm.py（纯 assert，无框架；可通过 exitcode 判断）
|
||
"""
|
||
import asyncio
|
||
import http.server
|
||
import json
|
||
import threading
|
||
|
||
from config.settings import settings
|
||
from tool.llm import LLMFailError, LLMClient
|
||
|
||
INDEX = json.dumps({"index": 0}).encode()
|
||
|
||
|
||
class MockLLMServer(http.server.BaseHTTPRequestHandler):
|
||
"""OpenAI 兼容 Mock：chat/embeddings/tags。model=boom:1b 返回 500 用于测降级。"""
|
||
|
||
def log_message(self, *args):
|
||
pass
|
||
|
||
def _reply(self, code: int, obj: dict):
|
||
body = json.dumps(obj).encode()
|
||
self.send_response(code)
|
||
self.send_header("Content-Type", "application/json")
|
||
self.send_header("Content-Length", str(len(body)))
|
||
self.end_headers()
|
||
self.wfile.write(body)
|
||
|
||
def do_GET(self): # ollama /api/tags 健康探测
|
||
self._reply(200, {"models": [{"name": "qwen2.5:7b"}]})
|
||
|
||
def do_POST(self):
|
||
n = int(self.headers.get("Content-Length", 0))
|
||
req = json.loads(self.rfile.read(n))
|
||
if self.path == "/v1/chat/completions":
|
||
model = req["model"]
|
||
if model.startswith("boom"):
|
||
self._reply(500, {"error": "model overloaded"})
|
||
return
|
||
self._reply(200, {"choices": [{"message": {"content": f"reply:{model}"}}]})
|
||
elif self.path == "/v1/embeddings":
|
||
texts = req["input"]
|
||
self._reply(200, {"data": [{"embedding": [float(i + 1)] * 4} for i in range(len(texts))]})
|
||
else:
|
||
self._reply(404, {"error": "not found"})
|
||
|
||
|
||
def start_server() -> tuple[threading.Thread, int]:
|
||
srv = http.server.ThreadingHTTPServer(("127.0.0.1", 0), MockLLMServer)
|
||
port = srv.server_address[1]
|
||
t = threading.Thread(target=srv.serve_forever, daemon=True)
|
||
t.start()
|
||
return t, port
|
||
|
||
|
||
def cfg(base_url: str, **kw):
|
||
"""从 .env 加载的 settings.llm 派生测试配置（只覆盖用例所需字段）。"""
|
||
kw.setdefault("mode", "ollama")
|
||
kw.setdefault("ollama_base", base_url)
|
||
kw.setdefault("max_retries", 2)
|
||
kw.setdefault("retry_backoff_sec", 0)
|
||
return settings.llm.model_copy(update=kw)
|
||
|
||
|
||
async def achat(c: LLMClient, text: str = "hi") -> str:
|
||
return await c.chat([{"role": "user", "content": text}])
|
||
|
||
|
||
def main():
|
||
_, port = start_server()
|
||
base = f"http://127.0.0.1:{port}"
|
||
|
||
# 1. ollama 模式 URL + 正常 chat
|
||
c = LLMClient(cfg(base, ollama_chat_model="ok:1b"))
|
||
assert c.is_ollama and c.base_url == f"{base}/v1", c.base_url
|
||
out = asyncio.run(achat(c))
|
||
assert out == "reply:ok:1b", out
|
||
|
||
# 2. 备用模型降级：主模型 500 → 自动切 fallback_chat_model
|
||
c2 = LLMClient(cfg(base, ollama_chat_model="boom:1b", ollama_embed_model="bge-m3",
|
||
fallback_chat_model="ok:1b"))
|
||
out = asyncio.run(achat(c2))
|
||
assert out == "reply:ok:1b", out
|
||
|
||
# 3. 全部失败 → LLMFailError（重试耗尽）
|
||
c3 = LLMClient(cfg(base, ollama_chat_model="boom:1b", max_retries=1))
|
||
try:
|
||
asyncio.run(achat(c3))
|
||
raise AssertionError("should have raised LLMFailError")
|
||
except LLMFailError:
|
||
pass
|
||
|
||
# 4. API 模式：端点/鉴权头/模型选择
|
||
c4 = LLMClient(settings.llm.model_copy(update={
|
||
"mode": "api", "api_base": base + "/v1", "api_key": "sk-test",
|
||
"api_chat_model": "gpt-x", "max_retries": 1,
|
||
}))
|
||
assert not c4.is_ollama and c4.base_url == f"{base}/v1"
|
||
assert c4._headers()["Authorization"] == "Bearer sk-test"
|
||
assert c4.chat_model == "gpt-x"
|
||
out = asyncio.run(achat(c4))
|
||
assert out == "reply:gpt-x", out
|
||
|
||
# 5. embedding 批量
|
||
c5 = LLMClient(cfg(base))
|
||
vecs = asyncio.run(c5.embed(["基金", "债券"]))
|
||
assert len(vecs) == 2 and len(vecs[0]) == 4, vecs
|
||
assert asyncio.run(c5.embed_one("货币")) == [1.0] * 4
|
||
|
||
# 6. 健康探测（ollama /api/tags）
|
||
asyncio.run(c5.check_health())
|
||
|
||
print("ALL LLM TESTS PASSED")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |