python/backend/services/rag/ingest_service.py - master - ksilisk/spbtechrun_hack

python/backend/services/rag/ingest_service.py
88 строк3 KB
Shaliko Salimov
code cleanup
01 дек 2025, 23:55
01 дек 2025, 23:556352613
from __future__ import annotations

import asyncio
import logging
from pathlib import Path
from typing import Iterable

from langchain_chroma import Chroma
from langchain_community.document_loaders import WebBaseLoader, TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter

from config import Settings


class RagIngestService:
    def __init__(
            self,
            *,
            settings: Settings,
            vectorstore: Chroma,
            logger: logging.Logger,
    ) -> None:
        self._settings = settings
        self._vectorstore = vectorstore
        self._logger = logger
        self._splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)

    async def ingest_urls(self, urls: list[str], source_tag: str | None = None) -> int:
        total_chunks = 0
        for url in urls:
            try:
                loader = WebBaseLoader(url)
                documents = await asyncio.get_running_loop().run_in_executor(None, loader.load)
            except Exception as exc:  # pragma: no cover - network interactions
                self._logger.error("Failed to load URL", extra={"url": url, "error": str(exc)})
                continue

            for document in documents:
                document.metadata.setdefault("source_url", url)
                document.metadata.setdefault("source_tag", source_tag or "url")

            chunks = self._splitter.split_documents(documents)
            self._vectorstore.add_documents(chunks)
            total_chunks += len(chunks)
            self._logger.info("Ingested URL", extra={"url": url, "chunks": len(chunks)})

        return total_chunks

    async def ingest_files(
            self,
            paths: Iterable[str],
            source_tag: str | None = None,
    ) -> int:
        """
        Ингестит локальные .md / .txt файлы в Chroma.
        paths — пути внутри контейнера/хоста.
        """
        total_chunks = 0

        for raw_path in paths:
            path = Path(raw_path)
            if not path.exists() or not path.is_file():
                self._logger.warning("File not found or not a file: %s", path)
                continue

            if path.suffix.lower() not in {".txt", ".md"}:
                self._logger.warning("Unsupported file type: %s", path)
                continue

            try:
                loader = TextLoader(str(path), encoding="utf-8")
                docs = loader.load()
            except Exception as exc:
                self._logger.exception("Failed to load file %s: %s", path, exc)
                continue

            for d in docs:
                d.metadata.setdefault("source", str(path))
                if source_tag:
                    d.metadata.setdefault("source_tag", source_tag)

            chunks = self._splitter.split_documents(docs)
            self._vectorstore.add_documents(chunks)
            total_chunks += len(chunks)

            self._logger.info("Ingested file %s as %d chunks", path, len(chunks))

        return total_chunks