from sentence_transformers import SentenceTransformer
import numpy as np
 
model = SentenceTransformer("BAAI/bge-m3")
 
route_embeddings = {
    "legal":   model.encode(["계약", "소송", "법률 해석", "판례"]),
    "medical": model.encode(["진단", "약 복용", "혈압", "증상"]),
    "tech":    model.encode(["코드 오류", "API", "버그", "아키텍처"]),
}
# 각 경로의 대표 임베딩을 평균으로 압축
route_centroids = {k: v.mean(axis=0) for k, v in route_embeddings.items()}
 
def semantic_route(query: str) -> str:
    q_emb = model.encode(query)
    scores = {
        route: np.dot(q_emb, centroid) / (np.linalg.norm(q_emb) * np.linalg.norm(centroid))
        for route, centroid in route_centroids.items()
    }
    return max(scores, key=scores.get)

import anthropic
import json
 
client = anthropic.Anthropic()
 
ROUTES = {
    "legal":   "법률, 계약서, 판례, 소송과 관련된 질문",
    "medical": "의학 정보, 약물, 증상, 진단과 관련된 질문",
    "tech":    "소프트웨어, 코드, 아키텍처, 인프라와 관련된 질문",
    "general": "위 카테고리에 해당하지 않는 일반 질문",
}
 
def llm_route(query: str) -> str:
    route_desc = "\n".join(f"- {k}: {v}" for k, v in ROUTES.items())
    prompt = f"""다음 쿼리를 분석하여 가장 적합한 카테고리를 선택하라.
카테고리 목록:
{route_desc}
 
쿼리: {query}
 
JSON 형식으로만 응답하라: {{"route": "<카테고리명>"}}"""
 
    response = client.messages.create(
        model="claude-haiku-4-5-20251001",
        max_tokens=64,
        messages=[{"role": "user", "content": prompt}],
    )
    return json.loads(response.content[0].text)["route"]

from dataclasses import dataclass
from typing import Optional
 
@dataclass
class SearchResult:
    docs: list
    score: float
 
THRESHOLD = 0.75
 
def fallback_route(query: str) -> SearchResult:
    # 1단계: 빠른 BM25 검색
    result = bm25_retriever.search(query)
    if result.score >= THRESHOLD:
        return result
 
    # 2단계: Dense retriever로 에스컬레이션
    result = dense_retriever.search(query)
    if result.score >= THRESHOLD:
        return result
 
    # 3단계: LLM 기반 고품질 검색으로 최종 에스컬레이션
    return llm_retriever.search(query)

import asyncio
from typing import List
 
async def ensemble_route(query: str) -> List[dict]:
    # 복수 retriever를 병렬 실행
    results = await asyncio.gather(
        bm25_retriever.async_search(query),
        dense_retriever.async_search(query),
        colbert_retriever.async_search(query),
    )
 
    # Reciprocal Rank Fusion으로 병합
    return reciprocal_rank_fusion(results)
 
def reciprocal_rank_fusion(result_lists: list, k: int = 60) -> list:
    scores: dict = {}
    for ranked_list in result_lists:
        for rank, doc in enumerate(ranked_list):
            doc_id = doc["id"]
            scores[doc_id] = scores.get(doc_id, 0) + 1 / (k + rank + 1)
    return sorted(scores.keys(), key=lambda d: scores[d], reverse=True)

class HierarchicalRouter:
    def __init__(self):
        self.root_router = SemanticRouter(routes=["tech", "legal", "general"])
        self.tech_router = SemanticRouter(routes=["backend", "frontend", "infra"])
 
    def route(self, query: str) -> str:
        top_level = self.root_router.classify(query)
 
        if top_level == "tech":
            sub_level = self.tech_router.classify(query)
            return f"tech/{sub_level}"
 
        return top_level  # "legal" | "general"