feat: InferenceEngine chat/chat_structured 全流程 + 文档补丁（补丁1/2/3）
This commit is contained in:
@@ -0,0 +1,185 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from typing import Any, Callable
|
||||
|
||||
from .client import LLMClient
|
||||
from .exceptions import LLMError
|
||||
from .prompt_registry import PromptRegistry
|
||||
from .token import make_estimator
|
||||
from .types import (
|
||||
ChatMessage,
|
||||
ChatResult,
|
||||
Prompt,
|
||||
StructuredResult,
|
||||
TokenUsage,
|
||||
)
|
||||
|
||||
|
||||
class InferenceEngine:
|
||||
"""统一 LLM 调用入口：模型选择/降级、重试、解析、Token 超限回调。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
client: LLMClient,
|
||||
models: Any | None = None,
|
||||
registry: PromptRegistry | None = None,
|
||||
estimator: Callable[[str], int] | None = None,
|
||||
truncate_cb: Callable[[str, dict], dict] | None = None,
|
||||
max_context_tokens: int = 32000,
|
||||
) -> None:
|
||||
self._client = client
|
||||
self._models = models
|
||||
self._registry = registry or PromptRegistry()
|
||||
self._estimator = estimator or make_estimator()
|
||||
self._truncate_cb = truncate_cb
|
||||
self._max_context_tokens = max_context_tokens
|
||||
|
||||
# ---------- 内部 ----------
|
||||
|
||||
def _render_prompt(self, prompt: Prompt | str, variables: dict) -> str:
|
||||
if isinstance(prompt, Prompt):
|
||||
return self._registry.render(prompt.template, variables) if variables else prompt.template
|
||||
return prompt
|
||||
|
||||
def _apply_truncation(self, text: str, variables: dict) -> dict:
|
||||
"""Token 超限时触发裁剪回调（注入），返回新 variables。"""
|
||||
if self._truncate_cb is not None:
|
||||
new_vars = self._truncate_cb(text, variables)
|
||||
if new_vars is not None:
|
||||
return new_vars
|
||||
return variables
|
||||
|
||||
def _model_names(self, model: str | None) -> list[str]:
|
||||
"""返回尝试顺序；显式指定 model 时只用它，否则 primary→fallback。"""
|
||||
if model:
|
||||
return [model]
|
||||
if self._models:
|
||||
names = []
|
||||
if getattr(self._models, "primary", None):
|
||||
names.append(self._models.primary.name)
|
||||
if getattr(self._models, "fallback", None):
|
||||
names.append(self._models.fallback.name)
|
||||
if names:
|
||||
return names
|
||||
return ["deepseek-chat"]
|
||||
|
||||
def _call(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
rendered: str,
|
||||
temperature: float,
|
||||
max_tokens: int,
|
||||
) -> tuple[str, TokenUsage]:
|
||||
messages = [ChatMessage(role="user", content=rendered)]
|
||||
return self._client.chat(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
|
||||
# ---------- 公开 ----------
|
||||
|
||||
def chat(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
prompt: Prompt | str,
|
||||
variables: dict,
|
||||
model: str | None = None,
|
||||
temperature: float = 0.2,
|
||||
max_tokens: int = 4096,
|
||||
) -> ChatResult:
|
||||
rendered = self._render_prompt(prompt, variables)
|
||||
if self._estimator(rendered) > self._max_context_tokens:
|
||||
variables = self._apply_truncation(rendered, variables)
|
||||
rendered = self._render_prompt(prompt, variables)
|
||||
|
||||
start = time.monotonic()
|
||||
last_error: str | None = None
|
||||
for idx, name in enumerate(self._model_names(model)):
|
||||
try:
|
||||
text, usage = self._call(
|
||||
model=name, rendered=rendered,
|
||||
temperature=temperature, max_tokens=max_tokens,
|
||||
)
|
||||
status = "ok" if idx == 0 else "fallback"
|
||||
return ChatResult(
|
||||
text=text, model=name, prompt_version=getattr(prompt, "version", "inline"),
|
||||
usage=usage, duration_ms=int((time.monotonic() - start) * 1000),
|
||||
status=status,
|
||||
)
|
||||
except LLMError as exc:
|
||||
last_error = str(exc)
|
||||
|
||||
return ChatResult(
|
||||
text="", model=name,
|
||||
prompt_version=getattr(prompt, "version", "inline"),
|
||||
usage=TokenUsage(), duration_ms=int((time.monotonic() - start) * 1000),
|
||||
status="failed", error=last_error,
|
||||
)
|
||||
|
||||
def chat_structured(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
prompt: Prompt | str,
|
||||
variables: dict,
|
||||
schema: dict,
|
||||
retry_count: int = 2,
|
||||
) -> StructuredResult:
|
||||
rendered = self._render_prompt(prompt, variables)
|
||||
# 追加 schema 约束说明（不强制模板支持）
|
||||
schema_hint = json.dumps(schema, ensure_ascii=False) if schema else ""
|
||||
base_rendered = rendered + (f'\n\n请输出符合以下 JSON Schema 的 JSON：{schema_hint}' if schema_hint else "")
|
||||
|
||||
start = time.monotonic()
|
||||
attempts = 0
|
||||
last_raw = ""
|
||||
last_error: str | None = None
|
||||
|
||||
while attempts <= retry_count:
|
||||
attempts += 1
|
||||
try:
|
||||
text, usage = self._call(
|
||||
model=self._model_names(None)[0], # 解析重试用首选模型
|
||||
rendered=base_rendered,
|
||||
temperature=0.0, max_tokens=4096,
|
||||
)
|
||||
last_raw = text
|
||||
data = json.loads(text)
|
||||
return StructuredResult(
|
||||
data=data, raw_text=text, parse_attempts=attempts,
|
||||
model=self._model_names(None)[0],
|
||||
prompt_version=getattr(prompt, "version", "inline"),
|
||||
usage=usage,
|
||||
duration_ms=int((time.monotonic() - start) * 1000),
|
||||
# fallback 语义保留给模型降级；解析重试成功仍为 ok
|
||||
status="ok",
|
||||
)
|
||||
except json.JSONDecodeError as exc:
|
||||
last_error = f"JSON 解析失败: {exc}"
|
||||
# 带错误信息重试
|
||||
base_rendered = base_rendered + f"\n\n上次解析失败：{exc}。请重新输出合法 JSON。"
|
||||
except LLMError as exc:
|
||||
return StructuredResult(
|
||||
data={}, raw_text="", parse_attempts=attempts,
|
||||
model=self._model_names(None)[0],
|
||||
prompt_version=getattr(prompt, "version", "inline"),
|
||||
usage=TokenUsage(),
|
||||
duration_ms=int((time.monotonic() - start) * 1000),
|
||||
status="failed", error=str(exc),
|
||||
)
|
||||
|
||||
return StructuredResult(
|
||||
data={}, raw_text=last_raw, parse_attempts=attempts,
|
||||
model=self._model_names(None)[0],
|
||||
prompt_version=getattr(prompt, "version", "inline"),
|
||||
usage=TokenUsage(),
|
||||
duration_ms=int((time.monotonic() - start) * 1000),
|
||||
status="parse_error", error=last_error,
|
||||
)
|
||||
Reference in New Issue
Block a user