simple_rag_chain.py

"""
간단한 RAG 체인 구현 (디버깅용)
"""
import os
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough


class SimpleRAGChain:
    def __init__(self, vector_store):
        """간단한 RAG 체인 초기화"""
        print("간단한 RAG 체인 초기화 중...")
        self.vector_store = vector_store

        # OpenAI API 키 확인
        openai_api_key = os.environ.get("OPENAI_API_KEY", "")
        print(f"API 키 설정됨: {bool(openai_api_key)}")

        # OpenAI 모델 초기화
        self.llm = ChatOpenAI(
            model_name="gpt-3.5-turbo",
            temperature=0.2,
            api_key=openai_api_key,
        )

        # 프롬프트 템플릿
        template = """
        다음 정보를 기반으로 질문에 정확하게 답변해주세요.

        질문: {question}

        참고 정보:
        {context}

        참고 정보에 답이 없는 경우 "제공된 문서에서 해당 정보를 찾을 수 없습니다."라고 답변하세요.
        """

        self.prompt = PromptTemplate.from_template(template)

        # 체인 구성
        self.chain = (
                {"context": self._retrieve, "question": RunnablePassthrough()}
                | self.prompt
                | self.llm
                | StrOutputParser()
        )
        print("간단한 RAG 체인 초기화 완료")

    def _retrieve(self, query):
        """문서 검색"""
        try:
            docs = self.vector_store.similarity_search(query, k=3)
            return "\n\n".join(doc.page_content for doc in docs)
        except Exception as e:
            print(f"검색 중 오류: {e}")
            return "문서 검색 중 오류가 발생했습니다."

    def run(self, query):
        """쿼리 처리"""
        try:
            return self.chain.invoke(query)
        except Exception as e:
            print(f"실행 중 오류: {e}")
            return f"오류 발생: {str(e)}"