rag_chain.py

"""
LangChain을 활용한 RAG 체인 구현
"""
from typing import List, Dict, Any
from langchain.schema import Document
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_community.chat_models import ChatOllama
from langchain_openai import ChatOpenAI

from config import (
    OLLAMA_HOST, LLM_MODEL, USE_OPENAI,
    OPENAI_API_KEY, TOP_K_RETRIEVAL, TOP_K_RERANK
)
from vector_store import VectorStore
from reranker import Reranker


class RAGChain:
    def __init__(self, vector_store: VectorStore, use_reranker: bool = True):
        """
        RAG 체인 초기화 (환경에 따른 LLM 선택)

        Args:
            vector_store: 벡터 스토어 인스턴스
            use_reranker: 리랭커 사용 여부
        """
        try:
            print("RAGChain 초기화 시작...")
            self.vector_store = vector_store
            self.use_reranker = use_reranker
            print(f"리랭커 사용 여부: {use_reranker}")

            if use_reranker:
                try:
                    self.reranker = Reranker()
                    print("리랭커 초기화 성공")
                except Exception as e:
                    print(f"리랭커 초기화 실패: {str(e)}")
                    self.reranker = None
                    self.use_reranker = False
            else:
                self.reranker = None

            # 환경에 따른 LLM 모델 설정
            if USE_OPENAI or IS_HUGGINGFACE:
                print(f"OpenAI 모델 초기화: {LLM_MODEL}")
                print(f"API 키 존재 여부: {'있음' if OPENAI_API_KEY else '없음'}")
                try:
                    self.llm = ChatOpenAI(
                        model_name=LLM_MODEL,
                        temperature=0.2,
                        api_key=OPENAI_API_KEY,
                    )
                    print("OpenAI 모델 초기화 성공")
                except Exception as e:
                    print(f"OpenAI 모델 초기화 실패: {str(e)}")
                    raise
            else:
                try:
                    print(f"Ollama 모델 초기화: {LLM_MODEL}")
                    self.llm = ChatOllama(
                        model=LLM_MODEL,
                        temperature=0.2,
                        base_url=OLLAMA_HOST,
                    )
                    print("Ollama 모델 초기화 성공")
                except Exception as e:
                    print(f"Ollama 모델 초기화 실패: {str(e)}")
                    raise

            # RAG 체인 구성 및 프롬프트 설정
            print("RAG 체인 설정 시작...")
            self.setup_chain()
            print("RAG 체인 설정 완료")
        except Exception as e:
            print(f"RAGChain 초기화 중 상세 오류: {str(e)}")
            import traceback
            traceback.print_exc()
            raise

    def setup_chain(self) -> None:
        """
        RAG 체인 및 프롬프트 설정
        """
        # 프롬프트 템플릿 정의
        template = """
        다음 정보를 기반으로 질문에 정확하게 답변해주세요.

        질문: {question}

        참고 정보:
        {context}

        참고 정보에 답이 없는 경우 "제공된 문서에서 해당 정보를 찾을 수 없습니다."라고 답변하세요.
        답변은 정확하고 간결하게 제공하되, 참고 정보에서 근거를 찾아 설명해주세요.
        참고 정보의 출처도 함께 알려주세요.
        """

        self.prompt = PromptTemplate.from_template(template)

        # RAG 체인 정의
        self.chain = (
                {"context": self._retrieve, "question": RunnablePassthrough()}
                | self.prompt
                | self.llm
                | StrOutputParser()
        )

    def _retrieve(self, query: str) -> str:
        """
        쿼리에 대한 관련 문서 검색 및 컨텍스트 구성

        Args:
            query: 사용자 질문

        Returns:
            검색 결과를 포함한 컨텍스트 문자열
        """
        # 벡터 검색 수행
        docs = self.vector_store.similarity_search(query, k=TOP_K_RETRIEVAL)

        # 리랭커 적용 (선택적)
        if self.use_reranker and docs:
            docs = self.reranker.rerank(query, docs, top_k=TOP_K_RERANK)

        # 검색 결과 컨텍스트 구성
        context_parts = []
        for i, doc in enumerate(docs, 1):
            source = doc.metadata.get("source", "알 수 없는 출처")
            page = doc.metadata.get("page", "")
            source_info = f"{source}"
            if page:
                source_info += f" (페이지: {page})"

            context_parts.append(f"[참고자료 {i}] - 출처: {source_info}\n{doc.page_content}\n")

        return "\n".join(context_parts)

    def run(self, query: str) -> str:
        """
        사용자 쿼리에 대한 RAG 파이프라인 실행

        Args:
            query: 사용자 질문

        Returns:
            모델 응답 문자열
        """
        return self.chain.invoke(query)