toolkit/webui.py

# -*- coding: utf-8 -*-
"""Веб-интерфейс ZuperAI: FastAPI + страница с генерацией кода.

Фичи:
  - два «мозга»: DeepSeek (frontier) + локальная ZuperAI (офлайн);
  - режим: «код» (только код) или «свободный» (код + пояснения);
  - самопроверка: генерирует -> проверяет (синтаксис/линтер) -> чинит -> отдаёт.
"""
import json
import os
import re
import tempfile
import threading
import urllib.request
from pathlib import Path

import torch
from fastapi import FastAPI
from fastapi.responses import HTMLResponse, JSONResponse
from pydantic import BaseModel

from selfcheck import check_python
from verify_util import (TEST_SYSTEM, VERIFY_ATTEMPTS, VERIFY_TIMEOUT, run_tests,
                         looks_like_tests, tail, ensure_unittest_header)

BASE = "Qwen/Qwen2.5-Coder-3B-Instruct"
LORA = "ZuperAI-12.5.0.2"

SYSTEM_CODE = ("Ты — опытный программист. Отвечай только рабочим, лаконичным кодом "
               "без лишних пояснений. Комментарии и докстринги пиши на русском.")
SYSTEM_FREE = ("Ты — опытный программист. Пиши рабочий код и кратко поясняй решение, "
               "если уместно. Комментарии и докстринги — на русском.")

app = FastAPI(title="ZuperAI 12.5.0.2")
_lock = threading.Lock()
_model = None
_tokenizer = None
_ready = False
INDEX = Path(__file__).resolve().parent / "static" / "index.html"


def _load_env():
    if os.environ.get("DEEPSEEK_API_KEY"):
        return
    candidates = [Path.cwd() / ".env", Path.home() / ".env"]
    local = os.environ.get("LOCALAPPDATA")
    if local:
        candidates.append(Path(local) / "hermes" / ".env")
    wanted = ("DEEPSEEK_API_KEY", "DEEPSEEK_MODEL", "DEEPSEEK_BASE_URL")
    for p in candidates:
        if not p.exists():
            continue
        try:
            for line in p.read_text(encoding="utf-8", errors="ignore").splitlines():
                line = line.strip()
                if not line or line.startswith("#") or "=" not in line:
                    continue
                k, _, v = line.partition("=")
                k, v = k.strip(), v.strip().strip('"').strip("'")
                if k in wanted and k not in os.environ and v:
                    os.environ[k] = v
        except Exception:
            continue


_load_env()


def deepseek_available():
    return bool(os.environ.get("DEEPSEEK_API_KEY"))


def deepseek_generate(messages, max_tokens=1024):
    key = os.environ["DEEPSEEK_API_KEY"]
    base = os.environ.get("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1")
    model = os.environ.get("DEEPSEEK_MODEL", "deepseek-chat")
    payload = json.dumps({"model": model, "messages": messages,
                          "max_tokens": max_tokens, "temperature": 0.2}).encode()
    req = urllib.request.Request(base.rstrip("/") + "/chat/completions", data=payload,
                                 headers={"Content-Type": "application/json",
                                          "Authorization": "Bearer " + key})
    r = urllib.request.urlopen(req, timeout=180)
    data = json.loads(r.read())
    return data["choices"][0]["message"]["content"].strip()


LMSTUDIO_BASE = os.environ.get("LMSTUDIO_BASE_URL", "http://127.0.0.1:1234/v1")


def lmstudio_available():
    """Проверяет, запущен ли LM Studio-сервер (OpenAI-совместимый, порт 1234)."""
    try:
        r = urllib.request.urlopen(LMSTUDIO_BASE.rstrip("/") + "/models", timeout=3)
        return r.status == 200
    except Exception:
        return False


def lmstudio_generate(messages, max_tokens=512):
    """Генерация через LM Studio (OpenAI-совместимый API, без ключа)."""
    payload = json.dumps({"model": "local-model", "messages": messages,
                          "max_tokens": max_tokens, "temperature": 0.2}).encode()
    req = urllib.request.Request(LMSTUDIO_BASE.rstrip("/") + "/chat/completions",
                                 data=payload,
                                 headers={"Content-Type": "application/json"})
    r = urllib.request.urlopen(req, timeout=300)
    data = json.loads(r.read())
    return data["choices"][0]["message"]["content"].strip()


def load_model():
    global _model, _tokenizer, _ready
    with _lock:
        if _ready:
            return
        from peft import PeftModel
        from transformers import AutoModelForCausalLM, AutoTokenizer
        _tokenizer = AutoTokenizer.from_pretrained(BASE, trust_remote_code=True)
        _model = AutoModelForCausalLM.from_pretrained(BASE, dtype=torch.bfloat16,
                                                      trust_remote_code=True).to("cuda")
        _model = PeftModel.from_pretrained(_model, LORA)
        _model.eval()
        _ready = True


def local_generate(messages, max_tokens=512):
    if not _ready:
        load_model()
    enc = _tokenizer.apply_chat_template(messages, tokenize=True,
                                         add_generation_prompt=True, return_tensors="pt")
    ids = enc["input_ids"].to(_model.device)
    with torch.no_grad():
        out = _model.generate(ids, max_new_tokens=max_tokens, do_sample=False)
    return _tokenizer.decode(out[0][ids.shape[1]:], skip_special_tokens=True)


def _strip_fences(code):
    """Снимает ```python ... ``` обёртку, если она есть."""
    code = code.strip()
    m = re.match(r"```[a-zA-Z0-9+#\-]*\s*\n", code)
    if m:
        code = code[m.end():]
    if code.rstrip().endswith("```"):
        code = code.rstrip()[:-3]
    return code.strip()


def _looks_python(code):
    head = code.strip().lower()
    return ("def " in head or "class " in head or "import " in head
            or "from " in head or "async def" in head or "@app" in head)


class GenReq(BaseModel):
    prompt: str
    max_tokens: int = 512
    backend: str = "auto"      # auto | deepseek | local
    mode: str = "code"         # code | free
    selfcheck: bool = True
    verify: bool = False       # исполняемая самопроверка: тесты -> код -> прогон -> чинка


def _verify(req, run_once, backend_used):
    """Исполняемая самопроверка: тесты -> код -> прогон в песочнице -> чинка."""
    prompt = req.prompt
    judge_deepseek = deepseek_available()

    def judge_gen(msgs):
        if judge_deepseek:
            return _strip_fences(deepseek_generate(msgs, req.max_tokens))
        if os.environ.get("ZUPERAI_NO_LOAD") == "1":
            raise RuntimeError("судьи нет: DeepSeek недоступен, локальная отключена")
        return _strip_fences(local_generate(msgs, req.max_tokens))

    sandbox = tempfile.mkdtemp(prefix="zuperai_verify_")
    try:
        t_msgs = [{"role": "system", "content": TEST_SYSTEM},
                  {"role": "user",
                   "content": f"Задача: {prompt}\n\nНапиши unittest-тесты, которые "
                              f"импортируют решение из solution и проверяют корректность. "
                              f"Верни ТОЛЬКО код тестов, без пояснений и без ```."}]
        tests = ensure_unittest_header(judge_gen(t_msgs))

        code_msgs = [{"role": "system", "content": SYSTEM_CODE},
                     {"role": "user",
                      "content": f"Задача: {prompt}\n\nТесты (unittest):\n{tests}\n\n"
                                 f"Напиши код (solution.py), который проходит все эти "
                                 f"тесты. Верни ТОЛЬКО код, без пояснений и без ```."}]
        code = ""
        for attempt in range(1, VERIFY_ATTEMPTS + 1):
            if not code:
                code = _strip_fences(run_once(code_msgs))
            ok, out = run_tests(code, tests, sandbox, VERIFY_TIMEOUT)
            if ok:
                return {"code": code, "tests": tests, "backend": backend_used,
                        "verified": True, "attempts": attempt,
                        "judge": "deepseek" if judge_deepseek else "local"}
            code_msgs.append({"role": "assistant", "content": code})
            code_msgs.append({"role": "user",
                              "content": f"Решение не прошло тесты. Вывод тестов:\n"
                                         f"{tail(out)}\n\n"
                                         f"Определи, где ошибка: в коде или в тесте. "
                                         f"Верни исправленный КОД либо исправленный ТЕСТ — "
                                         f"только одно, без пояснений и без ```."})
            fix = _strip_fences(judge_gen(code_msgs))
            if looks_like_tests(fix):
                tests = ensure_unittest_header(fix)
                code = ""
                code_msgs = [{"role": "system", "content": SYSTEM_CODE},
                             {"role": "user",
                              "content": f"Задача: {prompt}\n\nТесты (unittest):\n{tests}\n\n"
                                         f"Напиши код (solution.py), который проходит все "
                                         f"эти тесты. Верни ТОЛЬКО код, без ```."}]
            else:
                code = fix
                code_msgs[-2] = {"role": "assistant", "content": code}

        return {"code": code, "tests": tests, "backend": backend_used,
                "verified": False, "attempts": VERIFY_ATTEMPTS,
                "judge": "deepseek" if judge_deepseek else "local",
                "note": f"тесты не прошли за {VERIFY_ATTEMPTS} попыток"}
    except Exception as e:
        return JSONResponse({"error": f"Проверка тестами: {e}"}, status_code=502)


@app.get("/", response_class=HTMLResponse)
def index():
    return HTMLResponse(INDEX.read_text(encoding="utf-8"))


@app.get("/api/status")
def status():
    return {"ready": _ready, "deepseek": deepseek_available(),
            "lmstudio": lmstudio_available(), "name": "ZuperAI 12.5.0.2"}


@app.post("/api/generate")
def generate(req: GenReq):
    system = SYSTEM_FREE if req.mode == "free" else SYSTEM_CODE
    messages = [{"role": "system", "content": system},
                {"role": "user", "content": req.prompt}]

    def use_deepseek():
        return req.backend == "deepseek" or \
               (req.backend == "auto" and deepseek_available())

    def run_once(msgs):
        if req.backend == "lmstudio":
            return lmstudio_generate(msgs, req.max_tokens)
        if use_deepseek():
            return deepseek_generate(msgs, req.max_tokens)
        if os.environ.get("ZUPERAI_NO_LOAD") == "1":
            raise RuntimeError("локальная модель отключена (ZUPERAI_NO_LOAD=1)")
        return local_generate(msgs, req.max_tokens)

    if req.backend == "lmstudio":
        backend_used = "lmstudio"
    else:
        backend_used = "deepseek" if use_deepseek() else "local"

    if req.verify:
        return _verify(req, run_once, backend_used)

    try:
        code = run_once(messages)
    except Exception as e:
        return JSONResponse({"error": f"Генерация: {e}"}, status_code=502)

    code = _strip_fences(code)

    if req.selfcheck and _looks_python(code):
        for attempt in range(1, 4):
            ok, msg = check_python(code)
            if ok:
                return {"code": code, "backend": backend_used,
                        "verified": True, "attempts": attempt}
            messages.append({"role": "assistant", "content": code})
            messages.append({"role": "user", "content":
                             f"Твой код не прошёл проверку. Ошибка: {msg}. "
                             f"Верни ТОЛЬКО исправленный код, без пояснений и без "
                             f"обрамления в ```."})
            try:
                code = _strip_fences(run_once(messages))
            except Exception as e:
                return JSONResponse({"error": f"Генерация (fix): {e}"}, status_code=502)
        return {"code": code, "backend": backend_used,
                "verified": False, "note": "не прошло самопроверку после 3 попыток"}

    return {"code": code, "backend": backend_used, "verified": None}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8001)