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

src/localscript_api/main.py
191 строка6 KB
Якуб
some fixes of api server
13 апр 2026, 11:24
13 апр 2026, 11:246827c5f
"""LocalScript FastAPI Server.

Запуск:
    uv run uvicorn src.localscript_api.main:app --host 0.0.0.0 --port 8080

Env vars:
    MODEL_NAME: Модель LLM (по умолчанию: qwen3.5:latest)
    OPENAI_BASE_URL: Базовый URL LLM API
    OPENAI_API_KEY: API ключ
    MAX_QA_ROUNDS: Максимум clarification раундов (по умолчанию: 3)
    MAX_OUTPUT_TOKENS: Максимум токенов на вывод (по умолчанию: 4096)
"""

from __future__ import annotations

import logging
import os
import uuid
from contextlib import asynccontextmanager
from typing import AsyncGenerator

from fastapi import Depends, FastAPI, HTTPException
from pydantic import BaseModel

from sse_starlette.sse import EventSourceResponse

from src.localscript_api.api import router
from src.localscript_api.engine import LocalScriptEngine
from src.localscript_api.schemas import GenerateRequest, GenerateResponse
from src.localscript_api.sessions import SessionManager
from src.localscript_api.sse import format_sse_event

# Настройка логов
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
)
logger = logging.getLogger("localscript")


class AnswerRequest(BaseModel):
    """Ответ на clarification вопрос."""
    answer: str


class SessionResponse(BaseModel):
    """Информация о сессии."""
    session_id: str


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
    """Lifespan: инициализация и очистка engine."""
    model = os.getenv("MODEL_NAME", "qwen3.5:latest")
    max_qa = int(os.getenv("MAX_QA_ROUNDS", "3"))
    max_tokens = int(os.getenv("MAX_OUTPUT_TOKENS", "4096"))
    base_url = os.getenv("OPENAI_BASE_URL", "http://localhost:1235/v1")

    session_manager = SessionManager()

    logger.info(f"Initializing engine with model={model}, base_url={base_url}, max_qa={max_qa}, max_tokens={max_tokens}")

    engine = LocalScriptEngine(
        model=model,
        max_qa_rounds=max_qa,
        max_output_tokens=max_tokens,
        openai_base_url=base_url,
        session_manager=session_manager,
    )
    await engine.initialize()
    app.state.engine = engine
    app.state.session_manager = session_manager

    logger.info("LocalScript engine initialized")
    yield

    # Cleanup
    if hasattr(app.state, "engine"):
        await app.state.engine.cleanup()
        logger.info("LocalScript engine cleaned up")


app = FastAPI(
    title="LocalScript API",
    version="1.0.0",
    lifespan=lifespan,
)
app.include_router(router)


def _get_engine() -> LocalScriptEngine:
    """Dependency для получения engine."""
    return app.state.engine


def _get_session_manager() -> SessionManager:
    """Dependency для получения session manager."""
    return app.state.session_manager


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------


@app.post("/generate", response_model=GenerateResponse, tags=["Code Generation"])
async def generate_code(
    req: GenerateRequest,
    engine: LocalScriptEngine = Depends(_get_engine),
) -> GenerateResponse:
    """Сгенерировать код по тексту задачи.

    Полный цикл: LangGraph → Lua validation → {code: "..."}
    Clarification: если LLM задаёт вопрос → автоматически NO CLARIFICATION.
    """
    result = await engine.generate_code(
        prompt=req.prompt,
        language=req.language,
        save_path=req.save_path,
    )

    if result.error:
        logger.error(f"Generation error: {result.error}")
        return GenerateResponse(code="", saved_path=result.saved_path)

    return GenerateResponse(code=result.code, saved_path=result.saved_path)


@app.post("/generate/stream", tags=["Code Generation"])
async def generate_code_stream(
    req: GenerateRequest,
    engine: LocalScriptEngine = Depends(_get_engine),
    session_mgr: SessionManager = Depends(_get_session_manager),
):
    """Streaming генерация через SSE.

    Events: question, answer, validating, code, error, timeout

    Clarification: если LLM задаёт вопрос → создаётся сессия →
    клиент шлёт POST /session/{id}/answer → генерация продолжается.
    """
    session_id = str(uuid.uuid4())[:8]
    session = session_mgr.create(session_id)

    async def event_generator():
        try:
            yield format_sse_event(
                event_type="session_created",
                data=session_id,
                meta={"session_id": session_id},
            )
            async for event in engine.generate_code_stream(
                prompt=req.prompt,
                language=req.language,
                session=session,
                save_path=req.save_path,
            ):
                yield format_sse_event(
                    event_type=event["type"],
                    data=event["data"],
                    meta=event.get("meta"),
                )
        finally:
            session_mgr.finish(session_id)

    return EventSourceResponse(event_generator())


@app.post("/session/{session_id}/answer", response_model=SessionResponse, tags=["Sessions"])
async def submit_answer(
    session_id: str,
    req: AnswerRequest,
    session_mgr: SessionManager = Depends(_get_session_manager),
) -> SessionResponse:
    """Отправить ответ на clarification вопрос.

    Вызывается клиентом после получения SSE event типа "question".
    """
    session = session_mgr.get(session_id)
    if session is None:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    await session.answer_queue.put(req.answer)
    logger.info(f"Answer submitted for session {session_id}")
    return SessionResponse(session_id=session_id)


@app.get("/health", tags=["Health"])
async def health():
    """Проверка здоровья."""
    return {"status": "ok"}