342 lines
12 KiB
Python
342 lines
12 KiB
Python
from __future__ import annotations
|
||
|
||
from genesis.inference.engine import InferenceEngine
|
||
from genesis.inference.exceptions import LLMError
|
||
from genesis.inference.prompt_registry import PromptRegistry
|
||
from genesis.inference.token import approximate_token_count
|
||
from genesis.inference.types import Prompt
|
||
from tests.inference_helpers import FakeLLMClient
|
||
|
||
|
||
class Models:
|
||
"""模拟 config.InferenceModels（pydantic 结构，测试中直接构建）"""
|
||
def __init__(self):
|
||
import types as t
|
||
self.primary = t.SimpleNamespace(name="deepseek-chat", provider="deepseek")
|
||
self.fallback = t.SimpleNamespace(name="qwen-max", provider="qwen")
|
||
|
||
|
||
class ModelsNoFallback:
|
||
"""无 fallback 的模型配置（覆盖 getattr(fallback) 假值分支）"""
|
||
def __init__(self):
|
||
import types as t
|
||
self.primary = t.SimpleNamespace(name="deepseek-chat", provider="deepseek")
|
||
self.fallback = None
|
||
|
||
|
||
class ModelsEmpty:
|
||
"""主/备用模型均为空（覆盖 names 为空 → 返回内置默认）"""
|
||
def __init__(self):
|
||
self.primary = None
|
||
self.fallback = None
|
||
|
||
|
||
def make_engine(client=None, *, constants=None):
|
||
eng = InferenceEngine(
|
||
client=client or FakeLLMClient(),
|
||
models=Models(),
|
||
registry=PromptRegistry(),
|
||
estimator=approximate_token_count,
|
||
)
|
||
if constants:
|
||
eng._max_context_tokens = constants # 内部测试钩子
|
||
return eng
|
||
|
||
|
||
# ---------- chat 主路径 ----------
|
||
|
||
def test_chat_ok():
|
||
client = FakeLLMClient([("ok", "正文")])
|
||
eng = make_engine(client)
|
||
r = eng.chat(
|
||
session_id="s1",
|
||
prompt=Prompt(name="writer", version="v1", template="章节：{{ chapter }}"),
|
||
variables={"chapter": "DB設計"},
|
||
)
|
||
assert r.status == "ok" and r.text == "正文"
|
||
assert client.calls[0]["model"] == "deepseek-chat"
|
||
|
||
|
||
def test_chat_fallback_after_primary_failure():
|
||
client = FakeLLMClient([("raise_timeout", ""), ("ok", "备用输出")])
|
||
eng = make_engine(client)
|
||
r = eng.chat(
|
||
session_id="s1",
|
||
prompt=Prompt(name="p", version="v1", template="t:{{ x }}"),
|
||
variables={"x": "1"},
|
||
)
|
||
assert r.status == "fallback"
|
||
# 先主模型，后备模型（修正：断言首次调用为主模型）
|
||
assert client.calls[0]["model"] == "deepseek-chat"
|
||
|
||
|
||
def test_chat_all_failed_returns_failed():
|
||
client = FakeLLMClient([("raise_network", ""), ("raise_network", "")])
|
||
eng = make_engine(client)
|
||
r = eng.chat(
|
||
session_id="s1",
|
||
prompt=Prompt(name="p", version="v1", template="t"),
|
||
variables={},
|
||
)
|
||
assert r.status == "failed" and r.error
|
||
|
||
|
||
# ---------- error_code 透传（与 error 同源取最后一次异常） ----------
|
||
|
||
def test_chat_failed_error_code_timeout():
|
||
# 主模型超时 + 备用也失败：error_code 与 error 同源（取最后一次异常）
|
||
client = FakeLLMClient([("raise_timeout", ""), ("raise_timeout", "")])
|
||
eng = InferenceEngine(client=client, models=Models())
|
||
r = eng.chat(session_id="s1", prompt=Prompt(name="p", version="v1", template="t"), variables={})
|
||
assert r.status == "failed"
|
||
assert r.error_code == "LLM_TIMEOUT"
|
||
|
||
|
||
def test_chat_failed_error_code_network_last():
|
||
# 主模型超时（第一次）、备用网络失败（最后一次）→ error_code 取最后一次 = LLM_NETWORK_ERROR
|
||
client = FakeLLMClient([("raise_timeout", ""), ("raise_network", "")])
|
||
eng = InferenceEngine(client=client, models=Models())
|
||
r = eng.chat(session_id="s1", prompt=Prompt(name="p", version="v1", template="t"), variables={})
|
||
assert r.status == "failed"
|
||
assert r.error_code == "LLM_NETWORK_ERROR"
|
||
|
||
|
||
def test_chat_failed_error_code_not_configured():
|
||
# 备用模型 Key 缺失（最后一次）→ LLM_NOT_CONFIGURED
|
||
# 用自定义异常客户端模拟 NotConfigured
|
||
class NotConfiguredClient:
|
||
def __init__(self):
|
||
self.calls = []
|
||
def chat(self, *, model, messages, temperature, max_tokens):
|
||
self.calls.append(model)
|
||
from genesis.inference.exceptions import LLMNotConfiguredError
|
||
raise LLMNotConfiguredError("no key")
|
||
eng = InferenceEngine(client=NotConfiguredClient(), models=Models())
|
||
r = eng.chat(session_id="s1", prompt=Prompt(name="p", version="v1", template="t"), variables={})
|
||
assert r.status == "failed"
|
||
assert r.error_code == "LLM_NOT_CONFIGURED"
|
||
|
||
|
||
def test_chat_structured_parse_error_code():
|
||
client = FakeLLMClient([("parse_fail", ""), ("parse_fail", "")])
|
||
eng = InferenceEngine(client=client, models=Models())
|
||
r = eng.chat_structured(
|
||
session_id="s1", prompt=Prompt(name="p", version="v1", template="提取"),
|
||
variables={}, schema={}, retry_count=1,
|
||
)
|
||
assert r.status == "parse_error"
|
||
assert r.error_code == "LLM_PARSE_ERROR"
|
||
|
||
|
||
def test_chat_structured_failed_error_code_network():
|
||
client = FakeLLMClient([("raise_network", "")])
|
||
eng = InferenceEngine(client=client, models=Models())
|
||
r = eng.chat_structured(
|
||
session_id="s1", prompt=Prompt(name="p", version="v1", template="提取"),
|
||
variables={}, schema={},
|
||
)
|
||
assert r.status == "failed"
|
||
assert r.error_code == "LLM_NETWORK_ERROR"
|
||
|
||
|
||
def test_chat_plain_string_prompt():
|
||
eng = make_engine(FakeLLMClient([("ok", "hi")]))
|
||
r = eng.chat(session_id="s1", prompt="直接文本", variables={})
|
||
assert r.text == "hi" and r.status == "ok"
|
||
|
||
|
||
def test_chat_truncation_callback_triggered():
|
||
seen = {}
|
||
|
||
def truncate_cb(prompt_text, variables):
|
||
seen["called"] = True
|
||
seen["len"] = len(prompt_text)
|
||
return {**variables, "chapter": "裁剪版"}
|
||
|
||
eng = InferenceEngine(
|
||
client=FakeLLMClient([("ok", "x")]),
|
||
models=Models(),
|
||
registry=PromptRegistry(),
|
||
estimator=approximate_token_count,
|
||
truncate_cb=truncate_cb,
|
||
)
|
||
eng._max_context_tokens = 2 # 强制超限（复习：'abcd很长很长的标题' 约 3 token > 2）
|
||
r = eng.chat(
|
||
session_id="s1",
|
||
prompt=Prompt(name="p", version="v1", template="abcd{{ chapter }}"),
|
||
variables={"chapter": "很长很长的标题"},
|
||
)
|
||
assert seen["called"] is True
|
||
assert r.status == "ok"
|
||
|
||
|
||
# ---------- chat_structured 主路径 ----------
|
||
|
||
def test_chat_structured_ok():
|
||
client = FakeLLMClient([("ok", '{"a": 1}')])
|
||
eng = make_engine(client)
|
||
r = eng.chat_structured(
|
||
session_id="s1", prompt=Prompt(name="p", version="v1", template="提取"),
|
||
variables={"text": "内容"},
|
||
schema={"type": "object", "properties": {"a": {"type": "number"}}},
|
||
)
|
||
assert r.status == "ok" and r.data == {"a": 1}
|
||
|
||
|
||
def test_chat_structured_retry_parse():
|
||
client = FakeLLMClient([("parse_fail", ""), ("ok", '{"a": 2}')])
|
||
eng = make_engine(client)
|
||
r = eng.chat_structured(
|
||
session_id="s1", prompt=Prompt(name="p", version="v1", template="提取"),
|
||
variables={}, schema={},
|
||
)
|
||
assert r.status == "ok" and r.data == {"a": 2} and r.parse_attempts == 2
|
||
|
||
|
||
def test_chat_structured_parse_error_returns_raw():
|
||
client = FakeLLMClient([("parse_fail", ""), ("parse_fail", "")])
|
||
eng = make_engine(client)
|
||
r = eng.chat_structured(
|
||
session_id="s1", prompt=Prompt(name="p", version="v1", template="提取"),
|
||
variables={}, schema={}, retry_count=1,
|
||
)
|
||
assert r.status == "parse_error"
|
||
assert r.raw_text == "NOT JSON"
|
||
assert r.parse_attempts == 2
|
||
|
||
|
||
def test_chat_structured_failed_on_network():
|
||
client = FakeLLMClient([("raise_network", "")])
|
||
eng = make_engine(client)
|
||
r = eng.chat_structured(
|
||
session_id="s1", prompt=Prompt(name="p", version="v1", template="提取"),
|
||
variables={}, schema={},
|
||
)
|
||
assert r.status == "failed"
|
||
|
||
|
||
def test_chat_structured_truncation_callback_triggered():
|
||
# chat_structured 超限时同样触发 truncate_cb（与 chat 流程一致）
|
||
seen = {}
|
||
|
||
def truncate_cb(prompt_text, variables):
|
||
seen["called"] = True
|
||
return {**variables, "chapter": "裁剪版"}
|
||
|
||
eng = InferenceEngine(
|
||
client=FakeLLMClient([("ok", '{"a": 1}')]),
|
||
models=Models(),
|
||
registry=PromptRegistry(),
|
||
estimator=approximate_token_count,
|
||
truncate_cb=truncate_cb,
|
||
)
|
||
eng._max_context_tokens = 2 # 强制超限（渲染约 3 token > 2）
|
||
r = eng.chat_structured(
|
||
session_id="s1",
|
||
prompt=Prompt(name="p", version="v1", template="abcd{{ chapter }}"),
|
||
variables={"chapter": "很长很长的标题"},
|
||
schema={"type": "object", "properties": {"a": {"type": "number"}}},
|
||
)
|
||
assert seen["called"] is True
|
||
assert r.status == "ok"
|
||
|
||
|
||
# ---------- 补充分支覆盖（defensive / 缺失配置） ----------
|
||
|
||
def test_chat_fallback_all_failed_when_only_primary():
|
||
"""models 无 fallback：仅尝试主模型，失败后直接 failed。"""
|
||
client = FakeLLMClient([("raise_network", "")])
|
||
eng = InferenceEngine(
|
||
client=client,
|
||
models=ModelsNoFallback(),
|
||
registry=PromptRegistry(),
|
||
estimator=approximate_token_count,
|
||
)
|
||
r = eng.chat(session_id="s1", prompt="t", variables={})
|
||
assert r.status == "failed"
|
||
assert len(client.calls) == 1
|
||
assert client.calls[0]["model"] == "deepseek-chat"
|
||
|
||
|
||
def test_chat_empty_models_uses_default():
|
||
"""主/备用均为空配置时回退内置默认模型 deepseek-chat。"""
|
||
eng = InferenceEngine(
|
||
client=FakeLLMClient([("ok", "x")]),
|
||
models=ModelsEmpty(),
|
||
registry=PromptRegistry(),
|
||
estimator=approximate_token_count,
|
||
)
|
||
r = eng.chat(session_id="s1", prompt="t", variables={})
|
||
assert r.status == "ok"
|
||
assert eng._model_names(None) == ["deepseek-chat"]
|
||
|
||
|
||
def test_chat_defaults_without_registry_estimator():
|
||
"""registry/estimator/models 均未配置时使用内置默认（PromptRegistry + approximate + deepseek-chat）。"""
|
||
client = FakeLLMClient([("ok", "默认")])
|
||
eng = InferenceEngine(client=client, models=None)
|
||
r = eng.chat(
|
||
session_id="s1",
|
||
prompt=Prompt(name="p", version="v1", template="tt"),
|
||
variables={},
|
||
)
|
||
assert r.status == "ok" and r.text == "默认"
|
||
assert eng._model_names(None) == ["deepseek-chat"]
|
||
|
||
|
||
def test_chat_explicit_model_param():
|
||
"""显式指定 model，跳过主/备选择，仅调用该模型。"""
|
||
client = FakeLLMClient([("ok", "c")])
|
||
eng = make_engine(client)
|
||
r = eng.chat(session_id="s1", prompt="t", variables={}, model="qwen-custom")
|
||
assert r.status == "ok"
|
||
assert client.calls[0]["model"] == "qwen-custom"
|
||
assert len(client.calls) == 1
|
||
|
||
|
||
def test_chat_truncation_without_callback_keeps_variables():
|
||
"""超限但无 truncate_cb，variables 原样保留（裁剪分支不触发）。"""
|
||
client = FakeLLMClient([("ok", "ok")])
|
||
eng = make_engine(client)
|
||
eng._max_context_tokens = 2 # 强制超限（渲染后约 11 token）
|
||
r = eng.chat(
|
||
session_id="s1",
|
||
prompt=Prompt(name="p", version="v1", template="abcd{{ chapter }}"),
|
||
variables={"chapter": "很长很长的标题"},
|
||
)
|
||
assert r.status == "ok"
|
||
assert client.calls[0]["model"] == "deepseek-chat"
|
||
|
||
|
||
def test_chat_structured_empty_schema_no_hint():
|
||
"""schema 为空时不追加 schema 提示，仍可正常解析。"""
|
||
client = FakeLLMClient([("ok", '{"v": true}')])
|
||
eng = make_engine(client)
|
||
r = eng.chat_structured(
|
||
session_id="s1", prompt="直接文本", variables={}, schema={},
|
||
)
|
||
assert r.status == "ok" and r.data == {"v": True}
|
||
|
||
|
||
def test_chat_truncation_callback_returns_none_keeps_variables():
|
||
"""truncate_cb 返回 None 时回退原 variables（覆盖 new_vars is None 分支）。"""
|
||
def truncate_cb(prompt_text, variables):
|
||
return None
|
||
|
||
client = FakeLLMClient([("ok", "x")])
|
||
eng = InferenceEngine(
|
||
client=client,
|
||
models=Models(),
|
||
registry=PromptRegistry(),
|
||
estimator=approximate_token_count,
|
||
truncate_cb=truncate_cb,
|
||
)
|
||
eng._max_context_tokens = 2
|
||
r = eng.chat(
|
||
session_id="s1",
|
||
prompt=Prompt(name="p", version="v1", template="abcd{{ chapter }}"),
|
||
variables={"chapter": "很长很长的标题"},
|
||
)
|
||
assert r.status == "ok"
|
||
# 未替换变量：发送内容仍为原样渲染
|
||
assert "很长很长的标题" in client.calls[0]["messages"][0] |