toolkit/verify.py

# -*- coding: utf-8 -*-
"""ZuperAI с исполняемой самопроверкой (test-driven): тесты → код → прогон → чинка.

Поток:
  1. «судья» пишет юнит-тесты ПО ЗАДАЧЕ (контракт, без видения кода);
  2. локальная модель пишет код, который должен пройти эти тесты;
  3. тесты гоняются в песочнице (отдельный subprocess, таймаут);
  4. при провале судья решает, где ошибка — в коде или в тесте, и чинит её;
  5. повтор, пока тесты не пройдут (или не выйдет лимит попыток).

Кто есть кто:
  --code local|lmstudio     кто пишет КОД. local = transformers+LoRA (грузит 3B в память),
                            lmstudio = LM Studio (GGUF, быстрее, без torch). По умолчанию local.
  --judge deepseek|local|lmstudio  кто пишет ТЕСТЫ и судит. По умолчанию deepseek, если ключ есть.

  Практичные связки:
    --code lmstudio --judge deepseek   → код локально (быстро), тесты/судья сильные
    --code lmstudio --judge lmstudio   → полностью офлайн, без DeepSeek и без torch

Запуск:
  venv/Scripts/python.exe verify.py "Напиши функцию, которая ..."

Ключи:
  --attempts N   число попыток (по умолчанию 5)
  --timeout N    таймаут прогона тестов, сек (по умолчанию 30)
  --keep         не удалять песочницу (покажет путь к solution.py/test_*.py)
"""
import sys
import tempfile

from webui import (deepseek_available, deepseek_generate,
                   lmstudio_available, lmstudio_generate)
from verify_util import (TEST_SYSTEM, VERIFY_ATTEMPTS, VERIFY_TIMEOUT,
                         run_tests, looks_like_tests, tail, ensure_unittest_header)

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


def _cli_args():
    args = sys.argv[1:]
    attempts, timeout, keep = VERIFY_ATTEMPTS, VERIFY_TIMEOUT, False
    code = "local"
    judge = "deepseek" if deepseek_available() else "local"
    prompt_parts = []
    i = 0
    while i < len(args):
        a = args[i]
        if a == "--attempts" and i + 1 < len(args):
            attempts = int(args[i + 1]); i += 2; continue
        if a == "--timeout" and i + 1 < len(args):
            timeout = int(args[i + 1]); i += 2; continue
        if a == "--judge" and i + 1 < len(args):
            judge = args[i + 1].lower(); i += 2; continue
        if a == "--code" and i + 1 < len(args):
            code = args[i + 1].lower(); i += 2; continue
        if a == "--keep":
            keep = True; i += 1; continue
        prompt_parts.append(a); i += 1
    if judge not in ("deepseek", "local", "lmstudio"):
        judge = "local"
    if code not in ("local", "lmstudio"):
        code = "local"
    if judge == "deepseek" and not deepseek_available():
        print("[!] DeepSeek недоступен — судья переключён на local", file=sys.stderr)
        judge = "local"
    if judge == "lmstudio" and not lmstudio_available():
        print("[!] LM Studio не запущен (порт 1234) — судья переключён на local", file=sys.stderr)
        judge = "local"
    if code == "lmstudio" and not lmstudio_available():
        print("[!] LM Studio не запущен (порт 1234) — код переключён на local", file=sys.stderr)
        code = "local"
    return (" ".join(prompt_parts) or
            "Напиши функцию, которая принимает список чисел и возвращает их среднее",
            attempts, timeout, keep, judge, code)


def main():
    prompt, max_attempts, timeout, keep, judge, code_backend = _cli_args()

    # грузим torch-модель только если кто-то (код или судья) работает через transformers
    tok = model = None
    if code_backend == "local" or judge == "local":
        from infer_smart import load, gen, _strip_fences as _sf
        tok, model = load()
    else:
        from infer_smart import _strip_fences as _sf

    def code_gen(msg, max_tokens=768):
        if code_backend == "lmstudio":
            return _sf(lmstudio_generate(msg, max_tokens=max_tokens))
        return _sf(gen(tok, model, msg, max_tokens=max_tokens))

    def judge_gen(msg, max_tokens=768):
        if judge == "deepseek":
            return _sf(deepseek_generate(msg, max_tokens=max_tokens))
        if judge == "lmstudio":
            return _sf(lmstudio_generate(msg, max_tokens=max_tokens))
        return _sf(gen(tok, model, msg, max_tokens=max_tokens))

    print(f"[*] код: {code_backend}, судья: {judge}", file=sys.stderr)
    if judge == "local":
        print("[!] честно: локальная 3B может ошибаться в ТЕСТАХ (неверные ожидания). "
              "Для надёжности используйте --judge deepseek.", file=sys.stderr)

    sandbox = tempfile.mkdtemp(prefix="zuperai_verify_")

    # --- шаг 1: судья пишет тесты ПО ЗАДАЧЕ (TDD) ---
    t_messages = [{"role": "system", "content": TEST_SYSTEM},
                  {"role": "user",
                   "content": f"Задача: {prompt}\n\nНапиши unittest-тесты, которые "
                              f"импортируют решение из solution и проверяют корректность. "
                              f"Верни ТОЛЬКО код тестов, без пояснений и без ```."}]
    tests = ensure_unittest_header(judge_gen(t_messages))

    # --- шаг 2+: модель пишет код, гоняем, чиним ---
    messages = [{"role": "system", "content": SYSTEM},
                {"role": "user",
                 "content": f"Задача: {prompt}\n\nТесты (unittest):\n{tests}\n\n"
                            f"Напиши код (solution.py), который проходит все эти тесты. "
                            f"Верни ТОЛЬКО код, без пояснений и без ```."}]

    code = ""
    for attempt in range(1, max_attempts + 1):
        if not code:
            code = code_gen(messages)
            print(f"[попытка {attempt}/{max_attempts}] генерирую код…", file=sys.stderr)
        else:
            print(f"[попытка {attempt}/{max_attempts}] тестирую исправленный код…",
                  file=sys.stderr)

        ok, out = run_tests(code, tests, sandbox, timeout)
        if ok:
            print(code)
            print(f"\n[OK] тесты прошли, попытка {attempt}")
            print("----- сгенерированные тесты -----")
            print(tests)
            print("----- песочница -----")
            print(sandbox if keep else "(временная, удалена)")
            return

        print(f"[попытка {attempt}] тесты упали:\n{tail(out)}", file=sys.stderr)
        messages.append({"role": "assistant", "content": code})
        messages.append({"role": "user",
                         "content": f"Решение не прошло тесты. Вывод тестов:\n"
                                    f"{tail(out)}\n\n"
                                    f"Определи, где ошибка: в коде (неверная реализация) "
                                    f"или в тесте (неверное ожидаемое значение). "
                                    f"Верни исправленный КОД, либо исправленный ТЕСТ — "
                                    f"только одно, без пояснений и без ```."})
        fix = judge_gen(messages)
        if looks_like_tests(fix):
            tests = ensure_unittest_header(fix)
            code = ""
            messages = messages[:2]
            messages[1] = {"role": "user",
                           "content": f"Задача: {prompt}\n\nТесты (unittest):\n{tests}\n\n"
                                      f"Напиши код (solution.py), который проходит все эти "
                                      f"тесты. Верни ТОЛЬКО код, без пояснений и без ```."}
            print("[*] судья решил, что ошибка в ТЕСТЕ — тест обновлён", file=sys.stderr)
        else:
            code = fix
            messages[-2] = {"role": "assistant", "content": code}

    print(code)
    print(f"\n[!] за {max_attempts} попыток тесты не прошли — выше последний вариант")
    print("----- последние тесты -----")
    print(tests)
    print("----- песочница -----")
    print(sandbox if keep else "(временная, удалена)")


if __name__ == "__main__":
    main()