Agentic RAG 첫걸음 — Query Routing과 Adaptive Retrieval

def naive_rag(query: str) -> str:
    """단순 RAG: 항상 벡터 검색만 수행합니다."""
    docs = vector_store.similarity_search(query, k=4)
    context = "\n".join(d.page_content for d in docs)
    return llm.invoke(f"Context:\n{context}\n\nQ: {query}")

# 실패 케이스: "최근 OpenAI 매출 추이가 어떻게 돼?"
# → 벡터 DB에 실시간 데이터가 없어서 엉뚱한 답변 생성

from pydantic import BaseModel, Field
from typing import Literal
from openai import OpenAI

client = OpenAI()

class QueryAnalysis(BaseModel):
    """사용자 질문을 분석한 결과 스키마입니다."""
    intent: Literal["factual", "analytical", "comparison", "temporal", "opinion"]
    complexity: Literal["simple", "multi_hop", "aggregation"]
    requires_retrieval: bool
    suggested_sources: list[Literal["vector_db", "web_search", "sql_db"]]
    sub_queries: list[str] = Field(default_factory=list)

SYSTEM_PROMPT = """당신은 사용자 질문을 분석하는 전문가입니다.
질문의 의도(intent), 복잡도(complexity), 검색 필요 여부(requires_retrieval),
적절한 소스(suggested_sources), 하위 질문(sub_queries)을 판단하세요.

규칙:
- 실시간 정보가 필요하면 web_search를 추천하세요.
- 수치/통계 질문은 sql_db를 추천하세요.
- 일반 상식이나 개념 설명은 requires_retrieval=false로 설정하세요.
- 복합 질문은 sub_queries로 분해하세요."""

def analyze_query(query: str) -> QueryAnalysis:
    """사용자 질문의 의도와 복잡도를 분석합니다."""
    response = client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": query}
        ],
        response_format=QueryAnalysis
    )
    return response.choices[0].message.parsed

# 예시 1: 단순 사실 확인
result = analyze_query("트랜스포머 모델의 어텐션 메커니즘이 뭐야?")
# → intent="factual", complexity="simple",
#   requires_retrieval=True, suggested_sources=["vector_db"]

# 예시 2: 실시간 정보 필요
result = analyze_query("최근 OpenAI 매출 추이가 어떻게 돼?")
# → intent="temporal", complexity="aggregation",
#   requires_retrieval=True, suggested_sources=["web_search"]

# 예시 3: 검색 불필요
result = analyze_query("파이썬에서 리스트와 튜플의 차이는?")
# → intent="comparison", complexity="simple",
#   requires_retrieval=False, suggested_sources=[]

import chromadb
from tavily import TavilyClient
import sqlite3

# 1. 벡터 검색: ChromaDB
chroma_client = chromadb.PersistentClient(path="./chroma_db")
collection = chroma_client.get_collection("documents")

def vector_search(query: str, k: int = 4) -> list[str]:
    """벡터 DB에서 유사 문서를 검색합니다."""
    results = collection.query(query_texts=[query], n_results=k)
    return results["documents"][0]

# 2. 웹 검색: Tavily
tavily_client = TavilyClient(api_key="tvly-...")

def web_search(query: str) -> list[str]:
    """실시간 웹 검색을 수행합니다."""
    response = tavily_client.search(query, max_results=3)
    return [r["content"] for r in response["results"]]

# 3. Text-to-SQL: SQLite
conn = sqlite3.connect("./company.db")

def sql_query(query: str) -> str:
    """자연어를 SQL로 변환하여 실행합니다."""
    # LLM으로 자연어 → SQL 변환
    sql = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": (
                "다음 스키마를 참고하여 SQL을 생성하세요.\n"
                "테이블: sales(date, product, revenue, region)\n"
                "SQL만 출력하세요. 설명 없이."
            )},
            {"role": "user", "content": query}
        ]
    ).choices[0].message.content.strip()

    # SQL 실행
    cursor = conn.execute(sql)
    rows = cursor.fetchall()
    columns = [desc[0] for desc in cursor.description]
    return f"SQL: {sql}\n결과: {[dict(zip(columns, row)) for row in rows]}"

def route_query(analysis: QueryAnalysis, query: str) -> list[str]:
    """분석 결과에 따라 적절한 소스에서 정보를 검색합니다."""
    results = []

    for source in analysis.suggested_sources:
        if source == "vector_db":
            results.extend(vector_search(query))
        elif source == "web_search":
            results.extend(web_search(query))
        elif source == "sql_db":
            results.append(sql_query(query))

    return results

def adaptive_rag(query: str) -> str:
    """검색 필요 여부를 판단하고, 필요한 경우에만 검색합니다."""
    analysis = analyze_query(query)

    # 검색이 필요 없으면 LLM 지식으로 직접 답변
    if not analysis.requires_retrieval:
        return client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": "당신의 지식으로 정확하게 답변하세요."},
                {"role": "user", "content": query}
            ]
        ).choices[0].message.content

    # 복합 질문이면 하위 질문으로 분해해서 각각 검색
    queries = analysis.sub_queries if analysis.sub_queries else [query]
    all_context = []
    for q in queries:
        sub_analysis = analyze_query(q) if q != query else analysis
        all_context.extend(route_query(sub_analysis, q))

    # 수집한 문맥을 기반으로 최종 답변 생성
    context_text = "\n---\n".join(all_context)
    return client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "주어진 문맥을 기반으로 정확하게 답변하세요."},
            {"role": "user", "content": f"문맥:\n{context_text}\n\n질문: {query}"}
        ]
    ).choices[0].message.content

from typing import TypedDict, Annotated
from langgraph.graph import StateGraph, END

class AgenticRAGState(TypedDict):
    """Agentic RAG 파이프라인의 상태를 정의합니다."""
    query: str                              # 사용자 원본 질문
    analysis: QueryAnalysis | None          # 질문 분석 결과
    documents: list[str]                    # 검색된 문서들
    generation: str                         # 최종 생성된 답변


def analyze_node(state: AgenticRAGState) -> dict:
    """질문을 분석하여 의도와 최적 소스를 결정합니다."""
    analysis = analyze_query(state["query"])
    return {"analysis": analysis}


def should_retrieve(state: AgenticRAGState) -> str:
    """검색이 필요한지 판단하여 다음 노드를 결정합니다."""
    if state["analysis"].requires_retrieval:
        return "retrieve"
    return "generate_direct"


def retrieve_node(state: AgenticRAGState) -> dict:
    """분석 결과에 따라 적절한 소스에서 문서를 검색합니다."""
    analysis = state["analysis"]
    query = state["query"]

    all_docs = []
    # 하위 질문이 있으면 각각 검색
    queries = analysis.sub_queries if analysis.sub_queries else [query]
    for q in queries:
        sub_analysis = analyze_query(q) if q != query else analysis
        all_docs.extend(route_query(sub_analysis, q))

    return {"documents": all_docs}


def generate_node(state: AgenticRAGState) -> dict:
    """검색된 문서를 기반으로 답변을 생성합니다."""
    context = "\n---\n".join(state["documents"])
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": (
                "주어진 문맥을 기반으로 정확하게 답변하세요. "
                "문맥에 없는 내용은 추측하지 마세요."
            )},
            {"role": "user", "content": f"문맥:\n{context}\n\n질문: {state['query']}"}
        ]
    ).choices[0].message.content
    return {"generation": response}


def generate_direct_node(state: AgenticRAGState) -> dict:
    """검색 없이 LLM 지식만으로 답변을 생성합니다."""
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "당신의 지식으로 정확하고 친절하게 답변하세요."},
            {"role": "user", "content": state["query"]}
        ]
    ).choices[0].message.content
    return {"generation": response}


# 그래프 구성
graph = StateGraph(AgenticRAGState)

# 노드 추가
graph.add_node("analyze", analyze_node)
graph.add_node("retrieve", retrieve_node)
graph.add_node("generate", generate_node)
graph.add_node("generate_direct", generate_direct_node)

# 엣지 연결
graph.set_entry_point("analyze")
graph.add_conditional_edges(
    "analyze",
    should_retrieve,
    {
        "retrieve": "retrieve",           # 검색 필요 → 검색 노드로
        "generate_direct": "generate_direct"  # 검색 불필요 → 직접 답변
    }
)
graph.add_edge("retrieve", "generate")    # 검색 후 → 답변 생성
graph.add_edge("generate", END)           # 답변 완료
graph.add_edge("generate_direct", END)    # 직접 답변 완료

# 컴파일 및 실행
app = graph.compile()

# 실행 예시
result = app.invoke({
    "query": "작년 대비 올해 서울 날씨 변화를 분석해줘",
    "analysis": None,
    "documents": [],
    "generation": ""
})
print(result["generation"])