191 lines
7.1 KiB
Python
191 lines
7.1 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
from urllib.error import HTTPError, URLError
|
|
from urllib.request import Request, urlopen
|
|
|
|
|
|
@dataclass
|
|
class ModelResponse:
|
|
ok: bool
|
|
content: str = ""
|
|
parsed: dict[str, Any] | None = None
|
|
duration_seconds: float = 0.0
|
|
error: str = ""
|
|
|
|
|
|
def _chat_completions_url(base_url: str) -> str:
|
|
base = base_url.rstrip("/")
|
|
if base.endswith("/chat/completions"):
|
|
return base
|
|
if base.endswith("/v1"):
|
|
return f"{base}/chat/completions"
|
|
return f"{base}/v1/chat/completions"
|
|
|
|
|
|
def parse_json_object(text: str) -> dict[str, Any] | None:
|
|
raw = text.strip()
|
|
fence = re.search(r"```(?:json)?\s*(.*?)```", raw, re.IGNORECASE | re.DOTALL)
|
|
if fence:
|
|
raw = fence.group(1).strip()
|
|
start = raw.find("{")
|
|
end = raw.rfind("}")
|
|
if start >= 0 and end > start:
|
|
raw = raw[start : end + 1]
|
|
try:
|
|
parsed = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
return None
|
|
return parsed if isinstance(parsed, dict) else None
|
|
|
|
|
|
def _network_error_message(label: str, exc: BaseException, timeout: float) -> str:
|
|
reason = getattr(exc, "reason", exc)
|
|
if isinstance(exc, TimeoutError) or isinstance(reason, TimeoutError):
|
|
detail = f"timed out after {timeout:g}s"
|
|
else:
|
|
detail = str(reason) or type(reason).__name__
|
|
return f"{label} -> {detail}"
|
|
|
|
|
|
class ModelClient:
|
|
def __init__(self, model_config: dict[str, Any], timeout_seconds: float = 30):
|
|
self.model_config = model_config
|
|
self.timeout_seconds = timeout_seconds
|
|
|
|
@property
|
|
def available(self) -> bool:
|
|
return bool(
|
|
self.model_config.get("controller_model_proxy_url")
|
|
or (self.model_config.get("vllm_base_url") and self.model_config.get("vllm_model"))
|
|
)
|
|
|
|
def chat_json(self, messages: list[dict[str, str]], temperature: float | None = None) -> ModelResponse:
|
|
if not self.available:
|
|
return ModelResponse(ok=False, error="Model config is incomplete.")
|
|
|
|
proxy_url = str(self.model_config.get("controller_model_proxy_url") or "").strip()
|
|
if proxy_url:
|
|
return self._chat_json_via_controller(proxy_url, messages, temperature)
|
|
|
|
url = _chat_completions_url(str(self.model_config.get("vllm_base_url") or ""))
|
|
payload = {
|
|
"model": self.model_config.get("vllm_model"),
|
|
"messages": messages,
|
|
"temperature": temperature if temperature is not None else self.model_config.get("temperature", 0),
|
|
}
|
|
body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
|
headers = {
|
|
"Accept": "application/json",
|
|
"Content-Type": "application/json; charset=utf-8",
|
|
}
|
|
api_key = str(self.model_config.get("api_key") or "")
|
|
if api_key:
|
|
headers["Authorization"] = f"Bearer {api_key}"
|
|
|
|
started = time.monotonic()
|
|
request = Request(url, data=body, headers=headers, method="POST")
|
|
try:
|
|
with urlopen(request, timeout=self.timeout_seconds) as response:
|
|
raw = response.read()
|
|
except HTTPError as exc:
|
|
error_body = exc.read().decode("utf-8", errors="replace")
|
|
return ModelResponse(
|
|
ok=False,
|
|
duration_seconds=round(time.monotonic() - started, 3),
|
|
error=f"{url} -> HTTP {exc.code}: {error_body or exc.reason}",
|
|
)
|
|
except (URLError, TimeoutError, OSError) as exc:
|
|
return ModelResponse(
|
|
ok=False,
|
|
duration_seconds=round(time.monotonic() - started, 3),
|
|
error=_network_error_message(url, exc, self.timeout_seconds),
|
|
)
|
|
|
|
duration = round(time.monotonic() - started, 3)
|
|
try:
|
|
data = json.loads(raw.decode("utf-8"))
|
|
choice = data["choices"][0]
|
|
message = choice.get("message") or {}
|
|
content = (
|
|
message.get("content")
|
|
or message.get("reasoning_content")
|
|
or message.get("reasoning")
|
|
or choice.get("text")
|
|
or ""
|
|
)
|
|
except (KeyError, IndexError, TypeError, json.JSONDecodeError) as exc:
|
|
return ModelResponse(ok=False, duration_seconds=duration, error=f"Invalid model response: {exc}")
|
|
if not isinstance(content, str):
|
|
content = json.dumps(content, ensure_ascii=False)
|
|
|
|
return ModelResponse(
|
|
ok=True,
|
|
content=str(content),
|
|
parsed=parse_json_object(str(content)),
|
|
duration_seconds=duration,
|
|
)
|
|
|
|
def _chat_json_via_controller(
|
|
self,
|
|
url: str,
|
|
messages: list[dict[str, str]],
|
|
temperature: float | None = None,
|
|
) -> ModelResponse:
|
|
payload = {
|
|
"messages": messages,
|
|
"temperature": temperature if temperature is not None else self.model_config.get("temperature", 0),
|
|
}
|
|
body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
|
headers = {
|
|
"Accept": "application/json",
|
|
"Content-Type": "application/json; charset=utf-8",
|
|
}
|
|
token = str(self.model_config.get("controller_model_proxy_token") or "")
|
|
if token:
|
|
headers["Authorization"] = f"Bearer {token}"
|
|
|
|
started = time.monotonic()
|
|
request = Request(url, data=body, headers=headers, method="POST")
|
|
try:
|
|
with urlopen(request, timeout=self.timeout_seconds) as response:
|
|
raw = response.read()
|
|
except HTTPError as exc:
|
|
error_body = exc.read().decode("utf-8", errors="replace")
|
|
return ModelResponse(
|
|
ok=False,
|
|
duration_seconds=round(time.monotonic() - started, 3),
|
|
error=f"controller model proxy -> HTTP {exc.code}: {error_body or exc.reason}",
|
|
)
|
|
except (URLError, TimeoutError, OSError) as exc:
|
|
return ModelResponse(
|
|
ok=False,
|
|
duration_seconds=round(time.monotonic() - started, 3),
|
|
error=_network_error_message("controller model proxy", exc, self.timeout_seconds),
|
|
)
|
|
|
|
duration = round(time.monotonic() - started, 3)
|
|
try:
|
|
data = json.loads(raw.decode("utf-8")) if raw else {}
|
|
except json.JSONDecodeError as exc:
|
|
return ModelResponse(ok=False, duration_seconds=duration, error=f"Invalid controller proxy response: {exc}")
|
|
|
|
content = str(data.get("content") or "")
|
|
parsed = data.get("parsed") if isinstance(data.get("parsed"), dict) else parse_json_object(content)
|
|
proxy_duration = data.get("duration_seconds")
|
|
try:
|
|
duration = float(proxy_duration) if proxy_duration is not None else duration
|
|
except (TypeError, ValueError):
|
|
pass
|
|
return ModelResponse(
|
|
ok=bool(data.get("ok")),
|
|
content=content,
|
|
parsed=parsed,
|
|
duration_seconds=round(duration, 3),
|
|
error=str(data.get("error") or ""),
|
|
)
|