toolkit/agent.py

# -*- coding: utf-8 -*-
"""ZuperAI — кодинг-агент: по ТЗ генерирует целое решение + тесты → гоняет → чинит.

Отличие от verify.py: ориентирован на ЦЕЛЫЙ модуль/приложение (не одну функцию).
Пишет решение из нескольких функций/классов, сам пишет unittest-тесты, гоняет
их в песочнице и доводит до зелёного. Это «агент, который делает работу», а не
«генератор куска кода».

Запуск:
  venv/Scripts/python.exe agent.py "Напиши модуль работы с CSV: чтение, фильтр, агрегация"
  venv/Scripts/python.exe agent.py "FastAPI-приложение для заметок с SQLite"

Ключи:
  --code deepseek|lmstudio|local  кто пишет код (по умолчанию deepseek, если ключ есть)
  --judge deepseek|lmstudio|local  кто пишет тесты/судит (по умолчанию deepseek)
  --attempts N   число попыток (по умолчанию 5)
  --keep         не удалять песочницу
"""
import os
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_AGENT = ("Ты — senior-разработчик. По описанию напиши ПОЛНОЕ решение: "
                "весь нужный код в одном файле (solution.py) — функции, классы, "
                "обработка ошибок, понятные русские комментарии и докстринги. "
                "Не используй внешние зависимости сверх необходимого минимума. "
                "Без пояснений вне кода и без ```.")


def _pick(flag, default, available):
    if flag not in ("deepseek", "lmstudio", "local"):
        return default
    if flag == "deepseek" and not available[0]:
        return default
    if flag == "lmstudio" and not available[1]:
        return default
    return flag


def _gen(backend, msgs, max_tokens):
    if backend == "deepseek":
        from webui import _strip_fences
        return _strip_fences(deepseek_generate(msgs, max_tokens=max_tokens))
    if backend == "lmstudio":
        from webui import _strip_fences
        return _strip_fences(lmstudio_generate(msgs, max_tokens=max_tokens))
    # local: transformers + LoRA 3B
    from infer_smart import load, gen, _strip_fences as _sf
    if not hasattr(_gen, "_model"):
        _gen._model = load()
    tok, model = _gen._model
    return _sf(gen(tok, model, msgs, max_tokens=max_tokens))


def main():
    args = sys.argv[1:]
    code_backend, judge = "deepseek", "deepseek"
    attempts, keep = VERIFY_ATTEMPTS, False
    parts = []
    i = 0
    while i < len(args):
        a = args[i]
        if a == "--code" and i + 1 < len(args):
            code_backend = args[i + 1]; i += 2; continue
        if a == "--judge" and i + 1 < len(args):
            judge = args[i + 1]; i += 2; continue
        if a == "--attempts" and i + 1 < len(args):
            attempts = int(args[i + 1]); i += 2; continue
        if a == "--keep":
            keep = True; i += 1; continue
        parts.append(a); i += 1
    prompt = " ".join(parts)
    if not prompt:
        prompt = "Модуль работы с CSV: чтение файла, фильтрация строк, агрегация по колонке"

    avail = (deepseek_available(), lmstudio_available())
    code_backend = _pick(code_backend, "deepseek" if avail[0] else "local", avail)
    judge = _pick(judge, "deepseek" if avail[0] else "local", avail)

    sandbox = tempfile.mkdtemp(prefix="zuperai_agent_")
    print(f"[*] код: {code_backend}, судья: {judge}")

    # 1) судья пишет тесты по ТЗ
    t_msgs = [{"role": "system", "content": TEST_SYSTEM},
              {"role": "user",
               "content": f"Задача: {prompt}\n\nНапиши unittest-тесты, которые "
                          f"импортируют решение из solution и проверяют корректность "
                          f"всех заявленных возможностей. Верни ТОЛЬКО код тестов."}]
    tests = ensure_unittest_header(_gen(judge, t_msgs, 1536))

    # 2) код-агент пишет решение, гоняем, чиним
    msgs = [{"role": "system", "content": SYSTEM_AGENT},
            {"role": "user",
             "content": f"Задача: {prompt}\n\nТесты:\n{tests}\n\n"
                        f"Напиши полное решение (solution.py), проходящее эти тесты."}]
    code = ""
    for attempt in range(1, attempts + 1):
        if not code:
            code = _gen(code_backend, msgs, 1536)
            print(f"[попытка {attempt}/{attempts}] генерирую решение…", file=sys.stderr)
        else:
            print(f"[попытка {attempt}/{attempts}] тестирую исправленное…", file=sys.stderr)

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

        print(f"[попытка {attempt}] упало:\n{tail(out)}", file=sys.stderr)
        msgs.append({"role": "assistant", "content": code})
        msgs.append({"role": "user",
                     "content": f"Решение не прошло тесты:\n{tail(out)}\n\n"
                                f"Определи, где ошибка (в решении или в тесте), и верни "
                                f"исправленный КОД либо исправленный ТЕСТ — одно из двух."})
        fix = _gen(judge, msgs, 1536)
        if looks_like_tests(fix):
            tests = ensure_unittest_header(fix)
            code = ""
            msgs = msgs[:2]
            msgs[1] = {"role": "user",
                       "content": f"Задача: {prompt}\n\nТесты:\n{tests}\n\n"
                                  f"Напиши полное решение (solution.py), проходящее тесты."}
        else:
            code = fix
            msgs[-2] = {"role": "assistant", "content": code}

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


if __name__ == "__main__":
    main()