graph_worker.py

"""LangGraph-цикл: explore → generate → test → (исправление) → конец.

Phase B: агент работает внутри sandbox с полным проектом.
Использует нативный Tool Calling (bind_tools) вместо JSON-эмуляции.
"""

import json as _json
import os
import time as _time
from typing import TypedDict

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from langgraph.graph import END, StateGraph
from loguru import logger

from docker_manager import ProjectSandbox
from core.context_inspector import ProjectContextInspector
from tools.file_tools import list_files, make_coding_tools

load_dotenv()


# ── Lazy-инициализация LLM ──────────────────────────────────────

_flash_llm = None

def _get_llm() -> ChatOpenAI:
    """Возвращает экземпляр LLM (lazy singleton — не требует API key при импорте)."""
    global _flash_llm
    if _flash_llm is None:
        _flash_llm = ChatOpenAI(
            model="deepseek/deepseek-v4-flash",
            api_key=os.getenv("ROUTERAI_API_KEY"),
            base_url=os.getenv("ROUTERAI_BASE_URL", "https://routerai.ru/api/v1"),
            timeout=float(os.getenv("ARCH_CODE_LLM_TIMEOUT", "120")),
            max_retries=1,
        )
    return _flash_llm


def _invoke_llm(llm, messages, *, timeout: float = 120.0):
    """Вызвать LLM с ретраями на таймауты и пустые ответы.

    Исторический контекст: раньше таймаут реализовывался через
    ThreadPoolExecutor + future.cancel(). Это не работало: cancel() не
    может остановить уже запущенный поток, поэтому при таймауте LLM-вызов
    продолжал работать в фоне — двойной биллинг и stale-ответы.
    Теперь полагаемся на нативный timeout= у ChatOpenAI (httpx), а
    tenacity ретраит именно исключения таймаута httpx.

    Args:
        llm: Экземпляр ChatOpenAI (или bind_tools()).
        messages: Сообщения для invoke.
        timeout: Максимальное время ожидания ответа (сек).

    Returns:
        Ответ LLM (AIMessage / ChatResult).

    Raises:
        TimeoutError: Если LLM не ответил за timeout секунд
            (после всех ретраев).
    """
    import httpx
    from tenacity import (
        retry,
        retry_if_exception_type,
        stop_after_attempt,
        wait_exponential,
    )

    # Перехватываем только таймауты httpx + пустой ответ (ValueError).
    # Логические/валидационные ошибки не ретраим — они не исчезнут
    # от повторного вызова.
    _RETRYABLE = (httpx.TimeoutException, httpx.ReadTimeout, ValueError)

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=10),
        retry=retry_if_exception_type(_RETRYABLE),
        reraise=True,
    )
    def _call():
        response = llm.invoke(messages)
        # Пустой ответ — не таймаут, но тоже невалидный результат:
        # пробрасываем как ошибку, чтобы ретрай сработал.
        # ВАЖНО: ответ с tool_calls валиден даже при пустом content —
        # модель решила вызвать инструмент, а не ответить текстом.
        content = getattr(response, "content", None)
        has_tool_calls = bool(getattr(response, "tool_calls", None))
        if not has_tool_calls and (content is None or not str(content).strip()):
            raise ValueError("Empty LLM response")
        return response

    try:
        return _call()
    except (httpx.TimeoutException, ValueError) as exc:
        raise TimeoutError(
            f"LLM не ответил за {timeout}с "
            f"(RouterAI/DeepSeek V4 Flash) — таймаут"
        ) from exc


# ── Фильтр cp1251-несовместимых символов ─────────────────────────

def _sanitize(text: str) -> str:
    """Заменяет символы, не поддерживаемые cp1251, на '?'."""
    return text.encode("cp1251", errors="replace").decode("cp1251")


# ── Состояние графа ──────────────────────────────────────────────

class AgentState(TypedDict):
    task_id: str          # уникальный ID задачи (для изоляции sandbox)
    sandbox_dir: str      # абсолютный путь к sandbox/{task_id}/
    project_dir: str      # исходный проект (для справки)
    task: str             # описание задачи от Архитектора
    code: str             # итоговый код (для обратной совместимости)
    test_code: str        # тестовый скрипт (опционально)
    test_passed: bool     # тест пройден?
    error: str            # последняя ошибка выполнения
    iterations: int       # счётчик попыток
    success: bool         # финальный успех
    changed_files: list   # список изменённых файлов
    # ── Run Verifier (smoke-проверка приложения) ──────────────
    app_type: str         # "web_app" | "cli_script" | "skipped"
    skip_smoke_test: bool # True → пропустить run_app (микро-задачи)
    health_endpoint: str  # путь health-check для web_app
    health_port: int      # порт health-check для web_app
    # ── Контекст проекта (карта + AST-сигнатуры) ─────────────
    project_context: str  # сводка ProjectContextInspector (~2000 токенов)
    # ── Подробности для панели мониторинга ───────────────────
    thought_steps: list   # структурированные мысли агента
    action_steps: list    # структурированные действия
    chain_of_thought: str # свободная цепочка рассуждений
    # ── Billing / метрики ────────────────────────────────────
    prompt_tokens: int    # всего токенов промпта
    completion_tokens: int  # всего токенов генерации
    model: str            # название модели


# ── Утилита для таймстампов ──────────────────────────────────────


def _now_iso() -> str:
    """Текущее время в ISO-формате."""
    return _time.strftime("%Y-%m-%dT%H:%M:%S", _time.gmtime())


def _sum_tokens(response: object) -> dict:
    """Извлечь количество токенов из LLM-ответа.

    Returns:
        dict с prompt_tokens, completion_tokens.
    """
    try:
        md = getattr(response, "response_metadata", {}) or {}
        usage = md.get("token_usage", md.get("usage", {}))
        if isinstance(usage, dict):
            return {
                "prompt_tokens": usage.get("prompt_tokens", usage.get("input_tokens", 0)),
                "completion_tokens": usage.get("completion_tokens", usage.get("output_tokens", 0)),
            }
    except Exception as exc:
        logger.debug(f"Не удалось извлечь токены из ответа LLM: {exc}")
    return {"prompt_tokens": 0, "completion_tokens": 0}


# ── Узел: исследование проекта ──────────────────────────────────

def explore_project(state: AgentState):
    """Первый запуск: собрать карту проекта через ProjectContextInspector.

    Вместо LLM-вызова (5-10 секунд + токены) использует детерминированный
    инспектор: дерево каталогов + AST-сигнатуры классов/функций.
    Результат пишется в state["project_context"] и пробрасывается
    в system_prompt узла execute_actions.
    """
    sandbox_dir = state["sandbox_dir"]

    # Детерминированный сбор контекста (pathlib + ast, без bash и LLM)
    summary = ProjectContextInspector().generate_summary(sandbox_dir)

    ts = _now_iso()
    action_steps = state.get("action_steps", []) + [
        {"node": "explore", "action": f"Сбор карты проекта ({len(summary)} символов)", "status": "completed", "timestamp": ts},
    ]

    return {
        "error": "",
        "iterations": state["iterations"] + 1,
        "project_context": summary,
        "thought_steps": state.get("thought_steps", []),
        "action_steps": action_steps,
        "chain_of_thought": state.get("chain_of_thought", ""),
        "prompt_tokens": state.get("prompt_tokens", 0),
        "completion_tokens": state.get("completion_tokens", 0),
        "model": "deepseek/deepseek-v4-flash",
    }


# ── Узел: выполнение действий (чтение/запись файлов) ────────────

# Максимальное количество вызовов инструментов в одном узле
_MAX_TOOL_CALLS = 25

# Таймаут на один LLM-вызов (сек). RouterAI/DeepSeek V4 Flash иногда
# отвечает медленно; жёсткий таймаут не даёт задаче зависнуть навсегда.
LLM_TIMEOUT = float(os.getenv("ARCH_CODE_LLM_TIMEOUT", "120"))


def execute_actions(state: AgentState):
    """Агент читает и пишет файлы через нативный Tool Calling.

    Использует llm.bind_tools() — модель получает инструменты через API,
    а не через текстовый JSON-промпт. LangChain сам парсит tool_calls,
    Pydantic валидирует аргументы.
    """
    sandbox_dir = state["sandbox_dir"]
    task = state["task"]
    last_error = state["error"]

    # Инструменты с привязкой к sandbox_dir (partial)
    tools = make_coding_tools(sandbox_dir)
    llm = _get_llm().bind_tools(tools)

    error_context = (
        f"\nПредыдущая попытка упала с ошибкой:\n{last_error}\n"
        f"Исправь код и вернись в done().\n"
        if last_error else ""
    )

    # Карта проекта от explore-узла (дерево + AST-сигнатуры)
    project_context = state.get("project_context", "")
    context_block = (
        f"\n\n=== КАРТА ПРОЕКТА ===\n{project_context}\n=== КОНЕЦ КАРТЫ ===\n"
        if project_context.strip()
        else ""
    )

    # ── Формируем system prompt ────────────────────────────────
    system_prompt = (
        "Ты — AI-разработчик. Ты работаешь в fullstack-копии Python-проекта."
        f"{context_block}\n\n"
        "Твои инструменты:\n"
        "- read_file(relative_path) — прочитать файл\n"
        "- write_file(relative_path, content) — записать новый или перезаписать существующий файл\n"
        "- list_files(relative_path) — показать содержимое папки\n"
        "- done(changed_files) — завершить задачу (передай список созданных/изменённых файлов)\n\n"
        "Правила:\n"
        "1. Сначала изучи структуру и прочитай существующие файлы.\n"
        "2. Каждый новый файл пиши ПОЛНОСТЬЮ — не используй '...' или 'остальное без изменений'.\n"
        "3. После внесения всех изменений ВСЕГДА вызови done() с путями изменённых файлов.\n"
        "4. Не трогай .env, базы данных, .git/, venv/, __pycache__.\n"
        "5. Если нужно создать новые папки — write_file создаст их автоматически."
    )

    messages = [
        ("system", system_prompt),
        ("human", (
            f"Задача: {task}\n{error_context}"
            f"Изучи проект, прочитай нужные файлы и реализуй требуемые изменения. "
            f"По завершении вызови done()."
        )),
    ]

    changed_files = []
    action_steps = state.get("action_steps", [])
    thought_steps = state.get("thought_steps", [])
    chain_of_thought = state.get("chain_of_thought", "")
    prompt_tokens = state.get("prompt_tokens", 0)
    completion_tokens = state.get("completion_tokens", 0)

    for _ in range(_MAX_TOOL_CALLS):
        response = _invoke_llm(llm, messages, timeout=LLM_TIMEOUT)

        # Собираем токены с каждого LLM-вызова
        tokens = _sum_tokens(response)
        prompt_tokens += tokens["prompt_tokens"]
        completion_tokens += tokens["completion_tokens"]

        # ── Проверяем наличие tool_calls ──────────────────────
        if hasattr(response, "tool_calls") and response.tool_calls:
            # Добавляем AIMessage в историю (с tool_calls)
            messages.append(response)
            ts = _now_iso()

            for tc in response.tool_calls:
                tool_name = tc.get("name", "")
                tool_args = tc.get("args", {})

                # Инструмент done — завершение
                if tool_name == "done":
                    changed_files = tool_args.get("changed_files", [])
                    action_steps.append({
                        "node": "execute",
                        "action": f"Вызов done() → изменено файлов: {len(changed_files)}",
                        "status": "completed",
                        "timestamp": ts,
                    })
                    return {
                        "success": True,
                        "changed_files": changed_files,
                        "iterations": state["iterations"] + 1,
                        "thought_steps": thought_steps,
                        "action_steps": action_steps,
                        "chain_of_thought": chain_of_thought,
                        "prompt_tokens": prompt_tokens,
                        "completion_tokens": completion_tokens,
                        "model": "deepseek/deepseek-v4-flash",
                    }

                # Выполняем инструмент
                result_text = ""
                for t in tools:
                    if t.name == tool_name:
                        try:
                            result_text = t.invoke(tool_args)
                        except Exception as e:
                            result_text = f"❌ Ошибка вызова {tool_name}: {e}"
                        break

                if not result_text:
                    result_text = f"❌ Инструмент '{tool_name}' не найден"

                # Записываем действие
                action_summary = result_text[:200].replace("\n", " ")
                action_steps.append({
                    "node": "execute",
                    "action": f"{tool_name}({_fmt_args(tool_args)}) → {action_summary}",
                    "status": "completed" if not result_text.startswith("❌") else "failed",
                    "timestamp": ts,
                })

                messages.append(
                    ToolMessage(content=result_text, tool_call_id=tc["id"])
                )

            # После выполнения tool_calls — идём на следующий виток
            continue

        # ── Нет tool_calls — возможно DONE в тексте ──────────
        content = response.content or ""
        messages.append(AIMessage(content=content))
        if content.strip():
            ts = _now_iso()
            thought_steps.append({
                "node": "execute",
                "thought": content[:500],
                "timestamp": ts,
            })
            chain_of_thought += f"\n## Execute\n{content}\n"

        if "DONE" in content.upper() or "done" in content.lower():
            break

        messages.append(
            HumanMessage(
                content="Ты не вызвал done(). Если задача выполнена — вызови done() "
                        "с путями изменённых файлов. Если нет — продолжай работу."
            )
        )

    return {
        "success": True,
        "changed_files": changed_files,
        "iterations": state["iterations"] + 1,
        "thought_steps": thought_steps,
        "action_steps": action_steps,
        "chain_of_thought": chain_of_thought,
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "model": "deepseek/deepseek-v4-flash",
    }


def _fmt_args(args: dict) -> str:
    """Форматировать аргументы для лога (обрезает длинные значения)."""
    parts = []
    for k, v in args.items():
        v_str = str(v)
        if len(v_str) > 60:
            v_str = v_str[:57] + "..."
        parts.append(f"{k}={v_str}")
    return ", ".join(parts)


# ── Узел: тестирование в Docker ─────────────────────────────────

def run_tests(state: AgentState):
    """Запустить тесты проекта в Docker-контейнере."""
    sandbox_dir = state["sandbox_dir"]
    ts = _now_iso()
    action_steps = state.get("action_steps", [])

    # Определяем тип проекта
    has_requirements = os.path.exists(os.path.join(sandbox_dir, "requirements.txt"))
    has_package_json = os.path.exists(os.path.join(sandbox_dir, "package.json"))
    has_pyproject = os.path.exists(os.path.join(sandbox_dir, "pyproject.toml"))
    has_setup_py = os.path.exists(os.path.join(sandbox_dir, "setup.py"))

    if has_requirements or has_pyproject or has_setup_py:
        sandbox_type = "python"
    elif has_package_json:
        sandbox_type = "node"
    else:
        sandbox_type = "unknown"

    action_steps.append({
        "node": "test",
        "action": f"Тип проекта: {sandbox_type}. Запуск Docker-тестов...",
        "status": "running",
        "timestamp": ts,
    })

    # Устанавливаем зависимости и запускаем тесты
    try:
        result = ProjectSandbox.run_project_tests(sandbox_dir, sandbox_type, task_id=state["task_id"])
    except Exception as e:
        action_steps.append({
            "node": "test",
            "action": f"Ошибка Docker: {e}",
            "status": "failed",
            "timestamp": _now_iso(),
        })
        return {
            "success": False,
            "error": f"Ошибка Docker: {e}",
            "test_passed": False,
            "action_steps": action_steps,
        }

    if result["status"] == "success":
        action_steps.append({
            "node": "test",
            "action": "✅ Тесты пройдены",
            "status": "completed",
            "timestamp": _now_iso(),
        })
        return {"success": True, "error": "", "test_passed": True, "action_steps": action_steps}
    else:
        action_steps.append({
            "node": "test",
            "action": f"❌ Тесты не пройдены: {result['output'][:300]}",
            "status": "failed",
            "timestamp": _now_iso(),
        })
        return {
            "success": False,
            "error": result["output"],
            "test_passed": False,
            "action_steps": action_steps,
        }


# ── Узел: запуск приложения (smoke-проверка) ────────────────────

# Ключевые слова, указывающие на веб-приложение (а не CLI-скрипт)
_WEB_FRAMEWORK_KEYWORDS = (
    "fastapi", "flask", "uvicorn", "aiohttp", "starlette",
    "django", "express", "koa", "fastify", "hapijs",
)


def detect_app_type(sandbox_dir: str) -> str:
    """Определить тип приложения: web_app или cli_script.

    Читает main.py (или app.js) и ищет ПРИЗНАКИ веб-фреймворка:
      - импорты (import fastapi / from flask import / require('express'))
      - создание приложения (FastAPI(...), Flask(__name__), Express())
      - запуск сервера (uvicorn.run, app.run(...), app.listen(...))

    ВАЖНО: раньше матчились подстроки по всему файлу, включая комментарии
    и строки документации. Для ai-core (копируется в песочницу, когда
    project_dir=None) в main.py есть комментарий «Telemetry Server
    (FastAPI sidecar)» — это давало ложное web_app, uvicorn main:app
    падал, и задача помечалась failed, хотя код был корректный.
    Теперь эвристика консервативная: реальный импорт или вызов.

    Если файл точки входа не найден или фреймворк не обнаружен —
    считаем CLI-скриптом (консервативная эвристика).
    """
    import re

    entry_candidates = ["main.py", "app.py", "app.js", "server.js", "index.js"]
    main_content = ""
    for candidate in entry_candidates:
        path = os.path.join(sandbox_dir, candidate)
        if os.path.isfile(path):
            try:
                with open(path, "r", encoding="utf-8", errors="replace") as f:
                    main_content = f.read()
                break
            except Exception as exc:
                logger.debug(f"detect_app_type: не удалось прочитать {candidate}: {exc}")
                continue

    if not main_content:
        return "cli_script"

    # Убираем комментарии и docstring-подобные строки, чтобы не ловить
    # упоминания фреймворков в тексте (# ... / """ ... """).
    # Приблизительно: строки, начинающиеся с #, и тройные кавычки.
    no_comments = re.sub(
        r'"""[\s\S]*?"""|\'\'\'[\s\S]*?\'\'\'|#[^\n]*', "", main_content
    )
    lowered = no_comments.lower()

    # Признаки реального использования веб-фреймворка:
    # 1. Импорт: import fastapi / from fastapi import / import uvicorn
    if re.search(r"^\s*import\s+(fastapi|flask|uvicorn|aiohttp|starlette|django|falcon)", lowered, re.M):
        return "web_app"
    if re.search(r"^\s*from\s+(fastapi|flask|aiohttp|starlette|django)\s+import", lowered, re.M):
        return "web_app"
    # Node: require('express') / import express from
    if re.search(r"require\s*\(\s*['\"](express|koa|fastify|hapi|restify)", lowered):
        return "web_app"
    if re.search(r"from\s+(express|koa|fastify|hapi)\s+import", lowered):
        return "web_app"

    # 2. Создание приложения: FastAPI(...), Flask(__name__), express()
    if re.search(r"\b(fastapi|flask)\([^)]*\)", lowered):
        return "web_app"
    if re.search(r"\b(express|koa|fastify|hapi)\s*\(\s*\)", lowered):
        return "web_app"

    # 3. Запуск сервера: uvicorn.run, app.run(host=..., port=...), app.listen
    if re.search(r"uvicorn\.run|uvicorn\.Config|app\.run\s*\(|app\.listen\s*\(", lowered):
        return "web_app"

    return "cli_script"


def run_application(state: AgentState):
    """Запустить приложение в Docker и проверить, что оно работает.

    Двухфазный запуск (см. ProjectSandbox.run_application):
      1. install — зависимости с сетью
      2. run — приложение без сети (network_mode="none")

    Тип приложения определяется эвристикой detect_app_type():
      - web_app:   запуск сервера + HTTP-запрос к health-эндпоинту
      - cli_script: запуск с таймаутом + проверка exit code / Traceback

    При падении полный traceback записывается в state["error"], чтобы
    узел execute_actions на следующей итерации увидел точную причину.
    """
    sandbox_dir = state["sandbox_dir"]
    ts = _now_iso()
    action_steps = state.get("action_steps", [])

    # Флаг skip_smoke_test — микро-задачи без точки входа пропускают smoke
    if state.get("skip_smoke_test"):
        action_steps.append({
            "node": "run_app",
            "action": "⏭️ Smoke-тест пропущен (skip_smoke_test=True)",
            "status": "completed",
            "timestamp": ts,
        })
        return {
            "success": True,
            "error": "",
            "app_type": "skipped",
            "action_steps": action_steps,
        }

    # Определяем тип проекта (python / node)
    has_requirements = os.path.exists(os.path.join(sandbox_dir, "requirements.txt"))
    has_package_json = os.path.exists(os.path.join(sandbox_dir, "package.json"))
    has_pyproject = os.path.exists(os.path.join(sandbox_dir, "pyproject.toml"))
    has_setup_py = os.path.exists(os.path.join(sandbox_dir, "setup.py"))

    if has_requirements or has_pyproject or has_setup_py:
        project_type = "python"
    elif has_package_json:
        project_type = "node"
    else:
        project_type = "unknown"

    app_type = detect_app_type(sandbox_dir)
    state_app_type = app_type

    action_steps.append({
        "node": "run_app",
        "action": f"Тип проекта: {project_type}, приложение: {app_type}. "
                  f"Запуск smoke-проверки...",
        "status": "running",
        "timestamp": ts,
    })

    health_endpoint = state.get("health_endpoint", "/health")
    health_port = state.get("health_port", 8000)

    try:
        result = ProjectSandbox.run_application(
            sandbox_dir=sandbox_dir,
            project_type=project_type,
            app_type=app_type,
            task_id=state["task_id"],
            health_endpoint=health_endpoint,
            health_port=health_port,
        )
    except Exception as e:
        action_steps.append({
            "node": "run_app",
            "action": f"❌ Ошибка Docker при запуске: {e}",
            "status": "failed",
            "timestamp": _now_iso(),
        })
        return {
            "success": False,
            "error": f"Ошибка Docker при запуске приложения: {e}",
            "app_type": state_app_type,
            "action_steps": action_steps,
        }

    if result["status"] == "success":
        action_steps.append({
            "node": "run_app",
            "action": f"✅ Приложение работает ({app_type})",
            "status": "completed",
            "timestamp": _now_iso(),
        })
        return {
            "success": True,
            "error": "",
            "app_type": state_app_type,
            "action_steps": action_steps,
        }
    else:
        # Пробрасываем полный traceback/stderr в state["error"], чтобы
        # execute_actions на следующей итерации увидел точную причину.
        err_output = result.get("output", "")[:2000]
        action_steps.append({
            "node": "run_app",
            "action": f"❌ Приложение не работает: {err_output[:300]}",
            "status": "failed",
            "timestamp": _now_iso(),
        })
        return {
            "success": False,
            "error": f"Приложение не запустилось:\n{err_output}",
            "app_type": state_app_type,
            "action_steps": action_steps,
        }


# ── Маршрутизация ────────────────────────────────────────────────

def route_next_step(state: AgentState):
    if state["success"]:
        return END
    if state["iterations"] >= 3:
        return END
    return "execute"


def route_after_test(state: AgentState):
    """Роутинг после узла test:
    - если тесты не прошли → execute (исправление)
    - если skip_smoke_test=True → END (экономия 15-20 сек)
    - иначе → run_app (smoke-проверка запуска)
    """
    if not state.get("success"):
        return "execute"
    if state.get("skip_smoke_test"):
        return END
    return "run_app"


# ── Сборка графа ─────────────────────────────────────────────────

workflow = StateGraph(AgentState)

workflow.add_node("explore", explore_project)
workflow.add_node("execute", execute_actions)
workflow.add_node("test", run_tests)
workflow.add_node("run_app", run_application)

workflow.set_entry_point("explore")
workflow.add_edge("explore", "execute")
workflow.add_edge("execute", "test")
workflow.add_conditional_edges("test", route_after_test)
workflow.add_conditional_edges("run_app", route_next_step)

coding_graph = workflow.compile()