python/backend/api/v1/routes_rag.py - master - ksilisk/spbtechrun_hack

python/backend/api/v1/routes_rag.py

74 строки2 KB

Shaliko Salimov

implementation of the MCP server and refactoring of the previous intent schema

29 ноя 2025, 19:05

29 ноя 2025, 19:0546519e0

from __future__ import annotations

from pathlib import Path
from typing import List

from fastapi import APIRouter, Depends, UploadFile, File, Form, HTTPException
from pydantic import BaseModel, Field

from config import Settings, get_settings
from infrastructure.logging.logger import get_logger
from infrastructure.vectorstore.chroma_store import get_chroma_vectorstore
from services.rag.ingest_service import RagIngestService

router = APIRouter(prefix="/v1/rag", tags=["rag"])


class RagIngestRequest(BaseModel):
    urls: list[str] = Field(default_factory=list)
    source_tag: str | None = None


def get_rag_ingest_service(
        settings: Settings = Depends(get_settings),
) -> RagIngestService:
    vectorstore = get_chroma_vectorstore(settings)
    logger = get_logger("rag_ingest")
    return RagIngestService(settings=settings, vectorstore=vectorstore, logger=logger)


@router.post("/ingest")
async def ingest(
        payload: RagIngestRequest,
        service: RagIngestService = Depends(get_rag_ingest_service),
) -> dict:
    count = await service.ingest_urls(payload.urls, source_tag=payload.source_tag)
    return {"ingested_chunks": count}


@router.post("/upload")
async def upload_and_ingest_files(
        files: List[UploadFile] = File(...),
        source_tag: str | None = Form(default=None),
        service: RagIngestService = Depends(get_rag_ingest_service),
):
    """
    Принимает файлы (md/txt) через multipart/form-data, сохраняет их в rag_upload_dir
    и сразу инжестит в Chroma.
    """
    settings = get_settings()
    upload_dir = Path(settings.rag_upload_dir)
    upload_dir.mkdir(parents=True, exist_ok=True)

    saved_paths: list[str] = []

    for upload in files:
        suffix = Path(upload.filename or "").suffix.lower()
        if suffix not in {".txt", ".md"}:
            raise HTTPException(
                status_code=400,
                detail=f"Unsupported file type: {suffix}. Allowed: .txt, .md",
            )

        target_path = upload_dir / upload.filename

        content = await upload.read()
        target_path.write_bytes(content)
        saved_paths.append(str(target_path))

    ingested = await service.ingest_files(saved_paths, source_tag=source_tag)

    return {
        "uploaded_files": saved_paths,
        "ingested_chunks": ingested,
    }