python/backend/core/retrieval/rag_chain.py - master - ksilisk/spbtechrun_hack

python/backend/core/retrieval/rag_chain.py

61 строка2 KB

Shaliko Salimov

all prompts moving to the prompts directory and README.me refactoring

02 дек 2025, 21:18

02 дек 2025, 21:180511ade

from __future__ import annotations

import logging
from typing import Any

from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.retrieval import create_retrieval_chain
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.prompts import ChatPromptTemplate

from core.prompt.prompt_store import PromptStore
from schemas.facts import RagFacts

_RAG_CHAIN_PROMPT = "rag_chain.yaml"


def build_rag_qa_chain(vectorstore: Chroma, llm: BaseChatModel, prompt_store: PromptStore):
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", prompt_store.load(_RAG_CHAIN_PROMPT)),
            ("human", "{input}"),
        ]
    )

    combine_docs_chain = create_stuff_documents_chain(llm, prompt)

    retrieval_chain = create_retrieval_chain(
        retriever=vectorstore.as_retriever(),
        combine_docs_chain=combine_docs_chain
    )

    return retrieval_chain


class RagService:
    def __init__(self, qa_chain, logger: logging.Logger) -> None:
        self._qa_chain = qa_chain
        self._logger = logger

    async def ask(self, question: str) -> RagFacts:
        """
        Возвращает (ответ, мета), где мета содержит источники.
        """
        result: dict[str, Any] = await self._qa_chain.ainvoke({"input": question})
        # по умолчанию create_retrieval_chain кладёт в:
        # - result["answer"] — строка
        # - result["context"] — список Document
        answer: str = result.get("answer", "")
        docs: list[Document] = result.get("context") or []

        sources = [
            {
                "source": doc.metadata.get("source_url") or doc.metadata.get("source"),
                "snippet": doc.page_content,
            }
            for doc in docs
        ]

        return RagFacts(answer=answer, sources=sources)