core/agent.py

"""
LangGraph 기반 AI 메뉴 검색 에이전트

모든 검색은 항상 이 그래프를 통해 실행된다.
streamlit_app.py에서 직접 engine.search()를 호출하지 않고
run_agent()만 사용하는 단일 경로 구조.

그래프 구조:
  START
    └─► analyze_query      (규칙 기반 카테고리/의도 감지, LLM 없음)
          └─► search_menus  (MenuSearchEngine 재사용)
                └─► evaluate_results  (quality_score 계산)
                      ├─(만족 or retry>=2)─► generate_response ─► END
                      └─(불만족 and retry<2)─► rewrite_query
                                                   └─► search_menus  (루프백)
"""
import re
import sys
from pathlib import Path
from typing import Optional

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

from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END

from config import AZURE_KEY, AZURE_ENDPOINT, AZURE_API_VERSION, LLM_MODEL, DEFAULT_TOP_N, SIMILARITY_THRESHOLD
from core.search_engine import MenuSearchEngine

# ---------------------------------------------------------------------------
# AgentState
# ---------------------------------------------------------------------------

class AgentState(TypedDict):
    # ── 사용자 입력 (Streamlit → Agent) ──────────────────────────────────────
    original_query:     str             # 사용자 원본 입력 (불변)
    use_reranker:       bool            # Cohere Rerank 4 리랭킹 여부 (기본 True)
    use_hyde:           bool            # HyDE 활성 여부 (기본 True)
    top_n:              int             # 반환할 결과 수 (기본 DEFAULT_TOP_N)
    threshold:          float           # 최소 유사도 임계값 (기본 SIMILARITY_THRESHOLD)
    category_filter:    Optional[str]   # 사이드바 카테고리 필터 (기본 None)

    # ── Agent 내부 상태 ────────────────────────────────────────────────────
    detected_category:  Optional[str]   # 규칙 기반 감지 카테고리 ("국내주식" | ... | None)
    intent:             str             # "주문" | "잔고조회" | "시세" | "차트" | "기타"
    refined_query:      str             # 실제 검색 쿼리 (재작성 시 변경)
    search_results:     list            # MenuSearchEngine.search() 반환값
    quality_score:      float           # top1*0.7 + top3avg*0.3
    is_satisfactory:    bool            # quality_score >= QUALITY_THRESHOLD
    retry_count:        int             # 재검색 횟수 (최대 2회)
    rewrite_history:    list            # 재작성 쿼리 이력 (디버깅용)
    final_response:     str             # 최종 출력 텍스트
    error_message:      Optional[str]   # 오류 메시지


# ---------------------------------------------------------------------------
# 규칙 테이블 (카테고리 / 의도 감지)
# ---------------------------------------------------------------------------

CATEGORY_RULES = {
    "국내주식":       ["국내", "주식", "코스피", "코스닥", "kospi", "종목", "일봉", "차트", "호가", "체결", "잔고", "손익", "현재가", "예약주문", "정정", "취소",
                      "투자자별", "수급", "매매동향", "매매현황", "외국인", "기관", "프로그램매매"],
    "해외주식":       ["해외", "미국", "나스닥", "nasdaq", "뉴욕", "s&p", "sp500", "해외주식",
                      "해외잔고", "해외손익", "해외체결", "해외거래"],
    "파생상품":       ["파생", "선물", "옵션", "futures", "선물옵션", "지수선물",
                      "cfd", "국내cfd", "해외cfd", "금현물", "미국옵션"],
    "자산/뱅킹":      ["자산", "계좌", "이체", "잔고통합", "통합조회", "계좌간", "뱅킹"],
    "금융상품":       ["펀드", "채권", "etf", "금융상품", "랩", "wrap", "랩계좌"],
    "주식더모으기":   ["더모으기", "적립", "자동매수", "정기매수"],
    "랭킹/영웅전":    ["랭킹", "영웅전", "순위"],
    "대출/카드/보험": ["대출", "카드", "보험"],
    "환전":           ["환전", "외환", "환율"],
    "생활/혜택":      ["생활", "혜택", "포인트", "이벤트"],
}

INTENT_RULES = {
    "주문":     ["주문", "매수", "매도", "사다", "팔다", "사고", "팔고", "예약", "취소", "정정", "일괄"],
    "잔고조회": ["잔고", "보유", "얼마", "수익", "손익", "수익률", "올랐", "내렸", "얼마야"],
    "수급":     ["투자자별", "수급", "매매동향", "매매현황", "외국인", "기관", "거래원", "세력", "프로그램"],
    "시세":     ["시세", "현재가", "호가", "체결", "가격", "얼마에", "지금"],
    "차트":     ["차트", "일봉", "주봉", "월봉", "캔들", "그래프"],
    "조건검색": ["조건", "스크리닝", "조건식", "실시간"],
}

# 카테고리 감지 최소 신뢰도: 키워드 2개 이상 매칭 시에만 필터 적용
# 근거: "취소" 1개 → False positive 방지(전체 검색), "국내 주식 취소" 3개 → 신뢰 필터 적용
CATEGORY_CONFIDENCE_MIN_HITS = 2


# ---------------------------------------------------------------------------
# 종목 시장 감지 (LLM 보강)
# ---------------------------------------------------------------------------
# 역할: 규칙 기반으로 카테고리를 잡지 못한 경우, 쿼리에 영어 종목명이 있으면
#       LLM으로 국내/해외 여부를 판단하여 category_filter를 보강한다.
#
# 트리거 조건: detected_category가 None인 상태 + 영어 단어가 포함된 쿼리
# 예시:
#   "AMD 아직 체결 안됐어"   → 해외주식  (티커)
#   "SPY ETF 분석해줘"      → 해외주식  (티커)
#   "S&P500 실시간 알고싶어" → 해외주식  (지수)
#   "현대차 차트 보여줘"     → 국내주식  (규칙으로 이미 잡히므로 트리거 안됨)

# 영어 단어 추출 패턴 (2글자 이상 연속 영문자)
_STOCK_EN_RE = re.compile(r'[A-Za-z]{2,}')
# S&P500, S&P 처럼 &로 이어진 금융 지수 패턴
_FINANCIAL_INDEX_RE = re.compile(r'[A-Za-z]&[A-Za-z]')

# 종목명이 아닌 일반 영어 약어/단어 (이것들만 있으면 LLM 호출 안 함)
_NON_STOCK_EN = {
    "etf", "hts", "mts", "ipo", "per", "roe", "roa", "eps",
    "api", "rrf", "ok", "or", "in", "of", "to", "at", "an",
    "the", "and", "is", "it", "us", "kb", "nh", "sk", "lg", "kt",
    "ls", "db", "ibk", "els", "dls",
}

_MARKET_DETECT_SYSTEM = (
    "당신은 주식 시장 분류 전문가입니다. "
    "사용자 질문에 언급된 종목·회사·지수가 "
    "국내(KOSPI/KOSDAQ) 종목인지 해외(미국·해외거래소) 종목인지 판단하세요. "
    "'국내주식' 또는 '해외주식' 중 하나만 출력하세요. "
    "종목 정보가 없거나 불분명하면 '알수없음'을 출력하세요."
)
_MARKET_DETECT_USER = "질문: {query}"


def _has_foreign_stock_hint(query: str) -> bool:
    """
    쿼리에 종목명으로 보이는 영어 단어가 있는지 빠른 휴리스틱 감지.

    감지 패턴:
      - 2글자 이상 영문자 중 일반 약어/불용어가 아닌 것 (AMD, NVDA, Tesla 등)
      - S&P500 같은 금융 지수 패턴 (X&Y 형태)
    """
    tokens = _STOCK_EN_RE.findall(query)
    if any(t.lower() not in _NON_STOCK_EN for t in tokens):
        return True
    # S&P500, D&O 같은 X&Y 금융 지수/약어 패턴
    if _FINANCIAL_INDEX_RE.search(query):
        return True
    return False


def _detect_market_by_llm(query: str) -> Optional[str]:
    """
    LLM으로 국내/해외 종목 시장 분류.

    Returns:
        "국내주식" | "해외주식" | None (불분명 or 오류)
    """
    if not AZURE_KEY or not AZURE_ENDPOINT:
        return None
    try:
        from openai import AzureOpenAI
        client = AzureOpenAI(
            api_key=AZURE_KEY,
            azure_endpoint=AZURE_ENDPOINT,
            api_version=AZURE_API_VERSION,
        )
        resp = client.chat.completions.create(
            model=LLM_MODEL,
            messages=[
                {"role": "system", "content": _MARKET_DETECT_SYSTEM},
                {"role": "user",   "content": _MARKET_DETECT_USER.format(query=query)},
            ],
            max_completion_tokens=50,
            reasoning_effort="low",
        )
        result = resp.choices[0].message.content.strip()
        if "해외" in result:
            return "해외주식"
        if "국내" in result:
            return "국내주식"
        return None
    except Exception as e:
        print(f"[시장감지] LLM 호출 실패: {e}")
        return None


def _detect_category(query: str) -> Optional[str]:
    q = query.lower()

    # 1단계: 고신뢰 단일 키워드 — 이것 하나만 있어도 카테고리 확정
    # CFD, 금현물, 랩 등은 타 카테고리와 겹치지 않아 1개로 충분
    _HIGH_CONF: list[tuple[str, str]] = [
        # 더 긴 복합 키워드 우선 (짧은 것보다 먼저 체크)
        ("미국옵션", "파생상품"), ("국내cfd", "파생상품"), ("해외cfd", "파생상품"),
        ("금현물", "파생상품"),
        # 랩(Wrap)은 금융상품 > Wrap (파생상품 아님)
        ("랩계좌", "금융상품"), ("wrap", "금융상품"),
        # 단일이어도 주식거래 앱 맥락상 충분히 고신뢰
        ("옵션", "파생상품"),   # 옵션잔고미체결, 옵션호가주문 등
        ("cfd", "파생상품"),    # 소문자로 입력해도 탐지
        # "미국 주식/시장" → 해외주식 (미국옵션은 위에서 이미 파생으로 분기됨)
        ("미국", "해외주식"),   # 미국 주식 예상치 상회, 미국 조회상위 등
        ("나스닥", "해외주식"), ("뉴욕", "해외주식"),
    ]
    for kw, cat in _HIGH_CONF:
        if kw in q:
            return cat

    # 2단계: 일반 count 스코어링 (기존 방식 유지)
    scores: dict[str, int] = {}
    for cat, keywords in CATEGORY_RULES.items():
        scores[cat] = sum(1 for kw in keywords if kw in q)
    best_score = max(scores.values())
    if best_score < CATEGORY_CONFIDENCE_MIN_HITS:
        return None

    winners = [c for c, s in scores.items() if s == best_score]

    # 3단계: 동점 처리 — 더 특수한 도메인 우선 (해외 > 파생 > 국내)
    # "해외주식잔고" → 국내주식("주식","잔고")=2, 해외주식("해외","해외주식")=2 → 해외주식 선택
    if "해외주식" in winners and "해외" in q:
        return "해외주식"
    if "파생상품" in winners:
        return "파생상품"

    return winners[0]


def _detect_intent(query: str) -> str:
    q = query.lower()
    scores = {}
    for intent, keywords in INTENT_RULES.items():
        scores[intent] = sum(1 for kw in keywords if kw in q)
    best = max(scores, key=scores.get)
    return best if scores[best] > 0 else "기타"


# ---------------------------------------------------------------------------
# 노드 1: analyze_query  (LLM 없음)
# ---------------------------------------------------------------------------

def analyze_query(state: AgentState) -> AgentState:
    query = state["original_query"]
    detected_category = _detect_category(query)
    intent = _detect_intent(query)

    # ── 종목 시장 LLM 보강 ────────────────────────────────────────────────
    # 규칙 기반으로 카테고리를 못 잡았고(None) + 영어 종목명이 있으면
    # LLM으로 국내/해외 판단하여 category_filter를 채운다.
    # (이미 카테고리가 잡혔으면 규칙 결과를 신뢰하여 LLM 호출 안 함)
    if detected_category is None and _has_foreign_stock_hint(query):
        market = _detect_market_by_llm(query)
        if market:
            detected_category = market
            print(f"[시장감지] '{query[:30]}' → {market} (LLM)")

    return {
        **state,
        "detected_category": detected_category,
        "intent": intent,
        "refined_query": query,   # 초기에는 원본 쿼리 그대로
        "retry_count": 0,
        "rewrite_history": [],
        "search_results": [],
        "quality_score": 0.0,
        "is_satisfactory": False,
        "final_response": "",
        "error_message": None,
    }


# ---------------------------------------------------------------------------
# 노드 2: search_menus  (LLM 없음)
# ---------------------------------------------------------------------------

def search_menus(state: AgentState) -> AgentState:
    try:
        engine = MenuSearchEngine.get_instance()

        # 카테고리 필터: 규칙 기반 감지 카테고리 우선, 없으면 사이드바 필터 사용
        effective_category = state.get("detected_category") or state.get("category_filter")

        results = engine.search(
            query=state["refined_query"],
            top_n=state.get("top_n", DEFAULT_TOP_N),
            category_filter=effective_category,
            threshold=state.get("threshold", SIMILARITY_THRESHOLD),
            use_reranker=state.get("use_reranker", True),
            use_hyde=state.get("use_hyde", True),
        )
        return {**state, "search_results": results, "error_message": None}
    except Exception as e:
        return {**state, "search_results": [], "error_message": str(e)}


# ---------------------------------------------------------------------------
# 노드 3: evaluate_results  (LLM 없음)
# ---------------------------------------------------------------------------

QUALITY_THRESHOLD = 0.55   # RRF min-max 정규화 기준 (0~1 사이)
# RRF 점수는 min-max 정규화로 [0,1] 매핑됨.
# top1=1.0, 하위 결과들은 상대적 분포에 따라 결정.
# 실험 기준: 일반 쿼리 top1 = 0.85~1.0, 불분명한 쿼리 = 0.55~0.75
# 0.55 이상이면 만족할 수 있는 수준의 검색 결과로 판단.


def evaluate_results(state: AgentState) -> AgentState:
    results = state.get("search_results", [])
    if not results:
        return {**state, "quality_score": 0.0, "is_satisfactory": False}

    top1_score = results[0]["similarity"] if len(results) >= 1 else 0.0
    top3_avg = (
        sum(r["similarity"] for r in results[:3]) / min(3, len(results))
        if results else 0.0
    )
    quality_score = round(top1_score * 0.7 + top3_avg * 0.3, 4)
    is_satisfactory = quality_score >= QUALITY_THRESHOLD

    return {
        **state,
        "quality_score": quality_score,
        "is_satisfactory": is_satisfactory,
    }


# ---------------------------------------------------------------------------
# 노드 4: rewrite_query  (Azure GPT-4.1 mini 호출)
# ---------------------------------------------------------------------------

_REWRITE_SYSTEM = (
    "당신은 주식 HTS(영웅문 S#) 메뉴 검색 전문가입니다. "
    "사용자의 구어체 질문을 HTS 메뉴 검색에 최적화된 전문 용어 중심의 짧은 쿼리로 재작성하세요. "
    "출력은 재작성된 쿼리 문장 하나만 반환하세요. 설명 없이 쿼리만 출력하세요."
)

_REWRITE_USER_TMPL = (
    "원본 질문: {original_query}\n"
    "현재 검색 쿼리: {refined_query}\n"
    "감지된 카테고리: {category}\n"
    "감지된 의도: {intent}\n"
    "이전 검색 결과 품질 점수: {quality_score:.2f} (낮을수록 불만족)\n\n"
    "위 정보를 참고하여 더 나은 검색 쿼리를 한 문장으로 작성하세요."
)


def rewrite_query(state: AgentState) -> AgentState:
    try:
        from openai import AzureOpenAI
        client = AzureOpenAI(
            api_key=AZURE_KEY,
            azure_endpoint=AZURE_ENDPOINT,
            api_version=AZURE_API_VERSION,
        )
        user_msg = _REWRITE_USER_TMPL.format(
            original_query=state["original_query"],
            refined_query=state["refined_query"],
            category=state.get("detected_category") or "미감지",
            intent=state.get("intent", "기타"),
            quality_score=state.get("quality_score", 0.0),
        )
        response = client.chat.completions.create(
            model=LLM_MODEL,
            messages=[
                {"role": "system", "content": _REWRITE_SYSTEM},
                {"role": "user",   "content": user_msg},
            ],
            max_completion_tokens=100,
            reasoning_effort="low",
        )
        new_query = response.choices[0].message.content.strip()

        history = list(state.get("rewrite_history", []))
        history.append(new_query)

        return {
            **state,
            "refined_query": new_query,
            "retry_count": state.get("retry_count", 0) + 1,
            "rewrite_history": history,
        }
    except Exception as e:
        # LLM 오류 시 재시도 횟수만 올리고 원본 쿼리 유지
        return {
            **state,
            "retry_count": state.get("retry_count", 0) + 1,
            "error_message": f"rewrite_query 오류: {e}",
        }


# ---------------------------------------------------------------------------
# 노드 5: generate_response  (LLM 없음, 포매팅만)
# ---------------------------------------------------------------------------

def generate_response(state: AgentState) -> AgentState:
    results = state.get("search_results", [])
    retry_count = state.get("retry_count", 0)
    quality_score = state.get("quality_score", 0.0)
    rewrite_history = state.get("rewrite_history", [])
    error_message = state.get("error_message")

    lines = []

    if error_message:
        lines.append(f"[오류] {error_message}")

    if not results:
        lines.append("관련 메뉴를 찾지 못했습니다.")
    else:
        lines.append(f"검색 결과 (품질 점수: {quality_score:.2f}, 재시도: {retry_count}회)")
        lines.append("-" * 50)
        for rank, r in enumerate(results, 1):
            meta = r.get("metadata", {})
            menu_path = r.get("menu_path") or meta.get("menu_path", "")
            menu_name = r.get("menu_name") or meta.get("menu_name", r["menu_id"])
            similarity = r["similarity"]
            bar = "#" * int(similarity * 20)
            lines.append(f"  {rank}. [{similarity:.1%}] {menu_path or menu_name}")
            lines.append(f"     {bar}")

    if retry_count > 0 and rewrite_history:
        lines.append("")
        lines.append("[쿼리 재작성 이력]")
        for i, rq in enumerate(rewrite_history, 1):
            lines.append(f"  {i}차 재작성: {rq}")

    final_response = "\n".join(lines)
    return {**state, "final_response": final_response}


# ---------------------------------------------------------------------------
# 조건부 라우팅
# ---------------------------------------------------------------------------

def route_after_evaluation(state: AgentState) -> str:
    if state.get("is_satisfactory") or state.get("retry_count", 0) >= 2:
        return "generate_response"
    return "rewrite_query"


# ---------------------------------------------------------------------------
# 그래프 조립
# ---------------------------------------------------------------------------

def _build_graph() -> StateGraph:
    builder = StateGraph(AgentState)

    builder.add_node("analyze_query",    analyze_query)
    builder.add_node("search_menus",     search_menus)
    builder.add_node("evaluate_results", evaluate_results)
    builder.add_node("rewrite_query",    rewrite_query)
    builder.add_node("generate_response", generate_response)

    builder.add_edge(START,               "analyze_query")
    builder.add_edge("analyze_query",     "search_menus")
    builder.add_edge("search_menus",      "evaluate_results")
    builder.add_conditional_edges(
        "evaluate_results",
        route_after_evaluation,
        {
            "generate_response": "generate_response",
            "rewrite_query":     "rewrite_query",
        }
    )
    builder.add_edge("rewrite_query",     "search_menus")
    builder.add_edge("generate_response", END)

    return builder.compile()


# 모듈 로드 시 1회 컴파일 (재사용)
_graph = _build_graph()


# ---------------------------------------------------------------------------
# 공개 인터페이스
# ---------------------------------------------------------------------------

def run_agent(
    query: str,
    top_n: int = DEFAULT_TOP_N,
    threshold: float = SIMILARITY_THRESHOLD,
    category_filter: Optional[str] = None,
    use_reranker: bool = True,
    use_hyde: bool = True,
) -> AgentState:
    """
    LangGraph 그래프를 통한 메뉴 검색.
    Streamlit에서 항상 이 함수를 호출한다.

    사용 예:
        result = run_agent("내 주식 얼마나 올랐어?", top_n=5, threshold=0.3)
        search_results = result["search_results"]
        print(result["detected_category"])   # "국내주식"
        print(result["intent"])              # "잔고조회"
        print(result["retry_count"])         # 0~2
        print(result["rewrite_history"])     # ["재작성된 쿼리"]
        print(result["quality_score"])       # 0.0~1.0
    """
    initial_state: AgentState = {
        "original_query":    query,
        "use_reranker":      use_reranker,
        "use_hyde":          use_hyde,
        "top_n":             top_n,
        "threshold":         threshold,
        "category_filter":   category_filter,
        "detected_category": None,
        "intent":            "기타",
        "refined_query":     query,
        "search_results":    [],
        "quality_score":     0.0,
        "is_satisfactory":   False,
        "retry_count":       0,
        "rewrite_history":   [],
        "final_response":    "",
        "error_message":     None,
    }
    result = _graph.invoke(initial_state)
    return result