src/mcp_processor.py

"""
Model Context Protocol(MCP) 처리 모듈

Claude와의 통신을 위한 MCP 프로토콜 처리 클래스
"""

import os
import json
import logging
import uuid
from datetime import datetime
from typing import Dict, List, Any, Optional, Union

import mcp

from src import config
from src.utils import (
    MCPError,
    DataProcessingError,
    generate_unique_id,
    save_json_data,
    merge_metadata
)

logger = config.setup_logger("mcp_processor")

class MCPProcessor:
    """
    Claude와의 통신을 위한 MCP 프로토콜 처리기
    
    MCP 스펙에 맞게 컨텍스트를 준비하고 Claude와 통신하는 기능 제공
    """
    
    _instance = None
    
    def __new__(cls, *args, **kwargs):
        """싱글턴 패턴 구현"""
        if cls._instance is None:
            cls._instance = super(MCPProcessor, cls).__new__(cls)
        return cls._instance
    
    def __init__(self, model_name: str = None, api_key: str = None):
        """
        MCP 프로세서 초기화
        
        Args:
            model_name: 사용할 Claude 모델명
            api_key: API 키
        """
        # 이미 초기화된 경우 중복 초기화 방지
        if hasattr(self, 'initialized') and self.initialized:
            return
            
        # MCP SDK 환경 설정
        self.model_name = model_name or config.MCP_MODEL_NAME
        self.api_key = api_key or config.MCP_API_KEY
        
        try:
            # MCP 클라이언트 초기화
            self.client = mcp.Client(
                model=self.model_name,
                api_key=self.api_key if self.api_key else None  # API 키가 없으면 None 전달
            )
            
            self.conversations = {}  # 대화 내용 캐시
            self.data_dir = config.DATA_DIR
            self.initialized = True
            
            logger.info(f"MCP 프로세서 초기화됨 (모델: {self.model_name})")
            
        except Exception as e:
            logger.error(f"MCP 프로세서 초기화 실패: {str(e)}")
            raise MCPError(f"MCP 프로세서 초기화 실패: {str(e)}", "init")
    
    def process_query(
        self, 
        query: str, 
        context_data: Optional[Union[List[Dict[str, Any]], Dict[str, Any]]] = None, 
        conversation_id: Optional[str] = None,
        system_prompt: Optional[str] = None,
        temperature: float = 0.2,
        max_tokens: int = 2000
    ) -> Dict[str, Any]:
        """
        MCP를 사용하여 쿼리에 대한 응답 생성
        
        Args:
            query: 사용자 쿼리
            context_data: 컨텍스트 데이터
            conversation_id: 대화 ID
            system_prompt: 시스템 프롬프트
            temperature: 생성 온도(창의성)
            max_tokens: 최대 생성 토큰 수
            
        Returns:
            응답 및 메타데이터를 포함한 딕셔너리
            
        Raises:
            MCPError: MCP 처리 중 오류 발생 시
        """
        if not query:
            raise MCPError("쿼리가 비어 있습니다.", "empty_query")
            
        logger.info(f"MCP 쿼리 처리 중: {query[:50]}...")
        
        # 컨텍스트 처리
        context = None
        if context_data:
            context = self._prepare_context(context_data)
            logger.debug(f"컨텍스트 준비 완료: {len(context)}개 항목")
        
        # 대화 ID 처리
        if not conversation_id:
            conversation_id = generate_unique_id("conv")
            logger.debug(f"새 대화 ID 생성: {conversation_id}")
        
        # 시스템 프롬프트 준비
        if system_prompt:
            # 사용자 제공 시스템 프롬프트 사용
            pass
        else:
            # 기본 시스템 프롬프트 설정
            system_prompt = self._get_default_system_prompt()
        
        try:
            # MCP를 통한 응답 생성
            response = self.client.query(
                message=query,
                context=context,
                conversation_id=conversation_id,
                system_prompt=system_prompt,
                parameters={
                    "temperature": temperature,
                    "max_tokens": max_tokens
                }
            )
            
            logger.info(f"MCP 응답 생성 완료: {len(response.content)} 자")
            
            # 대화 이력 캐싱
            self._update_conversation_cache(conversation_id, query, response.content, context)
            
            # 필요한 경우 대화 저장
            # self.save_conversation(conversation_id)
            
            # 응답 및 사용된 컨텍스트 반환
            result = {
                "content": response.content,
                "conversation_id": response.conversation_id,
                "used_context": getattr(response, 'used_context', None),
                "query": query,
                "timestamp": datetime.now().isoformat()
            }
            
            return result
            
        except Exception as e:
            logger.error(f"MCP 쿼리 처리 중 오류 발생: {str(e)}")
            raise MCPError(f"MCP 쿼리 처리 중 오류 발생: {str(e)}", "query_processing")
    
    def _prepare_context(
        self, 
        context_data: Union[List[Dict[str, Any]], Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """
        MCP에 전달할 컨텍스트 준비
        
        Args:
            context_data: 컨텍스트 데이터
            
        Returns:
            MCP 형식의 컨텍스트 리스트
        """
        if isinstance(context_data, list):
            # 이미 리스트 형태로 전달된 경우
            return [self._format_context_item(item) for item in context_data]
        elif isinstance(context_data, dict):
            # 단일 아이템이 딕셔너리로 전달된 경우
            return [self._format_context_item(context_data)]
        else:
            logger.error(f"지원되지 않는 컨텍스트 데이터 형식: {type(context_data)}")
            raise MCPError(
                f"지원되지 않는 컨텍스트 데이터 형식: {type(context_data)}",
                "invalid_context_format"
            )
    
    def _format_context_item(self, item: Dict[str, Any]) -> Dict[str, Any]:
        """
        컨텍스트 아이템을 MCP 형식으로 변환
        
        Args:
            item: 원본 컨텍스트 아이템
            
        Returns:
            MCP 형식의 컨텍스트 아이템
        """
        # 필수 필드 검증 ('text' 또는 'content')
        raw_text = item.get("text") or item.get("content")
        if not raw_text:
            raise MCPError("컨텍스트 아이템에 'text' 또는 'content' 필드가 없습니다.", "missing_context_content")
        
        # 고유 ID 생성 또는 사용
        ctx_id = item.get("id", generate_unique_id("ctx"))
        
        # 메타데이터 추출
        metadata = item.get("metadata", {})
        source_code = metadata.get("code", "알 수 없음")
        relevance = item.get("relevance", metadata.get("relevance", 0))
        
        # XML 태그로 컨텍스트 래핑
        xml_content = f"""<evidence source="{source_code}" relevance="{relevance:.4f}">
{raw_text}
</evidence>"""
        
        # 기본 필드 설정 (MCP 표준 형식)
        formatted = {
            "id": ctx_id,
            "content": xml_content
        }
        
        # 메타데이터가 있으면 추가
        if "metadata" in item:
            formatted["metadata"] = item["metadata"]
            # 생성 시간 추가 (없는 경우)
            if "created_at" not in formatted["metadata"]:
                formatted["metadata"]["created_at"] = datetime.now().isoformat()
        else:
            # 기본 메타데이터 추가
            formatted["metadata"] = {
                "source": "KCSC",
                "created_at": datetime.now().isoformat()
            }
        
        return formatted
    
    def _get_default_system_prompt(self) -> str:
        """기본 시스템 프롬프트 생성"""
        return """한국건설기준센터(KCSC)의 건설기준 정보에 대한 질의응답을 수행합니다.
주어진 컨텍스트를 바탕으로 정확하고 객관적인 답변을 제공하세요.
건설기준에 명시된 내용을 우선적으로 참고하되, 명확한 정보가 없는 경우에는 이를 솔직하게 밝히고 
일반적인 건축구조 지식을 바탕으로 답변할 수 있습니다.

답변 시 참고할 사항:
1. 인용된 건설기준 코드와 출처를 명확히 기재하세요.
2. 관련 수식이나 도표가 필요한 경우 적절히 포맷팅하여 제공하세요.
3. 불완전하거나 모호한 정보에 대해서는 추가 정보를 요청하거나 제한사항을 명시하세요.
4. 가능한 한 관련 법규, 규정, 표준과 같은 공식 정보를 참조하세요.
5. 특정 설계 문제에 대한 최종 책임은 항상 공인된 전문가에게 있음을 명시하세요.

한국어로 답변하되, 필요한 경우 전문 용어는 한글과 영문을 병기할 수 있습니다."""
    
    def _update_conversation_cache(
        self,
        conversation_id: str,
        query: str,
        response: str,
        context: Optional[List[Dict[str, Any]]] = None
    ):
        """
        대화 이력 캐시 업데이트
        
        Args:
            conversation_id: 대화 ID
            query: 사용자 쿼리
            response: AI 응답
            context: 사용된 컨텍스트
        """
        if conversation_id not in self.conversations:
            self.conversations[conversation_id] = {
                "id": conversation_id,
                "messages": [],
                "contexts": [],
                "created_at": datetime.now().isoformat(),
                "updated_at": datetime.now().isoformat()
            }
        
        # 메시지 추가
        self.conversations[conversation_id]["messages"].append({
            "role": "user",
            "content": query,
            "timestamp": datetime.now().isoformat()
        })
        
        self.conversations[conversation_id]["messages"].append({
            "role": "assistant",
            "content": response,
            "timestamp": datetime.now().isoformat()
        })
        
        # 컨텍스트 추가 (중복 방지)
        if context:
            for ctx in context:
                ctx_id = ctx.get("id")
                if ctx_id:
                    # 이미 있는 컨텍스트인지 확인
                    existing_ctx = next(
                        (c for c in self.conversations[conversation_id]["contexts"] if c.get("id") == ctx_id),
                        None
                    )
                    
                    if not existing_ctx:
                        self.conversations[conversation_id]["contexts"].append(ctx)
        
        # 업데이트 시간 갱신
        self.conversations[conversation_id]["updated_at"] = datetime.now().isoformat()
    
    def get_conversation_history(self, conversation_id: str) -> Optional[Dict[str, Any]]:
        """
        대화 이력 조회
        
        Args:
            conversation_id: 대화 ID
            
        Returns:
            대화 이력 또는 None (찾지 못한 경우)
        """
        # 캐시에서 먼저 조회
        if conversation_id in self.conversations:
            return self.conversations[conversation_id]
        
        # MCP SDK의 대화 이력 조회 기능 사용
        try:
            logger.info(f"MCP에서 대화 이력 조회 중: {conversation_id}")
            history = self.client.get_conversation(conversation_id)
            
            # 캐시 업데이트
            self.conversations[conversation_id] = history
            
            return history
            
        except Exception as e:
            logger.error(f"대화 이력 조회 중 오류 발생: {str(e)}")
            return None
    
    def save_conversation(
        self, 
        conversation_id: str, 
        file_prefix: str = "conversation",
        include_context: bool = True
    ) -> Optional[str]:
        """
        대화 내용 저장
        
        Args:
            conversation_id: 저장할 대화 ID
            file_prefix: 파일명 접두사
            include_context: 컨텍스트 포함 여부
            
        Returns:
            저장된 파일 경로 또는 None (실패 시)
        """
        if not conversation_id:
            logger.warning("저장할 대화 ID가 없습니다.")
            return None
        
        try:
            # 대화 내용 가져오기
            conversation = self.get_conversation_history(conversation_id)
            
            if not conversation:
                logger.warning(f"대화 ID {conversation_id}에 해당하는 대화를 찾을 수 없습니다.")
                return None
            
            # 컨텍스트 필터링 (선택적)
            if not include_context and "contexts" in conversation:
                conversation = conversation.copy()
                conversation.pop("contexts", None)
            
            # 파일에 저장
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"{file_prefix}_{timestamp}.json"
            
            filepath = save_json_data(conversation, filename, self.data_dir)
            logger.info(f"대화 내용 저장 완료: {filepath}")
            
            return filepath
            
        except Exception as e:
            logger.error(f"대화 저장 중 오류 발생: {str(e)}")
            return None
    
    def delete_conversation(self, conversation_id: str) -> bool:
        """
        대화 삭제
        
        Args:
            conversation_id: 삭제할 대화 ID
            
        Returns:
            성공 여부
        """
        if not conversation_id:
            logger.warning("삭제할 대화 ID가 없습니다.")
            return False
        
        try:
            # 캐시에서 삭제
            if conversation_id in self.conversations:
                del self.conversations[conversation_id]
            
            # MCP 클라이언트에서 삭제 (API 지원 시)
            try:
                self.client.delete_conversation(conversation_id)
            except Exception as e:
                logger.warning(f"MCP 클라이언트에서 대화 삭제 실패 (기능 미지원): {str(e)}")
            
            logger.info(f"대화 ID {conversation_id} 삭제 완료")
            return True
            
        except Exception as e:
            logger.error(f"대화 삭제 중 오류 발생: {str(e)}")
            return False

# 실행 예시
if __name__ == "__main__":
    # 설정 초기화
    config.init()
    
    # MCP 프로세서 초기화
    processor = MCPProcessor()
    
    # 테스트 쿼리
    test_query = "콘크리트 기둥의 최소 철근비는 얼마인가요?"
    
    # 테스트 컨텍스트
    test_context = [
        {
            "id": "ctx-1",
            "content": "KDS 14 20 01 콘크리트구조 설계기준에 따르면, 기둥의 최소 철근비는 단면적의 1%이다.",
            "metadata": {
                "source": "KDS 14 20 01",
                "title": "콘크리트구조 설계기준",
                "relevance": 0.95
            }
        }
    ]
    
    # 쿼리 처리
    result = processor.process_query(test_query, test_context)
    print(result["content"])