rag_chain.py

"""
직접 DeepSeek API 호출을 위한 클라이언트 구현
"""
import os
import time
import logging
import requests
import json
from typing import Dict, Any, Optional, List

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

class DirectDeepSeekClient:
    """
    DeepSeek API를 직접 호출하는 클라이언트
    OpenAI 클라이언트를 우회하고 직접 HTTP 요청 사용
    """
    def __init__(self, api_key: str, model_name: str = "deepseek-chat"):
        """
        클라이언트 초기화

        Args:
            api_key: DeepSeek API 키
            model_name: 사용할 모델 이름 (기본값: "deepseek-chat")
        """
        self.api_key = api_key
        self.model_name = model_name
        self.endpoint = os.getenv("DEEPSEEK_ENDPOINT", "https://api.deepseek.com/v1/chat/completions")
        logger.info(f"DirectDeepSeekClient 초기화: 모델={model_name}, 엔드포인트={self.endpoint}")

    def generate(self,
                 prompt: str,
                 temperature: float = 0.3,
                 max_tokens: int = 1000,
                 max_retries: int = 3,
                 timeout: int = 60) -> Dict[str, Any]:
        """
        텍스트 생성 요청

        Args:
            prompt: 입력 프롬프트
            temperature: 생성 온도 (0.0 ~ 1.0)
            max_tokens: 최대 생성 토큰 수
            max_retries: 재시도 횟수
            timeout: 요청 타임아웃 (초)

        Returns:
            생성 결과 딕셔너리 (success, response, message 등)
        """
        # 메시지 구성 (단일 사용자 메시지)
        messages = [{"role": "user", "content": prompt}]
        return self.chat(messages, temperature, max_tokens, max_retries, timeout)

    def chat(self,
             messages: List[Dict[str, str]],
             temperature: float = 0.3,
             max_tokens: int = 1000,
             max_retries: int = 3,
             timeout: int = 60) -> Dict[str, Any]:
        """
        채팅 API 호출

        Args:
            messages: 채팅 메시지 리스트 (role, content 키를 가진 딕셔너리 리스트)
            temperature: 생성 온도 (0.0 ~ 1.0)
            max_tokens: 최대 생성 토큰 수
            max_retries: 재시도 횟수
            timeout: 요청 타임아웃 (초)

        Returns:
            생성 결과 딕셔너리 (success, response, message 등)
        """
        # API 요청 헤더 및 데이터
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

        payload = {
            "model": self.model_name,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens
        }

        # 재시도 로직
        retry_delay = 1.0
        attempt = 0

        while attempt < max_retries:
            attempt += 1
            try:
                logger.info(f"DeepSeek API 요청 시도 ({attempt}/{max_retries})...")

                # API 요청 전송
                response = requests.post(
                    self.endpoint,
                    headers=headers,
                    json=payload,
                    timeout=timeout
                )

                # 응답 확인
                if response.status_code == 200:
                    result = response.json()

                    # 응답 내용 추출
                    if "choices" in result and len(result["choices"]) > 0:
                        message_content = result["choices"][0].get("message", {}).get("content", "")
                        logger.info(f"DeepSeek API 응답 성공 (길이: {len(message_content)})")

                        return {
                            "success": True,
                            "response": message_content,
                            "status_code": response.status_code,
                            "raw_response": result
                        }
                    else:
                        logger.warning(f"DeepSeek API 응답은 성공했으나 예상치 못한 응답 형식: {result}")
                        return {
                            "success": False,
                            "message": "응답에서 메시지를 찾을 수 없습니다",
                            "status_code": response.status_code,
                            "raw_response": result
                        }
                else:
                    logger.error(f"DeepSeek API 오류: 상태 코드 {response.status_code}")

                    # 오류 메시지 추출
                    error_message = ""
                    try:
                        error_data = response.json()
                        error_message = error_data.get("error", {}).get("message", str(error_data))
                    except:
                        error_message = response.text

                    # 요청 한도 초과시 더 오래 대기
                    if response.status_code == 429:
                        retry_delay = min(retry_delay * 3, 15)
                    else:
                        retry_delay = min(retry_delay * 2, 10)

                    if attempt < max_retries:
                        logger.info(f"{retry_delay}초 후 재시도...")
                        time.sleep(retry_delay)
                    else:
                        # 모든 시도 실패
                        return {
                            "success": False,
                            "message": f"API 오류: {error_message}",
                            "status_code": response.status_code
                        }

            except requests.exceptions.Timeout:
                logger.error("DeepSeek API 요청 시간 초과")

                if attempt < max_retries:
                    logger.info(f"{retry_delay}초 후 재시도...")
                    time.sleep(retry_delay)
                    retry_delay = min(retry_delay * 2, 10)
                else:
                    return {
                        "success": False,
                        "message": "API 요청 시간 초과",
                        "status_code": None
                    }

            except requests.exceptions.ConnectionError:
                logger.error("DeepSeek API 연결 실패")

                if attempt < max_retries:
                    logger.info(f"{retry_delay}초 후 재시도...")
                    time.sleep(retry_delay)
                    retry_delay = min(retry_delay * 2, 10)
                else:
                    return {
                        "success": False,
                        "message": "API 서버 연결 실패",
                        "status_code": None
                    }

            except Exception as e:
                logger.error(f"DeepSeek API 요청 중 예상치 못한 오류: {e}")

                if attempt < max_retries:
                    logger.info(f"{retry_delay}초 후 재시도...")
                    time.sleep(retry_delay)
                    retry_delay = min(retry_delay * 2, 10)
                else:
                    return {
                        "success": False,
                        "message": f"예상치 못한 오류: {str(e)}",
                        "status_code": None
                    }

        # 모든 시도 실패
        return {
            "success": False,
            "message": "최대 재시도 횟수 초과",
            "status_code": None
        }

    def system_prompt_chat(self,
                           system_prompt: str,
                           user_prompt: str,
                           temperature: float = 0.3,
                           max_tokens: int = 1000,
                           max_retries: int = 3,
                           timeout: int = 60) -> Dict[str, Any]:
        """
        시스템 프롬프트와 사용자 프롬프트를 이용한 채팅 API 호출

        Args:
            system_prompt: 시스템 프롬프트
            user_prompt: 사용자 프롬프트
            temperature: 생성 온도 (0.0 ~ 1.0)
            max_tokens: 최대 생성 토큰 수
            max_retries: 재시도 횟수
            timeout: 요청 타임아웃 (초)

        Returns:
            생성 결과 딕셔너리
        """
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]

        return self.chat(messages, temperature, max_tokens, max_retries, timeout)


# 단독 실행을 위한 테스트 코드
if __name__ == "__main__":
    # 로깅 설정
    logging.basicConfig(level=logging.INFO)

    # API 키 확인
    api_key = os.environ.get("DEEPSEEK_API_KEY")
    if not api_key:
        print("환경 변수 DEEPSEEK_API_KEY가 설정되지 않았습니다.")
        exit(1)

    # 클라이언트 생성
    client = DirectDeepSeekClient(api_key)

    # 간단한 테스트
    response = client.generate("Hello, what can you do?")

    # 결과 출력
    if response["success"]:
        print("응답 성공!")
        print(response["response"])
    else:
        print(f"응답 실패: {response['message']}")