117 lines
4.2 KiB
Python
117 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import sys
|
|
import urllib.error
|
|
|
|
from common import ROOT, call_chat_completion, read_json
|
|
|
|
|
|
REGISTRY_INDEX = ROOT / "registry" / "index.json"
|
|
DEFAULT_BASE_URL = "http://docker-gpu.cin.su:8000"
|
|
PLUGIN_TASKS = {
|
|
"text": {"text", "chat", "summarization", "code", "tool-use"},
|
|
"translation": {"translation"},
|
|
"audio": {"speech-to-text", "speech-translation"},
|
|
"video": {"video", "image-understanding", "document-understanding", "visual-question-answering"},
|
|
"1c": {"1c", "1c-rag", "bsl-code", "metadata-safety", "1c-query"},
|
|
}
|
|
DEFAULT_PROMPTS = {
|
|
"text": "Кратко объясни, зачем нужен реестр локальных моделей.",
|
|
"translation": "Переведи на английский: Проверяем локальную модель перевода.",
|
|
"audio": "Составь чек-лист проверки качества распознавания речи.",
|
|
"video": "Какие тесты нужны для проверки модели анализа видео?",
|
|
"1c": "Какие метаданные 1С нужно запросить перед изменением документа?",
|
|
}
|
|
STATUS_RANK = {
|
|
"production": 0,
|
|
"staging": 1,
|
|
"candidate": 2,
|
|
"draft": 3,
|
|
"archived": 9,
|
|
}
|
|
|
|
|
|
def model_matches_plugin(model: dict, plugin: str) -> bool:
|
|
tasks = set(model.get("task") or [])
|
|
return bool(tasks & PLUGIN_TASKS.get(plugin, set()))
|
|
|
|
|
|
def model_score(model: dict, plugin: str) -> tuple[int, int, str]:
|
|
tasks = set(model.get("task") or [])
|
|
status = str(model.get("status") or "")
|
|
score = 0
|
|
if plugin == "text":
|
|
if tasks & {"text", "chat"}:
|
|
score -= 5
|
|
if tasks & {"1c", "1c-rag", "bsl-code"}:
|
|
score += 6
|
|
elif plugin == "1c":
|
|
if tasks & PLUGIN_TASKS["1c"]:
|
|
score -= 6
|
|
else:
|
|
if tasks & PLUGIN_TASKS.get(plugin, set()):
|
|
score -= 5
|
|
return (score, STATUS_RANK.get(status, 8), str(model.get("id") or ""))
|
|
|
|
|
|
def select_model(plugin: str, model_id: str | None) -> dict:
|
|
registry = read_json(REGISTRY_INDEX)
|
|
models = registry.get("models") or []
|
|
candidates = [model for model in models if model_matches_plugin(model, plugin)]
|
|
if model_id:
|
|
candidates = [model for model in candidates if model.get("id") == model_id]
|
|
candidates.sort(key=lambda model: model_score(model, plugin))
|
|
if not candidates:
|
|
raise ValueError(f"No model found for plugin={plugin!r} model={model_id!r}")
|
|
return candidates[0]
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Chat with a registry model through an OpenAI-compatible endpoint.")
|
|
parser.add_argument("--plugin", choices=sorted(PLUGIN_TASKS), default="text")
|
|
parser.add_argument("--model-id", help="Registry model id. Defaults to the first model for the plugin.")
|
|
parser.add_argument("--base-url", default=DEFAULT_BASE_URL)
|
|
parser.add_argument("--prompt")
|
|
parser.add_argument("--temperature", type=float, default=0.2)
|
|
parser.add_argument("--max-tokens", type=int, default=1000)
|
|
parser.add_argument("--system", default="Отвечай по-русски, кратко и по делу.")
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
model = select_model(args.plugin, args.model_id)
|
|
except ValueError as exc:
|
|
print(exc, file=sys.stderr)
|
|
return 1
|
|
|
|
served_model = model.get("served_model_name") or model.get("id")
|
|
prompt = args.prompt or DEFAULT_PROMPTS[args.plugin]
|
|
messages = [
|
|
{"role": "system", "content": args.system},
|
|
{"role": "user", "content": prompt},
|
|
]
|
|
|
|
print(f"Plugin: {args.plugin}")
|
|
print(f"Model: {model.get('id')} -> {served_model}")
|
|
print(f"Endpoint: {args.base_url}")
|
|
print("")
|
|
|
|
try:
|
|
answer = call_chat_completion(
|
|
base_url=args.base_url,
|
|
model=served_model,
|
|
messages=messages,
|
|
temperature=args.temperature,
|
|
max_tokens=args.max_tokens,
|
|
)
|
|
except (urllib.error.URLError, ValueError) as exc:
|
|
print(f"Chat request failed: {exc}", file=sys.stderr)
|
|
return 1
|
|
|
|
print(answer)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|