src/localscript_api/engine.py - master - Watashicuvu/agentic-tools

src/localscript_api/engine.py
423 строки17 KB
Якуб
cli fixes
14 апр 2026, 18:32
14 апр 2026, 18:326202d02
"""LocalScript Engine — обёртка над LangGraph + Lua validation.

Core-логика:
1. LangGraph генерация кода (use_langgraph=True)
2. Lua-валидация (syntax + lint + run) перед возвратом
3. Clarification handling (IronUser-подобная логика)
4. Retry при syntax/lint/run error (MAX_RETRIES=2)
"""

from __future__ import annotations

import asyncio
import logging
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import AsyncGenerator, Optional

from src.eval.langgraph_eval_shim import EvalLangGraphShim
from src.eval.lua_evaluator import LuaCodeEvaluator
from src.eval.models import EvalRole, EvalTask
from src.localscript_api.schemas import SseEventType
from src.localscript_api.sessions import Session, SessionManager
from src.swarm.langgraph.langchain_llm_client import LangChainLLMClient

logger = logging.getLogger(__name__)

MAX_RETRIES = 3
MAX_CLARIFICATION_ROUNDS = 3


@dataclass
class GenerationResult:
    """Результат генерации кода."""
    code: str
    error: Optional[str] = None
    clarifications: Optional[list[dict]] = None
    validation_ok: bool = False
    validation_output: Optional[str] = None
    saved_path: Optional[str] = None


class LocalScriptEngine:
    """Основной движок LocalScript API.

    Args:
        model: Имя модели LLM
        max_qa_rounds: Максимум clarification раундов
        max_output_tokens: Максимум токенов на вывод
    """

    def __init__(
        self,
        model: str,
        max_qa_rounds: int = MAX_CLARIFICATION_ROUNDS,
        max_output_tokens: int = 4096,
        openai_base_url: Optional[str] = None,
        session_manager: Optional[SessionManager] = None,
    ):
        self.model = model
        self.max_qa_rounds = max_qa_rounds
        self.max_output_tokens = max_output_tokens
        self.openai_base_url = openai_base_url
        self.session_manager = session_manager or SessionManager()

        self._langchain_client: Optional[LangChainLLMClient] = None
        self._shim: Optional[EvalLangGraphShim] = None

    async def initialize(self):
        """Инициализация: создание LLM клиента и LangGraph shim."""
        self._langchain_client = LangChainLLMClient(
            model=self.model,
            base_url=self.openai_base_url,
        )
        self._shim = EvalLangGraphShim(
            smart_client=self._langchain_client,
            model=self.model,
            mcp_url=None,  # Без MCP — только локальные инструменты
            max_qa_rounds=self.max_qa_rounds,
            max_output_tokens=self.max_output_tokens,
        )

    async def cleanup(self):
        """Очистка ресурсов."""
        if self._shim and self._shim._tool_orchestrator:
            await self._shim._tool_orchestrator.disconnect()
            self._shim._tool_orchestrator = None

    async def generate_code(
        self,
        prompt: str,
        language: str = "lua",
        save_path: Optional[str] = None,
    ) -> GenerationResult:
        """Сгенерировать код с валидацией.

        Args:
            prompt: Текст задачи
            language: Язык (по умолчанию lua)

        Returns:
            GenerationResult с кодом или ошибкой
        """
        if language.lower() != "lua":
            return GenerationResult(
                code="",
                error=f"Language '{language}' not supported. Only 'lua' is available.",
            )

        task = self._create_task(prompt, language, skip_auto_clarification=True)

        # Retry-loop: генерация + валидация
        last_error = None
        last_code = ""
        clarification_handled = False

        for attempt in range(MAX_RETRIES + 1):
            if attempt > 0 and last_error:
                logger.info(f"[LOCALSCRIPT] Retry {attempt + 1}/{MAX_RETRIES + 1}: {last_error[:100]}")
                task = EvalTask(
                    id=task.id,
                    role=task.role,
                    description=task.description,
                    expected_behavior=task.expected_behavior,
                    reference_code=task.reference_code,
                    context=task.context,
                    requires_clarification=task.requires_clarification,
                    error_feedback=(
                        f"\n\n⚠️ Попытка {attempt + 1}: Предыдущая ошибка:\n{last_error}\n\nИСПРАВЬ!"
                    ),
                )

            # Генерация через LangGraph
            if self._shim:
                result = await self._shim.run_task(task)
            last_code = result.code or ""

            if result.error:
                last_error = f"LangGraph error: {result.error}"
                continue

            # === CHECK FOR QUESTION: проверяем code И full_response ===
            question_text = None
            if last_code.strip().startswith("QUESTION:"):
                question_text = last_code[len("QUESTION:"):].strip()
            elif result.full_response and "QUESTION:" in result.full_response:
                # LLM кладёт вопрос в full_response, а code — пустой
                idx = result.full_response.find("QUESTION:")
                question_text = result.full_response[idx + len("QUESTION:"):].strip().split("\n")[0]
                last_code = ""  # code пуст — не считаем его результатом

            if question_text is not None:
                if not clarification_handled:
                    logger.info(f"[LOCALSCRIPT] LLM asked: {question_text[:100]} — responding NO CLARIFICATION")
                    feedback = (
                        f"\n\nNO CLARIFICATION — пользователь не может ответить на вопрос."
                        f"\nВопрос был: {question_text}"
                        f"\n\nСгенерируй ФИНАЛЬНЫЙ КОД без вопросов. Если не хватает информации — "
                        f"используй разумные предположения по умолчанию."
                    )
                    task = EvalTask(
                        id=task.id,
                        role=task.role,
                        description=task.description,
                        expected_behavior=task.expected_behavior,
                        reference_code=task.reference_code,
                        context=task.context,
                        requires_clarification=task.requires_clarification,
                        error_feedback=feedback,
                    )
                    clarification_handled = True
                    continue
                else:
                    last_error = "LLM asked a question again after NO CLARIFICATION"
                    continue
            # ============================================================

            # Lua-валидация
            valid, validation_output = await self._validate_lua(last_code)
            if not valid:
                last_error = validation_output
                continue

            # Успех
            if save_path and save_path.lower() in ('string', 'str'):
                save_path = None
            saved_path = self._save_code(last_code, save_path, task.id)
            return GenerationResult(
                code=last_code,
                validation_ok=True,
                validation_output=validation_output,
                clarifications=[],
                saved_path=saved_path,
            )

        # Все попытки исчерпаны
        saved_path = self._save_code(last_code, save_path, task.id) if last_code else None
        return GenerationResult(
            code=last_code,
            error=last_error or "All retry attempts failed",
            validation_ok=False,
            validation_output=last_error,
            saved_path=saved_path,
        )

    async def generate_code_stream(
        self,
        prompt: str,
        language: str = "lua",
        session: Optional[Session] = None,
        save_path: Optional[str] = None,
    ) -> AsyncGenerator[dict, None]:
        """Streaming генерация кода через SSE.

        Args:
            prompt: Текст задачи
            language: Язык (по умолчанию lua)
            session: Сессия для интерактивных clarification

        Yields:
            Словари с type и data для SSE
        """
        if language.lower() != "lua":
            yield {"type": "error", "data": f"Language '{language}' not supported"}
            return

        task = self._create_task(prompt, language, skip_auto_clarification=True)
        last_error = None
        last_code = ""
        questions_asked = 0

        for attempt in range(MAX_RETRIES + 1):
            if attempt > 0 and last_error:
                task = EvalTask(
                    id=task.id,
                    role=task.role,
                    description=task.description,
                    expected_behavior=task.expected_behavior,
                    reference_code=task.reference_code,
                    context=task.context,  # Сохраняем _skip_auto_clarification
                    requires_clarification=task.requires_clarification,
                    error_feedback=(
                        f"\n\n⚠️ Попытка {attempt + 1}: {last_error[:200]}\n\nИСПРАВЬ!"
                    ),
                )

            # Генерация
            if self._shim:
                result = await self._shim.run_task(task)
            last_code = result.code or ""

            if result.error:
                last_error = f"LangGraph error: {result.error}"
                continue

            # === CHECK FOR QUESTION: проверяем code И full_response ===
            question_text = None
            if last_code.strip().startswith("QUESTION:"):
                question_text = last_code[len("QUESTION:"):].strip()
            elif result.full_response and "QUESTION:" in result.full_response:
                idx = result.full_response.find("QUESTION:")
                question_text = result.full_response[idx + len("QUESTION:"):].strip().split("\n")[0]
                last_code = ""

            if question_text is not None:
                questions_asked += 1

                if session and questions_asked <= self.max_qa_rounds:
                    # Ждём ответ через сессию
                    logger.info(f"[STREAM] LLM asked: {question_text[:100]} — sending question event")
                    yield {
                        "type": SseEventType.QUESTION.value,
                        "data": question_text,
                        "meta": {"round": questions_asked},
                    }
                    await asyncio.sleep(0)  # Flush SSE event to client
                    logger.info(f"[STREAM] Waiting for answer (timeout=120s)...")
                    answer = await self.session_manager.wait_for_answer(session)

                    if answer is None:
                        # Таймаут — NO CLARIFICATION
                        yield {"type": "timeout", "data": "No answer received, regenerating without clarification"}
                        feedback = (
                            "\n\nNO CLARIFICATION — таймаут ожидания ответа."
                            "\nСгенерируй код без вопросов, используя разумные предположения."
                        )
                    else:
                        yield {"type": SseEventType.ANSWER.value, "data": answer}
                        feedback = (
                            f"\n\nОтвет пользователя: {answer}"
                            f"\n\nТеперь сгенерируй ФИНАЛЬНЫЙ КОД (без вопросов)!"
                        )

                    task = EvalTask(
                        id=task.id,
                        role=task.role,
                        description=task.description,
                        expected_behavior=task.expected_behavior,
                        reference_code=task.reference_code,
                        context=task.context,
                        requires_clarification=task.requires_clarification,
                        error_feedback=feedback,
                    )
                    continue
                else:
                    # Нет сессии или превышен лимит — NO CLARIFICATION
                    yield {"type": "info", "data": f"No clarification available (question: {question_text[:100]})"}
                    last_error = "LLM asked question but no session or max rounds exceeded"
                    continue
            # ==========================

            # Валидация
            yield {"type": SseEventType.VALIDATING.value, "data": "Checking syntax, lint, runtime..."}
            valid, validation_output = await self._validate_lua(last_code)

            if not valid:
                last_error = validation_output
                continue

            # Успех
            if save_path and save_path.lower() in ('string', 'str'):
                save_path = None
            saved_path = self._save_code(last_code, save_path, task.id)
            yield {
                "type": SseEventType.CODE.value,
                "data": last_code,
                "meta": {"validation": "ok", "saved_path": saved_path},
            }
            return

        # Все попытки исчерпаны
        yield {"type": "error", "data": last_error or "All retry attempts failed"}

    async def _validate_lua(self, code: str) -> tuple[bool, Optional[str]]:
        """Валидация Lua-кода: syntax + lint + run.

        Returns:
            (True, None) если код валиден
            (False, error_output) если ошибка
        """
        if not code.strip():
            return False, "Generated code is empty"

        evaluator = LuaCodeEvaluator(timeout=30)

        with tempfile.TemporaryDirectory() as temp_dir:
            result = await evaluator.evaluate(
                code=code,
                work_dir=temp_dir,
                run_code=True,
            )

        if not result.code_syntax_ok:
            return False, f"Syntax error:\n{result.code_syntax_output}"

        if not result.code_lint_ok:
            lint_errors = "\n".join(result.code_lint_errors[:5])
            if "error" in lint_errors.lower():
                return False, f"Lint errors:\n{lint_errors}"

        if not result.code_run_ok:
            return False, f"Runtime error:\n{result.code_run_output}"

        return True, None

    def _create_task(self, prompt: str, language: str, skip_auto_clarification: bool = False) -> EvalTask:
        """Создаёт EvalTask из prompt."""
        import hashlib

        task_id = f"localscript_{hashlib.md5(prompt.encode()).hexdigest()[:8]}"

        context = {}
        if skip_auto_clarification:
            context["_skip_auto_clarification"] = True

        return EvalTask(
            id=task_id,
            role=EvalRole.EXECUTOR,
            description=prompt,
            expected_behavior=f"Корректный {language.upper()} код для задачи: {prompt}",
            requires_clarification=True,
            reference_code=None,
            context=context if context else None,
            error_feedback=None,
        )

    def _save_code(self, code: str, save_path: Optional[str], task_id: str) -> Optional[str]:
        """Сохранить сгенерированный код.

        Args:
            code: Код для сохранения.
            save_path: Путь к директории или файлу.
            task_id: ID задачи для имени файла по умолчанию.

        Returns:
            Путь к сохранённому файлу или None.
        """
        if not code:
            return None

        try:
            if save_path:
                path = Path(save_path)
                # Если путь — директория, создаём имя файла
                if path.is_dir() or not path.suffix:
                    path.mkdir(parents=True, exist_ok=True)
                    path = path / f"{task_id}.lua"
                else:
                    path.parent.mkdir(parents=True, exist_ok=True)
                path.write_text(code, encoding="utf-8")
                logger.info(f"Code saved to: {path}")
                return str(path)
            else:
                # Временная директория
                with tempfile.NamedTemporaryFile(
                    mode="w", suffix=".lua", delete=False, prefix=f"{task_id}_"
                ) as f:
                    f.write(code)
                logger.info(f"Code saved to temp: {f.name}")
                return f.name
        except Exception as e:
            logger.warning(f"Failed to save code: {e}")
            return None