src/api/server.py - master - en_vui/ge

src/api/server.py
228 строк8 KB
Famel1x
init
11 фев 2026, 10:08
11 фев 2026, 10:081504a02
from contextlib import asynccontextmanager
from typing import Optional
from dataclasses import asdict
import time
import uuid
import logging 

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field 

from src.core.database import Base, engine
from src.services.genbi_service import GenBIService
from src.services.logger_service import UsageLogger
from src.core.logger import setup_logging 
from fastapi.responses import FileResponse
import os

# Настраиваем системное логирование
logger = setup_logging()

# Инициализация таблиц БД
Base.metadata.create_all(bind=engine)

class QueryRequest(BaseModel):
    text: str
    session_id: str = Field(default_factory=lambda: str(uuid.uuid4())) 
    
# Инициализация сервисов
service = GenBIService(interactive=False)
usage_logger = UsageLogger()

@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("🚀 Starting GenBI API Service...")
    service.setup_services()
    yield
    logger.info("🛑 Shutting down GenBI API Service...")

app = FastAPI(title="GenBI API", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.post("/query")
def query(req: QueryRequest):
    start_ts = time.time()
    request_id = str(uuid.uuid4())[:8] # ID для отладки конкретного запроса в логах
    
    logger.info(f"[{request_id}] Incoming request from session {req.session_id}: '{req.text}'")
    
    usage_acc = {"prompt": 0, "completion": 0, "total": 0}
    service.history.add_message(req.session_id, "user", req.text)

    try:
        # 1. Routing
        decision = service.router.route(req.text)
        logger.info(f"[{request_id}] Router decision: {decision.upper()}")
        
        # --- BLOCK BRANCH ---
        if decision == "block":
            censor_msg = service.get_censor_message()
            service.history.add_message(req.session_id, "assistant", censor_msg)
            
            logger.warning(f"[{request_id}] Request blocked by censor")
            
            duration = time.time() - start_ts
            usage_logger.log(
                session_id=req.session_id, 
                question=req.text, 
                sql_valid=False, 
                error="BLOCKED_BY_CENSOR",
                total_duration=duration,
                response_text=censor_msg
            )
            
            return {
                "success": False,
                "type": "block",
                "session_id": req.session_id,
                "response": censor_msg
            }

        # --- SQL BRANCH ---
        if decision == "sql":
            logger.info(f"[{request_id}] Executing SQL Agent flow...")
            result = service.sql_agent.execute(req.text)
            
            if result.success and result.payload:
                logger.info(f"[{request_id}] SQL Success. Rows fetched: {result.payload.row_count}")
                
                human_answer, usage = service.summarize_sql_result(
                    question=req.text,
                    sql_query=result.payload.query,
                    data=result.payload.data,
                    session_id=req.session_id
                )
                
                if usage:
                    usage_acc["prompt"] += usage.get("prompt_tokens", 0)
                    usage_acc["completion"] += usage.get("completion_tokens", 0)
                    usage_acc["total"] += usage.get("total_tokens", 0)
                
                duration = time.time() - start_ts
                
                usage_logger.log(
                    session_id=req.session_id,
                    question=req.text,
                    sql_query=result.payload.query,
                    sql_valid=result.success,
                    execution_time=result.payload.execution_time,
                    row_count=result.payload.row_count,
                    total_duration=duration,
                    response_text=human_answer,
                    prompt_tokens=usage_acc["prompt"],
                    completion_tokens=usage_acc["completion"],
                    total_tokens=usage_acc["total"]
                )

                logger.info(f"[{request_id}] Request completed in {duration:.2f}s")
                return {
                    "success": True,
                    "type": "sql",
                    "session_id": req.session_id,
                    "data": asdict(result.payload),
                    "response": human_answer
                }
            
            # Error handling
            logger.error(f"[{request_id}] SQL Agent failed: {result.error}")
            err_msg = f"Не удалось выполнить запрос: {result.error}"
            service.history.add_message(req.session_id, "assistant", err_msg)
            
            duration = time.time() - start_ts
            usage_logger.log(
                session_id=req.session_id,
                question=req.text,
                sql_valid=False, 
                error=result.error,
                total_duration=duration,
                response_text=err_msg
            )
            
            return {"success": False, "type": "error", "error": result.error}

        # --- CHAT BRANCH ---
        logger.info(f"[{request_id}] Executing Chat flow...")
        response_text, usage = service.chat_response(req.text, session_id=req.session_id)
        
        if usage:
            usage_acc["prompt"] += usage.get("prompt_tokens", 0)
            usage_acc["completion"] += usage.get("completion_tokens", 0)
            usage_acc["total"] += usage.get("total_tokens", 0)
        
        duration = time.time() - start_ts
        usage_logger.log(
            session_id=req.session_id,
            question=req.text,
            sql_valid=True,
            total_duration=duration,
            response_text=response_text,
            prompt_tokens=usage_acc["prompt"],
            completion_tokens=usage_acc["completion"],
            total_tokens=usage_acc["total"]
        )
        
        logger.info(f"[{request_id}] Request completed in {duration:.2f}s")
        return {
            "success": True, 
            "type": "chat", 
            "session_id": req.session_id, 
            "response": response_text
        }

    except Exception as e:
        logger.exception(f"[{request_id}] CRITICAL ERROR: {e}") # Логирует полный Traceback
        duration = time.time() - start_ts
        usage_logger.log(
            session_id=req.session_id, 
            question=req.text, 
            error=str(e),
            total_duration=duration
        )
        raise HTTPException(status_code=500, detail=str(e))
    
@app.get("/sessions")
def get_sessions():
    return service.history.list_sessions()

@app.get("/history/{session_id}")
def get_history(session_id: str):
    messages = service.history.get_session_messages(session_id)
    return {"session_id": session_id, "messages": messages}

@app.get("/logs/export")
def export_logs():
    logger.info("Exporting logs to Excel...")
    excel_file = usage_logger.get_excel_file()
    return StreamingResponse(
        excel_file,
        media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        headers={"Content-Disposition": "attachment; filename=genbi_logs.xlsx"}
    )

@app.get("/logs/json")
def get_logs_json():
    df = usage_logger.get_logs_df()
    df = df.fillna("") 
    data = df.to_dict(orient="records")
    return data

@app.get("/logs/system")
def get_system_log():
    """Скачивание полного технического лога (txt)"""
    log_path = "logs/genbi_system.log"
    if os.path.exists(log_path):
        return FileResponse(
            log_path, 
            media_type="text/plain", 
            filename="genbi_debug_trace.log"
        )
    raise HTTPException(status_code=404, detail="Log file not found")