mcp_server.py

"""
영웅문 S# AI 메뉴검색 — MCP 서버 (컨시어지 연동용)
================================================================
컨시어지(MCP 클라이언트)가 tools/call 로 호출하는 5개 tool을 노출한다.
프로토콜(initialize / tools/list / tools/call / JSON-RPC)은 FastMCP가
자동 처리하며, 여기서는 tool 함수 본체 + 코어 연결만 정의한다.

코어(core.agent.run_agent / search_engine)는 그대로 재사용한다.

실행:
    # stdio (컨시어지가 subprocess로 실행)
    .venv/Scripts/python.exe mcp_server.py

    # SSE/HTTP (별도 서비스로 띄우고 네트워크 연결)
    MCP_TRANSPORT=sse .venv/Scripts/python.exe mcp_server.py

전송 방식은 컨시어지 측과 합의 후 MCP_TRANSPORT 로 선택한다(기본 stdio).

[중요] stdio 전송에서는 stdout이 JSON-RPC 채널이라 print() 오염이 프로토콜을
깨뜨린다. 따라서 이 서버를 stdio로 띄울 때는 코어의 모든 로그를 stderr로
보내야 한다(아래 _silence_stdout 로 일괄 처리).
"""
import os
import sys
import io
import time
import uuid
import json
import logging
import contextlib
from pathlib import Path
from typing import Optional

sys.path.insert(0, str(Path(__file__).parent))

# ── 로깅: 항상 stderr로 (stdio 전송 시 stdout 오염 방지) ──────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s [mcp] %(message)s",
    stream=sys.stderr,
)
logger = logging.getLogger("menu-search-mcp")

from mcp.server.fastmcp import FastMCP

# ── 응답 스펙 버전 ────────────────────────────────────────────────────────────
SPEC_VERSION = "v1"
AGENT_NAME = "menu-search"
AGENT_VERSION = "1.0.0"

CATEGORIES = [
    "국내주식", "해외주식", "파생상품", "자산/뱅킹", "금융상품",
    "주식더모으기", "랭킹/영웅전", "대출/카드/보험", "환전", "생활/혜택",
]

mcp = FastMCP(
    AGENT_NAME,
    host=os.getenv("MCP_HOST", "0.0.0.0"),
    port=int(os.getenv("MCP_PORT", "9000")),
)

# ── stdout 오염 방지 ──────────────────────────────────────────────────────────
# stdio 전송 시 stdout은 JSON-RPC 채널이다. 코어 모듈들이 print()로 stdout에
# 찍으면 프로토콜이 깨지므로, 코어 호출 구간만 stdout을 stderr로 리다이렉트한다.
@contextlib.contextmanager
def _quiet_stdout():
    old = sys.stdout
    sys.stdout = sys.stderr
    try:
        yield
    finally:
        sys.stdout = old


# ── 코어 지연 로드 (싱글톤) ───────────────────────────────────────────────────
_engine = None
_raw_menus = None


def _get_engine():
    global _engine
    if _engine is None:
        with _quiet_stdout():
            from core.search_engine import MenuSearchEngine
            _engine = MenuSearchEngine.get_instance()
    return _engine


def _get_raw_menus() -> dict:
    """real_menus.json 을 menu_id → 메뉴 dict 로 로드(원본 데이터 조회용)."""
    global _raw_menus
    if _raw_menus is None:
        from config import RAW_DIR
        path = RAW_DIR / "real_menus.json"
        data = json.loads(path.read_text(encoding="utf-8"))
        _raw_menus = {m["menu_id"]: m for m in data}
    return _raw_menus


# ── 공통 응답 포맷 (v1) — 모든 tool 응답을 감싸는 공통 메타 ──────────────────────
def _envelope(payload: dict, trace_id: Optional[str], started: float,
              ok: bool = True, error: Optional[dict] = None) -> dict:
    return {
        "agent": AGENT_NAME,
        "agentVersion": AGENT_VERSION,
        "specVersion": SPEC_VERSION,
        "traceId": trace_id or str(uuid.uuid4()),
        "ok": ok,
        "elapsedMs": int((time.time() - started) * 1000),
        "error": error,        # {code, message, detail} or None
        **payload,
    }


def _err(code: str, message: str, detail: str = "") -> dict:
    return {"code": code, "message": message, "detail": detail}


# ── 컨시어지 게이트웨이 호환 tool ──────────────────────────────────────────────
# 게이트웨이 tools/list 스키마(search_menu / fetch_menu_raw / health)와 동일한
# tool 이름으로 노출하고, 응답은 우리 공통 응답 포맷(v1)으로 통일한다.
# 내부 검색 엔진: 임베딩(bge-m3) + BM25 + HyDE 하이브리드(hybrid-v1).
ENGINE_VERSION = "hybrid-v1"


# ════════════════════════════════════════════════════════════════════════════
# Tool) health — 상태(liveness/readiness) / 버전 / 데이터 건수
# ════════════════════════════════════════════════════════════════════════════
@mcp.tool()
def health() -> dict:
    """서버 상태(liveness/readiness)·버전·참조 데이터 건수·엔진을 공통 응답 포맷으로 반환한다."""
    started = time.time()
    try:
        engine = _get_engine()
        stats = engine.vectorstore.get_stats()
        return _envelope({"status": "ok", "live": True, "ready": True,
                          "menuCount": stats.get("total_menus", 0),
                          "engine": ENGINE_VERSION}, None, started)
    except Exception as e:
        logger.exception("health 실패")
        return _envelope({"status": "error", "live": True, "ready": False},
                         None, started, ok=False,
                         error=_err("ENGINE_NOT_READY", "검색엔진 미준비", str(e)))


# ════════════════════════════════════════════════════════════════════════════
# Tool) search_menu — 자연어 메뉴 검색
# ════════════════════════════════════════════════════════════════════════════
@mcp.tool()
def search_menu(
    query: str,
    limit: int = 5,
    agentHint: Optional[str] = None,
    version: str = "hybrid-v1",
    verbose: bool = False,
    debug: bool = False,
) -> dict:
    """키움 영웅문S# 메뉴를 검색한다. (엔진: 임베딩 bge-m3 + BM25 + HyDE 하이브리드)

    verbose/debug로 응답 상세도를 조절한다.

    Args:
        query: 사용자 질의 (예: "잔고 조회", "내 주식 얼마야").
        limit: 최대 결과 수 (기본 5, 최대 30).
        agentHint: 도메인 힌트(선택). 유효 카테고리면 필터, 아니면 자체 감지.
        version: 검색 엔진 버전 (기본 hybrid-v1).
        verbose: true면 각 결과에 category/keywords + 의도/품질 포함.
        debug: true면 재작성 이력·단계 정보 포함.

    Returns:
        공통 응답 포맷 + results:[{rank, menu_id, menu_name, menu_path, similarity}].
    """
    started = time.time()
    if not query or not query.strip():
        return _envelope({"query": query, "resultCount": 0, "results": []},
                         None, started, ok=False,
                         error=_err("BAD_REQUEST", "query가 비어있음"))
    try:
        from core.agent import run_agent
        limit = max(1, min(int(limit), 30))
        # agentHint가 유효 카테고리면 필터로, 아니면 자체 카테고리 감지 사용
        cat = agentHint if agentHint in CATEGORIES else None
        with _quiet_stdout():   # run_agent 내부 print → stderr (stdout 오염 방지)
            # HyDE 유지(의미 갭 브릿지). 리랭커는 OFF — CPU cross-encoder가
            # 후보 재점수에 6~11초 소요 + 짧은 메뉴명(예: "OTP")을 과소평가.
            result = run_agent(query=query, top_n=limit,
                               category_filter=cat,
                               use_hyde=True, use_reranker=False)
        hits = result.get("search_results", [])
        raw_menus = _get_raw_menus()   # menu_id → 원본 dict (category, keywords 조회용)

        results = []
        for i, r in enumerate(hits, 1):
            menu_id = r.get("menu_id", "")
            sim = round(r.get("similarity", 0.0), 4)
            raw = raw_menus.get(menu_id, {})
            item = {
                "rank": i,
                "id": menu_id,
                "name": r.get("menu_name", ""),
                "path": r.get("menu_path", ""),
                "score": sim,           # 공통 랭킹 점수 (0~1.0) — DIO/BIX 통합 필드
                "similarity": sim,      # BIX 고유 임베딩 유사도 (현재 score와 동일값)
                "matchType": "semantic",
                "category": raw.get("category", ""),
                "keywords": raw.get("keywords") or [],
            }
            results.append(item)

        payload = {"query": query, "resultCount": len(results), "results": results}
        if verbose:
            payload["detectedCategory"] = result.get("detected_category")
            payload["intent"] = result.get("intent")
            payload["qualityScore"] = round(result.get("quality_score", 0.0), 4)
        if debug:
            payload["debug"] = {
                "refinedQuery": result.get("refined_query", query),
                "retryCount": result.get("retry_count", 0),
                "rewriteHistory": result.get("rewrite_history", []),
                "detectedCategory": result.get("detected_category"),
                "intent": result.get("intent"),
            }
        return _envelope(payload, None, started)
    except Exception as e:
        logger.exception("search_menu 실패")
        return _envelope({"query": query, "resultCount": 0, "results": []},
                         None, started, ok=False,
                         error=_err("INTERNAL", "검색 중 오류", str(e)))


# ════════════════════════════════════════════════════════════════════════════
# Tool) fetch_menu_raw — 참조 데이터 전체 반환(대조 검증용)
# ════════════════════════════════════════════════════════════════════════════
@mcp.tool()
def fetch_menu_raw() -> dict:
    """메뉴 검색이 참조하는 원본 데이터(영웅문S# flat JSON) 전체를 반환한다.

    외부에서 참조 데이터 전체를 받아 검색 결과를 대조 검증하는 용도.

    Returns:
        공통 응답 포맷 + count + menus:[원본 메뉴 dict, ...].
    """
    started = time.time()
    try:
        menus = _get_raw_menus()
        return _envelope({"count": len(menus), "menus": list(menus.values())},
                         None, started)
    except Exception as e:
        logger.exception("fetch_menu_raw 실패")
        return _envelope({"count": 0, "menus": []}, None, started, ok=False,
                         error=_err("INTERNAL", "조회 오류", str(e)))


# ════════════════════════════════════════════════════════════════════════════
# ── Bearer 토큰 인증 미들웨어 ─────────────────────────────────────────────────
# MCP_TOKEN 이 설정돼 있으면 모든 HTTP 요청에 'Authorization: Bearer <token>'을
# 요구한다(공개 HF Space URL 보호용). 미설정이면 인증 비활성(로컬 개발).
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse


class _BearerAuthMiddleware(BaseHTTPMiddleware):
    def __init__(self, app, token: str):
        super().__init__(app)
        self._expected = f"Bearer {token}"

    async def dispatch(self, request, call_next):
        # keep-alive 핑(/healthz)은 인증 면제 — Space를 깨어있게만 유지
        if request.url.path == "/healthz":
            return await call_next(request)
        if request.headers.get("authorization", "") != self._expected:
            return JSONResponse(
                {"ok": False, "error": {"code": "UNAUTHORIZED",
                                        "message": "유효하지 않은 토큰"}},
                status_code=401,
            )
        return await call_next(request)


def _run_streamable_http():
    """streamable_http_app()에 토큰 인증 + keep-alive 엔드포인트를 붙여 실행."""
    import uvicorn
    from starlette.routing import Route
    from starlette.responses import PlainTextResponse

    app = mcp.streamable_http_app()

    # ── keep-alive 핑 엔드포인트 (인증 면제, 엔진 로드 없이 즉시 200) ──────────
    # 외부 cron이 주기적으로 GET /healthz 를 호출하면 HF Space가 sleep에 빠지지
    # 않아 콜드스타트(첫 MCP 호출 ~100초)를 방지한다.
    async def _healthz(request):
        return PlainTextResponse("ok")
    app.router.routes.append(Route("/healthz", _healthz, methods=["GET"]))

    token = os.getenv("MCP_TOKEN", "")
    if token:
        app.add_middleware(_BearerAuthMiddleware, token=token)
        logger.info("토큰 인증 활성 (MCP_TOKEN)")
    else:
        logger.warning("MCP_TOKEN 미설정 → 인증 비활성 (공개 배포 시 반드시 설정)")
    host = os.getenv("MCP_HOST", "0.0.0.0")
    port = int(os.getenv("MCP_PORT", "9000"))
    uvicorn.run(app, host=host, port=port, log_level="info")


def _preload_engine_os_quiet():
    """JSON-RPC(stdout) 시작 전에 엔진을 미리 로드한다.

    모델 라이브러리(sentence-transformers 등)는 C 레벨로 fd 1(stdout)에 직접
    출력하기도 해서, 파이썬 레벨 stdout 교체로는 못 막는다. 로드 구간 동안
    OS 레벨로 fd 1을 fd 2(stderr)로 dup2 하여 JSON-RPC 채널 오염을 차단한다.
    """
    saved = os.dup(1)
    try:
        sys.stdout.flush()
        os.dup2(2, 1)          # stdout(fd1) → stderr(fd2)
        _get_engine()
    except Exception:
        logger.exception("엔진 preload 실패")
    finally:
        sys.stdout.flush()
        os.dup2(saved, 1)      # fd1 복원 (이후 JSON-RPC는 깨끗한 stdout)
        os.close(saved)


if __name__ == "__main__":
    transport = os.getenv("MCP_TRANSPORT", "stdio").lower()
    logger.info(f"MCP 서버 기동 (transport={transport}, spec={SPEC_VERSION}) — 엔진 preload 중...")
    _preload_engine_os_quiet()
    logger.info("엔진 준비 완료, 서버 시작")
    if transport in ("streamable-http", "http"):
        _run_streamable_http()          # 토큰 인증 포함 (배포 기본)
    elif transport == "sse":
        mcp.run(transport="sse")
    else:
        mcp.run(transport="stdio")