simple_rag_chain.py

"""
간단한 RAG 체인 구현 (디버깅용) - 직접 DeepSeek API 호출 방식
"""
import os
import logging
import time
from typing import Dict, Any, List

# 직접 DeepSeek 클라이언트 사용
from direct_deepseek import DirectDeepSeekClient

# 로깅 설정
logger = logging.getLogger("SimpleRAGChain")

class SimpleRAGChain:
    def __init__(self, vector_store, api_key=None, model="deepseek-chat", endpoint=None):
        """간단한 RAG 체인 초기화"""
        logger.info("간단한 RAG 체인 초기화 중...")
        self.vector_store = vector_store

        # DeepSeek API 키 확인
        self.api_key = api_key or os.environ.get("DEEPSEEK_API_KEY", "")
        self.model = model or os.environ.get("DEEPSEEK_MODEL", "deepseek-chat")
        logger.info(f"API 키 설정됨: {bool(self.api_key)}")

        # DeepSeek 클라이언트 초기화
        if self.api_key:
            try:
                self.client = DirectDeepSeekClient(
                    api_key=self.api_key,
                    model_name=self.model
                )
                logger.info(f"DeepSeek 클라이언트 초기화 성공: {self.model}")
            except Exception as e:
                logger.error(f"DeepSeek 클라이언트 초기화 실패: {e}")
                self.client = None
        else:
            logger.warning("API 키가 설정되지 않아 클라이언트를 초기화할 수 없습니다.")
            self.client = None

        logger.info("간단한 RAG 체인 초기화 완료")

    def _retrieve(self, query: str) -> str:
        """문서 검색 및 컨텍스트 구성"""
        try:
            docs = self.vector_store.similarity_search(query, k=3)
            if not docs:
                return "관련 문서를 찾을 수 없습니다."

            # 검색 결과 컨텍스트 구성
            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")

            context = "\n".join(context_parts)

            # 길이 제한
            if len(context) > 6000:
                context = context[:2500] + "\n...(중략)...\n" + context[-2500:]

            return context
        except Exception as e:
            logger.error(f"검색 중 오류: {e}")
            return "문서 검색 중 오류가 발생했습니다."

    def _generate_prompt(self, query: str, context: str) -> List[Dict[str, str]]:
        """DeepSeek API용 프롬프트 생성"""
        # 시스템 프롬프트
        system_prompt = """다음 정보를 기반으로 질문에 정확하게 답변해주세요.
참고 정보에서 답을 찾을 수 없는 경우 "제공된 문서에서 해당 정보를 찾을 수 없습니다."라고 답변하세요.
정보 출처를 포함해서 대답하세요."""

        # 사용자 프롬프트
        user_prompt = f"""질문: {query}

참고 정보:
{context}"""

        # DeepSeek API 프롬프트 포맷
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]

        return messages

    def run(self, query: str) -> str:
        """쿼리 처리"""
        try:
            logger.info(f"SimpleRAGChain 실행: {query[:50]}...")

            # 문서 검색
            context = self._retrieve(query)

            # 클라이언트가 초기화되지 않은 경우
            if self.client is None:
                logger.warning("DeepSeek 클라이언트가 초기화되지 않음. 검색 결과만 반환.")
                return f"API 연결이 설정되지 않았습니다. 검색 결과:\n\n{context}"

            # 프롬프트 생성
            messages = self._generate_prompt(query, context)

            # API 호출
            start_time = time.time()
            response = self.client.chat(messages)
            logger.info(f"API 응답 시간: {time.time() - start_time:.2f}초")

            if response["success"]:
                logger.info("응답 생성 성공")
                return response["response"]
            else:
                logger.error(f"응답 생성 실패: {response['message']}")
                return f"응답 생성 실패: {response['message']}\n\n검색 결과:\n{context}"

        except Exception as e:
            logger.error(f"실행 중 오류: {e}")
            return f"오류 발생: {str(e)}"