RAG_helper/routers/query.py

import logging

from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.ext.asyncio import AsyncSession

from config import settings
from db.session import get_session
from models.requests import QueryRequest
from models.responses import QueryResponse, SourceInfo
from services import config_service, intent_document_service, intent_service
from services.llm_client import LLMClient
from services.rag_pipeline import rag_query

logger = logging.getLogger(__name__)

router = APIRouter(tags=["query"])

DEBUG_INTENT_CODE = "_debug"


@router.post("/query", response_model=QueryResponse)
async def query_rag(request: QueryRequest, session: AsyncSession = Depends(get_session)):
    from main import vectorstore_service

    if vectorstore_service is None:
        raise HTTPException(status_code=503, detail="Service not ready")

    if not settings.deepseek_api_key:
        raise HTTPException(status_code=500, detail="DEEPSEEK_API_KEY not configured")

    # Дефолт ветки — _debug (страница «Отладки»). При тесте из Настроек оператор передаёт
    # код выбранной ветки, и используются её промпт и подписки.
    intent_code = request.intent_code or DEBUG_INTENT_CODE
    intent = await intent_service.get_intent_by_code(session, intent_code)
    if intent is None:
        raise HTTPException(status_code=404, detail=f"Ветка {intent_code!r} не найдена.")

    # Системный промпт: если override задан — используем его (тест черновика из Настроек),
    # иначе — активный конфиг ветки. Если конфига нет и override пустой — 503.
    if request.system_prompt is not None:
        effective_system_prompt = request.system_prompt
        config_version = None
    else:
        active_cfg = await config_service.get_active_config_for_intent(session, intent.id)
        if active_cfg is None:
            raise HTTPException(
                status_code=503,
                detail=(
                    f"У ветки {intent_code!r} нет активной версии промпта. "
                    f"Зайдите в Настройки → выберите ветку → создайте и активируйте промпт."
                ),
            )
        effective_system_prompt = active_cfg.system_prompt
        config_version = active_cfg.version

    # document_ids: приоритет — явный параметр запроса. Иначе берём подписки ветки.
    # Для _debug дефолт пустой подписки — None (вся коллекция, удобство Отладки);
    # для пациентских и системных веток дефолт пустой — [] (= 0 чанков).
    if request.document_ids is not None:
        effective_doc_ids = request.document_ids
        subscribed_count = len(effective_doc_ids) if effective_doc_ids is not None else 0
    else:
        subscribed = await intent_document_service.list_documents_for_intent_code(
            session, intent_code,
        )
        if subscribed:
            effective_doc_ids = subscribed
        elif intent_code == DEBUG_INTENT_CODE:
            effective_doc_ids = None  # вся коллекция — только для _debug
        else:
            effective_doc_ids = []  # 0 чанков для остальных
        subscribed_count = len(subscribed)

    # disable_rag — пропускаем retrieval целиком (для веток без RAG, например _router).
    if request.disable_rag:
        effective_doc_ids = []  # пустой список → vectorstore.query вернёт []

    llm_client = LLMClient()
    try:
        result = await rag_query(
            vectorstore=vectorstore_service,
            llm_client=llm_client,
            question=request.text,
            top_k=request.top_k,
            document_ids=effective_doc_ids,
            temperature=request.temperature,
            max_tokens=request.max_tokens,
            system_prompt=effective_system_prompt,
        )
    except Exception as e:
        logger.exception("RAG query failed")
        raise HTTPException(status_code=500, detail=f"RAG query error: {e}")

    return QueryResponse(
        answer=result["answer"],
        sources=[SourceInfo(**s) for s in result["sources"]],
        model_used=result["model_used"],
        assembled_prompt=result.get("assembled_prompt", ""),
        intent_code=intent_code,
        config_version=config_version,
        rag_subscription={
            "subscribed_count": subscribed_count,
            "found_count": len(result["sources"]),
        },
    )