“ydy0615”
101 lines
3.4 KiB
Python
101 lines
3.4 KiB
Python
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from risk_config import RiskConfig
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LLMPolicy:
|
|
job_type: str
|
|
model: str
|
|
profile: str
|
|
max_input_chars: int
|
|
max_output_tokens: int
|
|
temperature: float
|
|
thinking: str | None
|
|
|
|
|
|
def _normalize_thinking(value: str | None, *, allow_high: bool) -> str | None:
|
|
candidate = (value or "").strip().lower()
|
|
if candidate in {"", "none", "off"}:
|
|
return None
|
|
if candidate not in {"low", "medium", "high"}:
|
|
return "low"
|
|
if candidate == "high" and not allow_high:
|
|
return "medium"
|
|
return candidate
|
|
|
|
|
|
def resolve_llm_policy(job_type: str, request_payload: dict[str, Any], config: RiskConfig) -> LLMPolicy:
|
|
if job_type == "completion":
|
|
return LLMPolicy(
|
|
job_type=job_type,
|
|
model=config.completion_model,
|
|
profile="completion",
|
|
max_input_chars=config.completion_max_input_chars,
|
|
max_output_tokens=config.completion_max_output_tokens,
|
|
temperature=config.completion_temperature,
|
|
thinking=_normalize_thinking(request_payload.get("model_thinking"), allow_high=False),
|
|
)
|
|
if job_type == "pro_completion":
|
|
return LLMPolicy(
|
|
job_type=job_type,
|
|
model=config.pro_model,
|
|
profile="pro",
|
|
max_input_chars=config.pro_max_input_chars,
|
|
max_output_tokens=config.pro_max_output_tokens,
|
|
temperature=config.pro_temperature,
|
|
thinking=_normalize_thinking(request_payload.get("pro_thinking"), allow_high=True) or "medium",
|
|
)
|
|
if job_type == "web_search":
|
|
return LLMPolicy(
|
|
job_type=job_type,
|
|
model=config.web_search_model,
|
|
profile="completion",
|
|
max_input_chars=config.web_search_max_input_chars,
|
|
max_output_tokens=config.web_search_max_output_tokens,
|
|
temperature=config.web_search_temperature,
|
|
thinking="low",
|
|
)
|
|
if job_type == "compress":
|
|
return LLMPolicy(
|
|
job_type=job_type,
|
|
model=config.completion_model,
|
|
profile="completion",
|
|
max_input_chars=config.compress_max_input_chars,
|
|
max_output_tokens=config.compress_max_output_tokens,
|
|
temperature=0.2,
|
|
thinking="low",
|
|
)
|
|
if job_type == "ocr":
|
|
return LLMPolicy(
|
|
job_type=job_type,
|
|
model=config.vision_model,
|
|
profile="vision",
|
|
max_input_chars=config.ocr_max_input_bytes,
|
|
max_output_tokens=config.completion_max_output_tokens,
|
|
temperature=0.0,
|
|
thinking=None,
|
|
)
|
|
if job_type == "tts":
|
|
return LLMPolicy(
|
|
job_type=job_type,
|
|
model=config.speech_tts_model,
|
|
profile="speech_tts",
|
|
max_input_chars=config.speech_tts_max_input_chars,
|
|
max_output_tokens=0,
|
|
temperature=0.0,
|
|
thinking=None,
|
|
)
|
|
if job_type == "asr":
|
|
return LLMPolicy(
|
|
job_type=job_type,
|
|
model=config.speech_asr_model,
|
|
profile="speech_asr",
|
|
max_input_chars=config.speech_asr_max_input_bytes,
|
|
max_output_tokens=0,
|
|
temperature=0.0,
|
|
thinking=None,
|
|
)
|
|
raise ValueError(f"unsupported llm policy job type: {job_type}")
|