34 lines
1.1 KiB
Python
34 lines
1.1 KiB
Python
#!/usr/bin/env python3
|
|
"""ai_brain 网关多 URL 回退 — 离线回归"""
|
|
import importlib.util
|
|
import os
|
|
|
|
AI_BRAIN_PATH = os.path.join(
|
|
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
|
|
"app", "agent", "ai_brain.py",
|
|
)
|
|
|
|
|
|
def _load_brain():
|
|
spec = importlib.util.spec_from_file_location("wp_ai_brain_ut", AI_BRAIN_PATH)
|
|
mod = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(mod)
|
|
return mod.AIBrain
|
|
|
|
|
|
def test_gateway_post_urls_primary_and_fallback():
|
|
Brain = _load_brain()
|
|
b = Brain(ai_api_url="http://127.0.0.1:3102", ai_api_key="k")
|
|
urls = b._gateway_post_urls()
|
|
assert urls[0] == "http://127.0.0.1:3102/api/gateway/chat"
|
|
assert "http://127.0.0.1:3102/v1/chat/completions" in urls
|
|
assert "http://127.0.0.1:18080/v1/chat/completions" in urls
|
|
|
|
|
|
def test_extract_gateway_content_openai_and_reply():
|
|
Brain = _load_brain()
|
|
oai = {"choices": [{"message": {"content": "OK"}}]}
|
|
assert Brain._extract_gateway_content(oai) == "OK"
|
|
reply = {"reply": "你好"}
|
|
assert Brain._extract_gateway_content(reply) == "你好"
|