306 lines
11 KiB
Python
306 lines
11 KiB
Python
"""MiniMax provider: OpenAI-compatible chat completions with JSON-schema output and model fallback."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import json
|
|
import logging
|
|
import re
|
|
from typing import Any, get_origin
|
|
|
|
import openai
|
|
from pydantic import ValidationError
|
|
|
|
from us2kr.config import Settings
|
|
from us2kr.llm.base import LLMError
|
|
from us2kr.llm.prompts import SYSTEM_PROMPT, build_user_message
|
|
from us2kr.models import MarketPayload, Report
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
# Errors from the openai SDK that mean "this model/attempt failed; try the next model".
|
|
OPENAI_ERRORS: tuple[type[Exception], ...] = (
|
|
openai.APIStatusError,
|
|
openai.APIConnectionError,
|
|
openai.APITimeoutError,
|
|
openai.RateLimitError,
|
|
)
|
|
|
|
_THINK_BLOCK = re.compile(r"<think>.*?</think>", re.DOTALL)
|
|
_CODE_FENCE = re.compile(r"```(?:json|JSON)?\s*(.*?)```", re.DOTALL)
|
|
|
|
REPAIR_INSTRUCTION = (
|
|
"위 응답은 JSON 스키마 검증에 실패했습니다. 아래 오류를 참고하여, 설명 없이 스키마에 맞는 "
|
|
"완전한 JSON 객체 하나만 다시 출력하세요.\n\n오류:\n"
|
|
)
|
|
|
|
|
|
def _strip_titles(node: Any) -> None:
|
|
"""Recursively remove JSON-schema "title" keys (noise for the model, not needed for validation)."""
|
|
if isinstance(node, dict):
|
|
node.pop("title", None)
|
|
for value in node.values():
|
|
_strip_titles(value)
|
|
elif isinstance(node, list):
|
|
for value in node:
|
|
_strip_titles(value)
|
|
|
|
|
|
def report_json_schema() -> dict[str, Any]:
|
|
"""JSON schema for `Report` suitable for `response_format` (no titles, closed object)."""
|
|
schema = copy.deepcopy(Report.model_json_schema())
|
|
_strip_titles(schema)
|
|
schema["additionalProperties"] = False
|
|
return schema
|
|
|
|
|
|
def report_response_format() -> dict[str, Any]:
|
|
return {"type": "json_schema", "json_schema": {"name": "report", "schema": report_json_schema()}}
|
|
|
|
|
|
def _first_json_object(text: str) -> str | None:
|
|
"""Return the first balanced `{...}` block in `text` (string-aware), or None."""
|
|
start = text.find("{")
|
|
if start < 0:
|
|
return None
|
|
depth = 0
|
|
in_str = False
|
|
escape = False
|
|
for i in range(start, len(text)):
|
|
ch = text[i]
|
|
if in_str:
|
|
if escape:
|
|
escape = False
|
|
elif ch == "\\":
|
|
escape = True
|
|
elif ch == '"':
|
|
in_str = False
|
|
continue
|
|
if ch == '"':
|
|
in_str = True
|
|
elif ch == "{":
|
|
depth += 1
|
|
elif ch == "}":
|
|
depth -= 1
|
|
if depth == 0:
|
|
return text[start : i + 1]
|
|
return None
|
|
|
|
|
|
def extract_json_text(raw: str) -> str:
|
|
"""Strip <think> blocks and code fences from a model reply, returning the JSON candidate text.
|
|
|
|
Handles: complete <think>...</think> blocks anywhere, an unterminated <think> (take text after the
|
|
last </think>, or the first {...} object), and ```json fences. Raises ValueError if nothing remains.
|
|
"""
|
|
text = _THINK_BLOCK.sub("", raw)
|
|
if "<think>" in text:
|
|
if "</think>" in raw:
|
|
text = raw.rsplit("</think>", 1)[1]
|
|
else:
|
|
text = _first_json_object(raw) or ""
|
|
fence = _CODE_FENCE.search(text)
|
|
if fence:
|
|
text = fence.group(1)
|
|
text = _first_json_object(text.strip()) or ""
|
|
if not text:
|
|
raise ValueError("no JSON object found in model reply")
|
|
return text
|
|
|
|
|
|
_LIST_FIELDS: frozenset[str] = frozenset(
|
|
name for name, info in Report.model_fields.items() if get_origin(info.annotation) is list
|
|
)
|
|
_BULLET_SPLIT = re.compile(r"(?:^|\n)\s*(?:[-*•]|\d+[.)])\s+")
|
|
_SENTENCE_SPLIT = re.compile(r"(?<=[.!?。])\s+")
|
|
|
|
|
|
def _split_prose(text: str) -> list[str]:
|
|
"""Split a prose string into list items: bullets/newlines first, then sentence boundaries."""
|
|
parts = [p.strip() for p in _BULLET_SPLIT.split(text) if p.strip()]
|
|
if len(parts) <= 1:
|
|
parts = [p.strip() for p in re.split(r"\n+", text) if p.strip()]
|
|
if len(parts) <= 1:
|
|
parts = [p.strip() for p in _SENTENCE_SPLIT.split(text) if p.strip()]
|
|
return parts or [text.strip()]
|
|
|
|
|
|
def coerce_report_data(data: Any) -> tuple[Any, list[str]]:
|
|
"""Fix the common MiniMax schema slip of emitting list fields as a single string.
|
|
|
|
Returns (data, names of coerced fields). Anything that is not a dict is returned untouched so that
|
|
pydantic can report the real error.
|
|
"""
|
|
if not isinstance(data, dict):
|
|
return data, []
|
|
coerced: list[str] = []
|
|
for name in _LIST_FIELDS:
|
|
value = data.get(name)
|
|
if isinstance(value, str):
|
|
data[name] = _split_prose(value)
|
|
coerced.append(name)
|
|
return data, coerced
|
|
|
|
|
|
def parse_report(raw: str) -> Report:
|
|
"""Turn raw model content into a validated Report (raises ValueError / ValidationError)."""
|
|
text = extract_json_text(raw)
|
|
try:
|
|
data = json.loads(text)
|
|
except json.JSONDecodeError:
|
|
return Report.model_validate_json(text) # let pydantic produce the canonical error
|
|
data, coerced = coerce_report_data(data)
|
|
if coerced:
|
|
log.warning("coerced string -> list for report fields: %s", ", ".join(sorted(coerced)))
|
|
return Report.model_validate(data)
|
|
|
|
|
|
class MiniMaxProvider:
|
|
"""LLMProvider backed by MiniMax's OpenAI-compatible endpoint with a model fallback chain."""
|
|
|
|
def __init__(self, settings: Settings, client: Any | None = None) -> None:
|
|
self._settings = settings
|
|
models = [settings.minimax_model, *settings.minimax_fallback_models]
|
|
self._models = list(dict.fromkeys(m for m in models if m))
|
|
if not self._models:
|
|
raise LLMError("no MiniMax model configured")
|
|
if client is None:
|
|
if not settings.minimax_api_key:
|
|
raise LLMError("MINIMAX_API_KEY is not set")
|
|
client = openai.OpenAI(
|
|
api_key=settings.minimax_api_key,
|
|
base_url=settings.minimax_base_url,
|
|
timeout=settings.minimax_timeout_seconds,
|
|
max_retries=2,
|
|
)
|
|
self._client = client
|
|
self._response_format = report_response_format()
|
|
|
|
@property
|
|
def models(self) -> list[str]:
|
|
return list(self._models)
|
|
|
|
# ------------------------------------------------------------------ public
|
|
def summarize(self, payload: MarketPayload) -> tuple[Report, str]:
|
|
user_message = build_user_message(payload)
|
|
base_messages: list[dict[str, Any]] = [
|
|
{"role": "system", "content": SYSTEM_PROMPT},
|
|
{"role": "user", "content": user_message},
|
|
]
|
|
attempts: list[str] = []
|
|
for model in self._models:
|
|
result = self._attempt_model(model, base_messages, attempts)
|
|
if result is not None:
|
|
return result, model
|
|
raise LLMError("all MiniMax models failed: " + " | ".join(attempts))
|
|
|
|
# ----------------------------------------------------------------- helpers
|
|
def _attempt_model(
|
|
self, model: str, messages: list[dict[str, Any]], attempts: list[str]
|
|
) -> Report | None:
|
|
"""One model: initial call plus at most one repair retry. Returns None when the model failed."""
|
|
content = self._call(model, messages, attempts)
|
|
if content is None:
|
|
return None
|
|
try:
|
|
return parse_report(content)
|
|
except (ValidationError, ValueError) as exc:
|
|
error_text = _short(str(exc), 800)
|
|
attempts.append(f"{model}: invalid report ({_short(str(exc), 160)})")
|
|
log.warning(
|
|
"model %s returned invalid report, retrying once with repair hint: %s",
|
|
model,
|
|
_short(str(exc), 300),
|
|
)
|
|
|
|
repair_messages = [
|
|
*messages,
|
|
{"role": "assistant", "content": _short(content, 12000)},
|
|
{"role": "user", "content": REPAIR_INSTRUCTION + error_text},
|
|
]
|
|
content = self._call(model, repair_messages, attempts, label="repair")
|
|
if content is None:
|
|
return None
|
|
try:
|
|
return parse_report(content)
|
|
except (ValidationError, ValueError) as exc:
|
|
attempts.append(f"{model} (repair): invalid report ({_short(str(exc), 160)})")
|
|
log.warning("model %s repair retry still invalid: %s", model, _short(str(exc), 300))
|
|
return None
|
|
|
|
def _call(
|
|
self, model: str, messages: list[dict[str, Any]], attempts: list[str], label: str = "initial"
|
|
) -> str | None:
|
|
"""Perform one chat completion. Returns content text, or None (attempt recorded) on failure."""
|
|
tag = model if label == "initial" else f"{model} ({label})"
|
|
extra_body: dict[str, Any] = {"reasoning_split": True}
|
|
if model.startswith("MiniMax-M3"):
|
|
extra_body["thinking"] = {"type": "adaptive"}
|
|
try:
|
|
completion = self._client.chat.completions.create(
|
|
model=model,
|
|
messages=messages,
|
|
max_completion_tokens=self._settings.minimax_max_completion_tokens,
|
|
response_format=self._response_format,
|
|
extra_body=extra_body,
|
|
)
|
|
except OPENAI_ERRORS as exc:
|
|
status = getattr(exc, "status_code", None)
|
|
detail = f"{type(exc).__name__}" + (f" {status}" if status else "") + f": {_short(str(exc), 160)}"
|
|
attempts.append(f"{tag}: {detail}")
|
|
log.warning("MiniMax call failed for %s: %s", tag, detail)
|
|
return None
|
|
|
|
choices = getattr(completion, "choices", None) or []
|
|
if not choices:
|
|
attempts.append(f"{tag}: empty choices")
|
|
log.warning("MiniMax %s returned no choices", tag)
|
|
return None
|
|
choice = choices[0]
|
|
finish_reason = getattr(choice, "finish_reason", None)
|
|
self._log_usage(model, completion, finish_reason)
|
|
if finish_reason == "length":
|
|
attempts.append(f"{tag}: truncated (finish_reason=length)")
|
|
log.warning("MiniMax %s output truncated (finish_reason=length)", tag)
|
|
return None
|
|
content = getattr(choice.message, "content", None)
|
|
if not content:
|
|
attempts.append(f"{tag}: empty content")
|
|
log.warning("MiniMax %s returned empty content", tag)
|
|
return None
|
|
return content
|
|
|
|
@staticmethod
|
|
def _log_usage(model: str, completion: Any, finish_reason: str | None) -> None:
|
|
usage = getattr(completion, "usage", None)
|
|
prompt_tokens = getattr(usage, "prompt_tokens", None)
|
|
completion_tokens = getattr(usage, "completion_tokens", None)
|
|
details = getattr(usage, "prompt_tokens_details", None)
|
|
cached = getattr(details, "cached_tokens", None) if details is not None else None
|
|
log.info(
|
|
"MiniMax response model=%s served_as=%s finish=%s prompt_tokens=%s completion_tokens=%s "
|
|
"cached_tokens=%s",
|
|
model,
|
|
getattr(completion, "model", None),
|
|
finish_reason,
|
|
prompt_tokens,
|
|
completion_tokens,
|
|
cached,
|
|
)
|
|
|
|
|
|
def _short(text: str, limit: int) -> str:
|
|
text = " ".join(text.split())
|
|
return text if len(text) <= limit else text[: limit - 1] + "…"
|
|
|
|
|
|
__all__ = [
|
|
"OPENAI_ERRORS",
|
|
"MiniMaxProvider",
|
|
"coerce_report_data",
|
|
"extract_json_text",
|
|
"parse_report",
|
|
"report_json_schema",
|
|
"report_response_format",
|
|
]
|