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()
|